120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
# [DEF:backend.src.core.superset_client:Module]
|
|
#
|
|
# @SEMANTICS: superset, api, client, database, metadata
|
|
# @PURPOSE: Extends the base SupersetClient with database-specific metadata fetching.
|
|
# @LAYER: Core
|
|
# @RELATION: INHERITS_FROM -> superset_tool.client.SupersetClient
|
|
#
|
|
# @INVARIANT: All database metadata requests must include UUID and name.
|
|
|
|
# [SECTION: IMPORTS]
|
|
from typing import List, Dict, Optional, Tuple
|
|
from backend.src.core.logger import belief_scope
|
|
from superset_tool.client import SupersetClient as BaseSupersetClient
|
|
from superset_tool.models import SupersetConfig
|
|
# [/SECTION]
|
|
|
|
# [DEF:SupersetClient:Class]
|
|
# @PURPOSE: Extended SupersetClient for migration-specific operations.
|
|
class SupersetClient(BaseSupersetClient):
|
|
|
|
# [DEF:get_databases_summary:Function]
|
|
# @PURPOSE: Fetch a summary of databases including uuid, name, and engine.
|
|
# @PRE: self.network must be initialized and authenticated.
|
|
# @POST: Returns a list of database dictionaries with 'engine' field.
|
|
# @RETURN: List[Dict] - Summary of databases.
|
|
def get_databases_summary(self) -> List[Dict]:
|
|
with belief_scope("SupersetClient.get_databases_summary"):
|
|
"""
|
|
Fetch a summary of databases including uuid, name, and engine.
|
|
"""
|
|
query = {
|
|
"columns": ["uuid", "database_name", "backend"]
|
|
}
|
|
_, databases = self.get_databases(query=query)
|
|
|
|
# Map 'backend' to 'engine' for consistency with contracts
|
|
for db in databases:
|
|
db['engine'] = db.pop('backend', None)
|
|
|
|
return databases
|
|
# [/DEF:get_databases_summary:Function]
|
|
|
|
# [DEF:get_database_by_uuid:Function]
|
|
# @PURPOSE: Find a database by its UUID.
|
|
# @PRE: db_uuid must be a string.
|
|
# @POST: Returns database metadata if found.
|
|
# @PARAM: db_uuid (str) - The UUID of the database.
|
|
# @RETURN: Optional[Dict] - Database info if found, else None.
|
|
def get_database_by_uuid(self, db_uuid: str) -> Optional[Dict]:
|
|
with belief_scope("SupersetClient.get_database_by_uuid", f"uuid={db_uuid}"):
|
|
"""
|
|
Find a database by its UUID.
|
|
"""
|
|
query = {
|
|
"filters": [{"col": "uuid", "op": "eq", "value": db_uuid}]
|
|
}
|
|
_, databases = self.get_databases(query=query)
|
|
return databases[0] if databases else None
|
|
# [/DEF:get_database_by_uuid:Function]
|
|
|
|
# [DEF:get_dashboards_summary:Function]
|
|
# @PURPOSE: Fetches dashboard metadata optimized for the grid.
|
|
# @PRE: self.network must be authenticated.
|
|
# @POST: Returns a list of dashboard dictionaries mapped to the grid schema.
|
|
# @RETURN: List[Dict]
|
|
def get_dashboards_summary(self) -> List[Dict]:
|
|
with belief_scope("SupersetClient.get_dashboards_summary"):
|
|
"""
|
|
Fetches dashboard metadata optimized for the grid.
|
|
Returns a list of dictionaries mapped to DashboardMetadata fields.
|
|
"""
|
|
query = {
|
|
"columns": ["id", "dashboard_title", "changed_on_utc", "published"]
|
|
}
|
|
_, dashboards = self.get_dashboards(query=query)
|
|
|
|
# Map fields to DashboardMetadata schema
|
|
result = []
|
|
for dash in dashboards:
|
|
result.append({
|
|
"id": dash.get("id"),
|
|
"title": dash.get("dashboard_title"),
|
|
"last_modified": dash.get("changed_on_utc"),
|
|
"status": "published" if dash.get("published") else "draft"
|
|
})
|
|
return result
|
|
# [/DEF:get_dashboards_summary:Function]
|
|
|
|
# [DEF:get_dataset:Function]
|
|
# @PURPOSE: Fetch full dataset structure including columns and metrics.
|
|
# @PRE: dataset_id must be a valid integer.
|
|
# @POST: Returns full dataset metadata from Superset API.
|
|
# @PARAM: dataset_id (int) - The ID of the dataset.
|
|
# @RETURN: Dict - The dataset metadata.
|
|
def get_dataset(self, dataset_id: int) -> Dict:
|
|
with belief_scope("SupersetClient.get_dataset", f"id={dataset_id}"):
|
|
"""
|
|
Fetch full dataset structure.
|
|
"""
|
|
return self.network.get(f"/api/v1/dataset/{dataset_id}").json()
|
|
# [/DEF:get_dataset:Function]
|
|
|
|
# [DEF:update_dataset:Function]
|
|
# @PURPOSE: Update dataset metadata.
|
|
# @PRE: dataset_id must be valid, data must be a valid Superset dataset payload.
|
|
# @POST: Dataset is updated in Superset.
|
|
# @PARAM: dataset_id (int) - The ID of the dataset.
|
|
# @PARAM: data (Dict) - The payload for update.
|
|
def update_dataset(self, dataset_id: int, data: Dict):
|
|
with belief_scope("SupersetClient.update_dataset", f"id={dataset_id}"):
|
|
"""
|
|
Update dataset metadata.
|
|
"""
|
|
self.network.put(f"/api/v1/dataset/{dataset_id}", json=data)
|
|
# [/DEF:update_dataset:Function]
|
|
|
|
# [/DEF:SupersetClient:Class]
|
|
|
|
# [/DEF:backend.src.core.superset_client:Module]
|