Files
ss-tools/backend/src/core/async_superset_client.py
2026-05-20 23:54:53 +03:00

616 lines
29 KiB
Python

# #region AsyncSupersetClientModule [C:3] [TYPE Module] [SEMANTICS superset, transform, dashboard, filter, async-superset-client]
# @BRIEF Parse a Superset dashboard URL and extract native filter state asynchronously.
# @LAYER: Core
import asyncio
import json
import re
from typing import Any, cast
from .config_models import Environment
from .logger import belief_scope, logger as app_logger
from .superset_client import SupersetClient
from .utils.async_network import AsyncAPIClient
# #region AsyncSupersetClient [C:3] [TYPE Class]
# @BRIEF Async sibling of SupersetClient for dashboard read paths.
# @RELATION INHERITS -> [SupersetClient]
# @RELATION DEPENDS_ON -> [AsyncAPIClient]
# @RELATION CALLS -> [AsyncAPIClient.request]
class AsyncSupersetClient(SupersetClient):
# #region AsyncSupersetClientInit [TYPE Function] [C:3]
# @PURPOSE: Initialize async Superset client with AsyncAPIClient transport.
# @PRE: env is valid Environment instance.
# @POST: Client uses async network transport and inherited projection helpers.
# @DATA_CONTRACT: Input[Environment] -> self.network[AsyncAPIClient]
# @RELATION: [DEPENDS_ON] ->[AsyncAPIClient]
def __init__(self, env: Environment):
self.env = env
auth_payload = {
"username": env.username,
"password": env.password,
"provider": "db",
"refresh": "true",
}
self.network = AsyncAPIClient(
config={"base_url": env.url, "auth": auth_payload},
verify_ssl=env.verify_ssl,
timeout=env.timeout,
)
self.delete_before_reimport = False
# #endregion AsyncSupersetClientInit
# #region AsyncSupersetClientClose [TYPE Function] [C:3]
# @PURPOSE: Close async transport resources.
# @POST: Underlying AsyncAPIClient is closed.
# @SIDE_EFFECT: Closes network sockets.
# @RELATION: [CALLS] ->[AsyncAPIClient.aclose]
async def aclose(self) -> None:
await self.network.aclose()
# #endregion AsyncSupersetClientClose
# #region get_dashboards_page_async [TYPE Function] [C:3]
# @PURPOSE: Fetch one dashboards page asynchronously.
# @POST: Returns total count and page result list.
# @DATA_CONTRACT: Input[query: Optional[Dict]] -> Output[Tuple[int, List[Dict]]]
# @RELATION: [CALLS] -> [AsyncAPIClient.request]
async def get_dashboards_page_async(
self, query: dict | None = None
) -> tuple[int, list[dict]]:
with belief_scope("AsyncSupersetClient.get_dashboards_page_async"):
validated_query = self._validate_query_params(query or {})
if "columns" not in validated_query:
validated_query["columns"] = [
"slug",
"id",
"url",
"changed_on_utc",
"dashboard_title",
"published",
"created_by",
"changed_by",
"changed_by_name",
"owners",
]
response_json = cast(
dict[str, Any],
await self.network.request(
method="GET",
endpoint="/dashboard/",
params={"q": json.dumps(validated_query)},
),
)
result = response_json.get("result", [])
total_count = response_json.get("count", len(result))
return total_count, result
# #endregion get_dashboards_page_async
# #region get_dashboard_async [TYPE Function] [C:3]
# @PURPOSE: Fetch one dashboard payload asynchronously.
# @POST: Returns raw dashboard payload from Superset API.
# @DATA_CONTRACT: Input[dashboard_id: int] -> Output[Dict]
# @RELATION: [CALLS] ->[AsyncAPIClient.request]
async def get_dashboard_async(self, dashboard_id: int) -> dict:
with belief_scope(
"AsyncSupersetClient.get_dashboard_async", f"id={dashboard_id}"
):
response = await self.network.request(
method="GET", endpoint=f"/dashboard/{dashboard_id}"
)
return cast(dict, response)
# #endregion get_dashboard_async
# #region get_chart_async [TYPE Function] [C:3]
# @PURPOSE: Fetch one chart payload asynchronously.
# @POST: Returns raw chart payload from Superset API.
# @DATA_CONTRACT: Input[chart_id: int] -> Output[Dict]
# @RELATION: [CALLS] ->[AsyncAPIClient.request]
async def get_chart_async(self, chart_id: int) -> dict:
with belief_scope("AsyncSupersetClient.get_chart_async", f"id={chart_id}"):
response = await self.network.request(
method="GET", endpoint=f"/chart/{chart_id}"
)
return cast(dict, response)
# #endregion get_chart_async
# #region get_dashboard_detail_async [TYPE Function] [C:3]
# @PURPOSE: Fetch dashboard detail asynchronously with concurrent charts/datasets requests.
# @POST: Returns dashboard detail payload for overview page.
# @DATA_CONTRACT: Input[dashboard_id: int] -> Output[Dict]
# @RELATION: [CALLS] ->[get_dashboard_async]
# @RELATION: [CALLS] ->[get_chart_async]
async def get_dashboard_detail_async(self, dashboard_id: int) -> dict:
with belief_scope(
"AsyncSupersetClient.get_dashboard_detail_async", f"id={dashboard_id}"
):
dashboard_response = await self.get_dashboard_async(dashboard_id)
dashboard_data = dashboard_response.get("result", dashboard_response)
charts: list[dict] = []
datasets: list[dict] = []
def extract_dataset_id_from_form_data(
form_data: dict | None,
) -> int | None:
if not isinstance(form_data, dict):
return None
datasource = form_data.get("datasource")
if isinstance(datasource, str):
matched = re.match(r"^(\d+)__", datasource)
if matched:
try:
return int(matched.group(1))
except ValueError:
return None
if isinstance(datasource, dict):
ds_id = datasource.get("id")
try:
return int(ds_id) if ds_id is not None else None
except (TypeError, ValueError):
return None
ds_id = form_data.get("datasource_id")
try:
return int(ds_id) if ds_id is not None else None
except (TypeError, ValueError):
return None
chart_task = self.network.request(
method="GET",
endpoint=f"/dashboard/{dashboard_id}/charts",
)
dataset_task = self.network.request(
method="GET",
endpoint=f"/dashboard/{dashboard_id}/datasets",
)
charts_response, datasets_response = await asyncio.gather(
chart_task,
dataset_task,
return_exceptions=True,
)
if not isinstance(charts_response, Exception):
charts_payload = (
charts_response.get("result", [])
if isinstance(charts_response, dict)
else []
)
for chart_obj in charts_payload:
if not isinstance(chart_obj, dict):
continue
chart_id = chart_obj.get("id")
if chart_id is None:
continue
form_data = chart_obj.get("form_data")
if isinstance(form_data, str):
try:
form_data = json.loads(form_data)
except Exception:
form_data = {}
dataset_id = extract_dataset_id_from_form_data(
form_data
) or chart_obj.get("datasource_id")
charts.append(
{
"id": int(chart_id),
"title": chart_obj.get("slice_name")
or chart_obj.get("name")
or f"Chart {chart_id}",
"viz_type": (
form_data.get("viz_type")
if isinstance(form_data, dict)
else None
),
"dataset_id": int(dataset_id)
if dataset_id is not None
else None,
"last_modified": chart_obj.get("changed_on"),
"overview": chart_obj.get("description")
or (
form_data.get("viz_type")
if isinstance(form_data, dict)
else None
)
or "Chart",
}
)
else:
app_logger.warning(
"[get_dashboard_detail_async][Warning] Failed to fetch dashboard charts: %s",
charts_response,
)
if not isinstance(datasets_response, Exception):
datasets_payload = (
datasets_response.get("result", [])
if isinstance(datasets_response, dict)
else []
)
for dataset_obj in datasets_payload:
if not isinstance(dataset_obj, dict):
continue
dataset_id = dataset_obj.get("id")
if dataset_id is None:
continue
db_payload = dataset_obj.get("database")
db_name = (
db_payload.get("database_name")
if isinstance(db_payload, dict)
else None
)
table_name = (
dataset_obj.get("table_name")
or dataset_obj.get("datasource_name")
or dataset_obj.get("name")
or f"Dataset {dataset_id}"
)
schema = dataset_obj.get("schema")
fq_name = f"{schema}.{table_name}" if schema else table_name
datasets.append(
{
"id": int(dataset_id),
"table_name": table_name,
"schema": schema,
"database": db_name
or dataset_obj.get("database_name")
or "Unknown",
"last_modified": dataset_obj.get("changed_on"),
"overview": fq_name,
}
)
else:
app_logger.warning(
"[get_dashboard_detail_async][Warning] Failed to fetch dashboard datasets: %s",
datasets_response,
)
if not charts:
raw_position_json = dashboard_data.get("position_json")
chart_ids_from_position = set()
if isinstance(raw_position_json, str) and raw_position_json:
try:
parsed_position = json.loads(raw_position_json)
chart_ids_from_position.update(
self._extract_chart_ids_from_layout(parsed_position)
)
except Exception:
app_logger.debug("Could not parse position JSON for dashboard %s", dashboard_ref)
elif isinstance(raw_position_json, dict):
chart_ids_from_position.update(
self._extract_chart_ids_from_layout(raw_position_json)
)
raw_json_metadata = dashboard_data.get("json_metadata")
if isinstance(raw_json_metadata, str) and raw_json_metadata:
try:
parsed_metadata = json.loads(raw_json_metadata)
chart_ids_from_position.update(
self._extract_chart_ids_from_layout(parsed_metadata)
)
except Exception:
app_logger.debug("Could not parse json_metadata for dashboard %s", dashboard_ref)
elif isinstance(raw_json_metadata, dict):
chart_ids_from_position.update(
self._extract_chart_ids_from_layout(raw_json_metadata)
)
fallback_chart_tasks = [
self.get_chart_async(int(chart_id))
for chart_id in sorted(chart_ids_from_position)
]
fallback_chart_responses = await asyncio.gather(
*fallback_chart_tasks,
return_exceptions=True,
)
for chart_id, chart_response in zip(
sorted(chart_ids_from_position), fallback_chart_responses
):
if isinstance(chart_response, Exception):
app_logger.warning(
"[get_dashboard_detail_async][Warning] Failed to resolve fallback chart %s: %s",
chart_id,
chart_response,
)
continue
chart_data = chart_response.get("result", chart_response)
charts.append(
{
"id": int(chart_id),
"title": chart_data.get("slice_name")
or chart_data.get("name")
or f"Chart {chart_id}",
"viz_type": chart_data.get("viz_type"),
"dataset_id": chart_data.get("datasource_id"),
"last_modified": chart_data.get("changed_on"),
"overview": chart_data.get("description")
or chart_data.get("viz_type")
or "Chart",
}
)
dataset_ids_from_charts = {
c.get("dataset_id") for c in charts if c.get("dataset_id") is not None
}
known_dataset_ids = {
d.get("id") for d in datasets if d.get("id") is not None
}
missing_dataset_ids = sorted(
int(item)
for item in dataset_ids_from_charts
if item not in known_dataset_ids
)
if missing_dataset_ids:
dataset_fetch_tasks = [
self.network.request(
method="GET", endpoint=f"/dataset/{dataset_id}"
)
for dataset_id in missing_dataset_ids
]
dataset_fetch_responses = await asyncio.gather(
*dataset_fetch_tasks,
return_exceptions=True,
)
for dataset_id, dataset_response in zip(
missing_dataset_ids, dataset_fetch_responses
):
if isinstance(dataset_response, Exception):
app_logger.warning(
"[get_dashboard_detail_async][Warning] Failed to backfill dataset %s: %s",
dataset_id,
dataset_response,
)
continue
dataset_data = (
dataset_response.get("result", dataset_response)
if isinstance(dataset_response, dict)
else {}
)
db_payload = dataset_data.get("database")
db_name = (
db_payload.get("database_name")
if isinstance(db_payload, dict)
else None
)
table_name = (
dataset_data.get("table_name")
or dataset_data.get("datasource_name")
or dataset_data.get("name")
or f"Dataset {dataset_id}"
)
schema = dataset_data.get("schema")
fq_name = f" {schema}.{table_name}" if schema else table_name
datasets.append(
{
"id": int(dataset_id),
"table_name": table_name,
"schema": schema,
"database": db_name
or dataset_data.get("database_name")
or "Unknown",
"last_modified": dataset_data.get("changed_on"),
"overview": fq_name,
}
)
return {
"id": int(dashboard_data.get("id") or dashboard_id),
"title": dashboard_data.get("dashboard_title")
or dashboard_data.get("title")
or f"Dashboard {dashboard_id}",
"slug": dashboard_data.get("slug"),
"url": dashboard_data.get("url"),
"description": dashboard_data.get("description"),
"last_modified": dashboard_data.get("changed_on_utc")
or dashboard_data.get("changed_on"),
"published": dashboard_data.get("published"),
"charts": charts,
"datasets": datasets,
"chart_count": len(charts),
"dataset_count": len(datasets),
}
# #endregion get_dashboard_detail_async
# #region get_dashboard_permalink_state_async [TYPE Function] [C:2]
# @PURPOSE: Fetch stored dashboard permalink state asynchronously.
# @POST: Returns dashboard permalink state payload from Superset API.
# @DATA_CONTRACT: Input[permalink_key: str] -> Output[Dict]
async def get_dashboard_permalink_state_async(self, permalink_key: str) -> dict:
with belief_scope(
"AsyncSupersetClient.get_dashboard_permalink_state_async",
f"key={permalink_key}",
):
response = await self.network.request(
method="GET", endpoint=f"/dashboard/permalink/{permalink_key}"
)
return cast(dict, response)
# #endregion get_dashboard_permalink_state_async
# #region get_native_filter_state_async [TYPE Function] [C:2]
# @PURPOSE: Fetch stored native filter state asynchronously.
# @POST: Returns native filter state payload from Superset API.
# @DATA_CONTRACT: Input[dashboard_id: Union[int, str], filter_state_key: str] -> Output[Dict]
async def get_native_filter_state_async(
self, dashboard_id: int, filter_state_key: str
) -> dict:
with belief_scope(
"AsyncSupersetClient.get_native_filter_state_async",
f"dashboard={dashboard_id}, key={filter_state_key}",
):
response = await self.network.request(
method="GET",
endpoint=f"/dashboard/{dashboard_id}/filter_state/{filter_state_key}",
)
return cast(dict, response)
# #endregion get_native_filter_state_async
# #region extract_native_filters_from_permalink_async [TYPE Function] [C:3]
# @PURPOSE: Extract native filters dataMask from a permalink key asynchronously.
# @POST: Returns extracted dataMask with filter states.
# @DATA_CONTRACT: Input[permalink_key: str] -> Output[Dict]
# @RELATION: [CALLS] ->[get_dashboard_permalink_state_async]
async def extract_native_filters_from_permalink_async(
self, permalink_key: str
) -> dict:
with belief_scope(
"AsyncSupersetClient.extract_native_filters_from_permalink_async",
f"key={permalink_key}",
):
permalink_response = await self.get_dashboard_permalink_state_async(
permalink_key
)
result = permalink_response.get("result", permalink_response)
state = result.get("state", result)
data_mask = state.get("dataMask", {})
extracted_filters = {}
for filter_id, filter_data in data_mask.items():
if not isinstance(filter_data, dict):
continue
extracted_filters[filter_id] = {
"extraFormData": filter_data.get("extraFormData", {}),
"filterState": filter_data.get("filterState", {}),
"ownState": filter_data.get("ownState", {}),
}
return {
"dataMask": extracted_filters,
"activeTabs": state.get("activeTabs", []),
"anchor": state.get("anchor"),
"chartStates": state.get("chartStates", {}),
"permalink_key": permalink_key,
}
# #endregion extract_native_filters_from_permalink_async
# #region extract_native_filters_from_key_async [TYPE Function] [C:3]
# @PURPOSE: Extract native filters from a native_filters_key URL parameter asynchronously.
# @POST: Returns extracted filter state with extraFormData.
# @DATA_CONTRACT: Input[dashboard_id: Union[int, str], filter_state_key: str] -> Output[Dict]
# @RELATION: [CALLS] ->[get_native_filter_state_async]
async def extract_native_filters_from_key_async(
self, dashboard_id: int, filter_state_key: str
) -> dict:
with belief_scope(
"AsyncSupersetClient.extract_native_filters_from_key_async",
f"dashboard={dashboard_id}, key={filter_state_key}",
):
filter_response = await self.get_native_filter_state_async(
dashboard_id, filter_state_key
)
result = filter_response.get("result", filter_response)
value = result.get("value")
if isinstance(value, str):
try:
parsed_value = json.loads(value)
except json.JSONDecodeError as e:
app_logger.warning(
"[extract_native_filters_from_key_async][Warning] Failed to parse filter state JSON: %s",
e,
)
parsed_value = {}
elif isinstance(value, dict):
parsed_value = value
else:
parsed_value = {}
extracted_filters = {}
if "id" in parsed_value and "extraFormData" in parsed_value:
filter_id = parsed_value.get("id", filter_state_key)
extracted_filters[filter_id] = {
"extraFormData": parsed_value.get("extraFormData", {}),
"filterState": parsed_value.get("filterState", {}),
"ownState": parsed_value.get("ownState", {}),
}
else:
for filter_id, filter_data in parsed_value.items():
if not isinstance(filter_data, dict):
continue
extracted_filters[filter_id] = {
"extraFormData": filter_data.get("extraFormData", {}),
"filterState": filter_data.get("filterState", {}),
"ownState": filter_data.get("ownState", {}),
}
return {
"dataMask": extracted_filters,
"dashboard_id": dashboard_id,
"filter_state_key": filter_state_key,
}
# #endregion extract_native_filters_from_key_async
# #region parse_dashboard_url_for_filters_async [TYPE Function] [C:3]
# @PURPOSE: Parse a Superset dashboard URL and extract native filter state asynchronously.
# @POST: Returns extracted filter state or empty dict if no filters found.
# @DATA_CONTRACT: Input[url: str] -> Output[Dict]
# @RELATION: [CALLS] ->[extract_native_filters_from_permalink_async]
# @RELATION: [CALLS] ->[extract_native_filters_from_key_async]
async def parse_dashboard_url_for_filters_async(self, url: str) -> dict:
with belief_scope(
"AsyncSupersetClient.parse_dashboard_url_for_filters_async", f"url={url}"
):
import urllib.parse
parsed_url = urllib.parse.urlparse(url)
query_params = urllib.parse.parse_qs(parsed_url.query)
path_parts = parsed_url.path.rstrip("/").split("/")
result = {
"url": url,
"dashboard_id": None,
"filter_type": None,
"filters": {},
}
# Check for permalink URL: /dashboard/p/{key}/
if "p" in path_parts:
try:
p_index = path_parts.index("p")
if p_index + 1 < len(path_parts):
permalink_key = path_parts[p_index + 1]
filter_data = (
await self.extract_native_filters_from_permalink_async(
permalink_key
)
)
result["filter_type"] = "permalink"
result["filters"] = filter_data
return result
except ValueError:
app_logger.debug("Permalink key not found in URL for dashboard %s", dashboard_ref)
# Check for native_filters_key in query params
native_filters_key = query_params.get("native_filters_key", [None])[0]
if native_filters_key:
dashboard_ref = None
if "dashboard" in path_parts:
try:
dash_index = path_parts.index("dashboard")
if dash_index + 1 < len(path_parts):
potential_id = path_parts[dash_index + 1]
if potential_id not in ("p", "list", "new"):
dashboard_ref = potential_id
except ValueError:
app_logger.debug("Dashboard not found in path_parts for native_filters_key")
if dashboard_ref:
# Resolve slug to numeric ID — the filter_state API requires a numeric ID
resolved_id = None
try:
resolved_id = int(dashboard_ref)
except (ValueError, TypeError):
try:
dash_resp = await self.get_dashboard_async(dashboard_ref)
dash_data = (
dash_resp.get("result", dash_resp)
if isinstance(dash_resp, dict)
else {}
)
raw_id = dash_data.get("id")
if raw_id is not None:
resolved_id = int(raw_id)
except Exception as e:
app_logger.warning(
"[parse_dashboard_url_for_filters_async][Warning] Failed to resolve dashboard slug '%s' to ID: %s",
dashboard_ref,
e,
)
if resolved_id is not None:
filter_data = await self.extract_native_filters_from_key_async(
resolved_id, native_filters_key
)
result["filter_type"] = "native_filters_key"
result["dashboard_id"] = resolved_id
result["filters"] = filter_data
return result
else:
app_logger.warning(
"[parse_dashboard_url_for_filters_async][Warning] Could not resolve dashboard_id from URL for native_filters_key"
)
return result
# Check for native_filters in query params (direct filter values)
native_filters = query_params.get("native_filters", [None])[0]
if native_filters:
try:
parsed_filters = json.loads(native_filters)
result["filter_type"] = "native_filters"
result["filters"] = {"dataMask": parsed_filters}
return result
except json.JSONDecodeError as e:
app_logger.warning(
"[parse_dashboard_url_for_filters_async][Warning] Failed to parse native_filters JSON: %s",
e,
)
return result
# #endregion parse_dashboard_url_for_filters_async
# #endregion AsyncSupersetClient
# #endregion AsyncSupersetClientModule