busya
f416583a8c
resolve_database_id → async (was sync but called async SupersetClient methods) - await client.get_database(db_id) - await client.get_databases(...) validate_target_table_schema → async (calls async resolve_database_id + execute_and_poll) - await executor.resolve_database_id(...) - await executor.execute_and_poll(...) Route check_target_schema → await validate_target_table_schema(...) MapperPlugin.execute → await executor.resolve_database_id(...) (execute() was already async def, just missing await) SupersetSqlLabExecutor.execute_sql → await self.resolve_database_id() (was sync call to now-async method)
429 lines
18 KiB
Python
429 lines
18 KiB
Python
# #region SupersetSqlLabExecutor [C:4] [TYPE Module] [SEMANTICS translate, superset, sqllab, execute, query]
|
|
# @BRIEF Submit SQL to Superset SQL Lab API and poll execution status.
|
|
# @LAYER Infrastructure
|
|
# @RELATION DEPENDS_ON -> [ConfigManager]
|
|
# @RELATION DEPENDS_ON -> [ConfigManager]
|
|
# @PRE Valid Superset environment configuration and authenticated client.
|
|
# @POST SQL is submitted to Superset SQL Lab; execution reference is returned.
|
|
# @SIDE_EFFECT Makes HTTP calls to Superset /api/v1/sqllab/execute/; polls status.
|
|
# @RATIONALE Direct SQL Lab API submission provides audit trail and execution monitoring within Superset.
|
|
# @REJECTED Direct database connection bypass — would skip Superset's SQL Lab audit and RBAC.
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any
|
|
import uuid
|
|
|
|
from ...core.config_manager import ConfigManager
|
|
from ...core.logger import belief_scope, logger
|
|
from ...core.superset_client import SupersetClient
|
|
|
|
|
|
# #region SupersetSqlLabExecutor [C:4] [TYPE Class]
|
|
# @BRIEF Submit SQL to Superset SQL Lab API with polling and status tracking.
|
|
# @PRE Valid environment ID and ConfigManager.
|
|
# @POST SQL submitted to Superset; execution reference recorded.
|
|
# @SIDE_EFFECT Makes HTTP POST to /api/v1/sqllab/execute/; polls GET for status.
|
|
class SupersetSqlLabExecutor:
|
|
|
|
def __init__(self, config_manager: ConfigManager, env_id: str):
|
|
self.config_manager = config_manager
|
|
self.env_id = env_id
|
|
self._client: SupersetClient | None = None
|
|
self._database_id: int | None = None
|
|
self._database_backend: str | None = None
|
|
self._database_name: str | None = None
|
|
|
|
# region _get_client [TYPE Function]
|
|
# @PURPOSE Lazy-initialize SupersetClient for the configured environment.
|
|
# @PRE env_id must correspond to a valid environment config.
|
|
# @POST Returns authenticated SupersetClient.
|
|
# @SIDE_EFFECT Authenticates against Superset API.
|
|
def _get_client(self) -> SupersetClient:
|
|
with belief_scope("SupersetSqlLabExecutor._get_client"):
|
|
if self._client is not None:
|
|
return self._client
|
|
|
|
environments = self.config_manager.get_environments()
|
|
env_config = next((e for e in environments if e.id == self.env_id), None)
|
|
if not env_config:
|
|
raise ValueError(f"Superset environment '{self.env_id}' not found")
|
|
|
|
self._client = SupersetClient(env_config)
|
|
return self._client
|
|
# endregion _get_client
|
|
|
|
# region resolve_database_id [TYPE Function]
|
|
# @PURPOSE Resolve the target database ID from the environment and fetch its backend engine.
|
|
# @PRE database_name or database_id should be known.
|
|
# @POST Returns database_id integer or raises ValueError. Also sets self._database_backend.
|
|
# @SIDE_EFFECT Fetches databases list from Superset; optionally fetches single DB for backend.
|
|
async def resolve_database_id(
|
|
self,
|
|
database_name: str | None = None,
|
|
target_database_id: str | None = None,
|
|
) -> int:
|
|
with belief_scope("SupersetSqlLabExecutor.resolve_database_id"):
|
|
client = self._get_client()
|
|
|
|
# If a specific target_database_id is provided, use it directly
|
|
if target_database_id:
|
|
try:
|
|
db_id = int(target_database_id)
|
|
# Fetch full DB info to get backend
|
|
db_info = await client.get_database(db_id)
|
|
result = db_info.get("result", db_info)
|
|
self._database_id = result.get("id", db_id)
|
|
self._database_backend = result.get("backend") or result.get("engine")
|
|
self._database_name = result.get("database_name", "")
|
|
logger.reason("Resolved database ID by target_database_id", {
|
|
"target_database_id": target_database_id,
|
|
"database_id": self._database_id,
|
|
"database_name": self._database_name,
|
|
"backend": self._database_backend,
|
|
"all_keys": list(result.keys()),
|
|
})
|
|
return self._database_id
|
|
except Exception as e:
|
|
logger.explore("Failed to resolve by target_database_id, falling back", {
|
|
"target_database_id": target_database_id,
|
|
"error": str(e),
|
|
})
|
|
# Fall through to default resolution
|
|
|
|
# Fetch full database list with backend info
|
|
_, databases = await client.get_databases(
|
|
query={"columns": ["id", "database_name", "backend"]}
|
|
)
|
|
if not databases:
|
|
raise ValueError("No databases found in Superset environment")
|
|
|
|
if database_name:
|
|
for db in databases:
|
|
if db.get("database_name", "").lower() == database_name.lower():
|
|
self._database_id = db["id"]
|
|
self._database_backend = db.get("backend")
|
|
self._database_name = db.get("database_name", "")
|
|
logger.reason("Resolved database ID by name", {
|
|
"database_name": database_name,
|
|
"database_id": self._database_id,
|
|
"actual_name": self._database_name,
|
|
"backend": self._database_backend,
|
|
})
|
|
return self._database_id
|
|
raise ValueError(
|
|
f"Database '{database_name}' not found in Superset environment"
|
|
)
|
|
|
|
# Default: use first database
|
|
self._database_id = databases[0]["id"]
|
|
self._database_backend = databases[0].get("backend")
|
|
self._database_name = databases[0].get("database_name", "")
|
|
logger.reason("Using default database", {
|
|
"database_name": self._database_name,
|
|
"database_id": self._database_id,
|
|
"backend": self._database_backend,
|
|
})
|
|
return self._database_id
|
|
# endregion resolve_database_id
|
|
|
|
# region get_database_backend [TYPE Function]
|
|
# @PURPOSE Return the cached database backend/engine string.
|
|
# @POST Returns backend string or None if not yet resolved.
|
|
def get_database_backend(self) -> str | None:
|
|
return self._database_backend
|
|
# endregion get_database_backend
|
|
|
|
# region get_database_name [TYPE Function]
|
|
# @PURPOSE Return the cached database name.
|
|
# @POST Returns database name string or None if not yet resolved.
|
|
def get_database_name(self) -> str | None:
|
|
return self._database_name
|
|
# endregion get_database_name
|
|
|
|
# region execute_sql [TYPE Function]
|
|
# @PURPOSE Submit SQL to Superset SQL Lab and return execution tracking info.
|
|
# @PRE sql is valid SQL string. database_id is a valid Superset DB ID.
|
|
# @POST Returns execution result dict with query_id and status.
|
|
# @SIDE_EFFECT Makes HTTP POST to /api/v1/sqllab/execute/.
|
|
async def execute_sql(
|
|
self,
|
|
sql: str,
|
|
database_id: int | None = None,
|
|
run_async: bool = True,
|
|
) -> dict[str, Any]:
|
|
with belief_scope("SupersetSqlLabExecutor.execute_sql"):
|
|
client = self._get_client()
|
|
db_id = database_id or self._database_id
|
|
if not db_id:
|
|
db_id = await self.resolve_database_id()
|
|
|
|
logger.reason("Submitting SQL to Superset SQL Lab", {
|
|
"database_id": db_id,
|
|
"sql_length": len(sql),
|
|
"run_async": run_async,
|
|
})
|
|
|
|
payload = {
|
|
"database_id": db_id,
|
|
"sql": sql,
|
|
"runAsync": run_async,
|
|
"schema": None,
|
|
"tab": "translation-insert",
|
|
"client_id": f"trl-{uuid.uuid4().hex[:7]}", # max 11 chars for Superset varchar(11)
|
|
}
|
|
|
|
# Use the canonical Superset SQL Lab execute endpoint.
|
|
# NOTE: Some older docs reference /sql_lab/execute/ (with underscore),
|
|
# but this endpoint does NOT exist in current Superset versions.
|
|
# The browser UI uses /api/v1/sqllab/execute/ — confirmed by
|
|
# fetch() calls in Superset SQL Lab.
|
|
endpoint = "/api/v1/sqllab/execute/"
|
|
try:
|
|
response = await client.network.request(
|
|
method="POST",
|
|
endpoint=endpoint,
|
|
data=json.dumps(payload),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
logger.reason("SQL Lab endpoint succeeded", extra={
|
|
"endpoint": endpoint,
|
|
})
|
|
except Exception as ep_err:
|
|
logger.explore("SQL Lab execute failed", extra={
|
|
"error": str(ep_err),
|
|
"endpoint": endpoint,
|
|
})
|
|
raise ValueError(
|
|
f"Superset SQL Lab execute failed: {ep_err}"
|
|
) from ep_err
|
|
|
|
# Parse response — try multiple known Superset response formats
|
|
result = response if isinstance(response, dict) else {}
|
|
|
|
# Superset may return query_id at top level, nested in "result", as "id", or in "query" block
|
|
query_id = (
|
|
result.get("query_id")
|
|
or result.get("id")
|
|
or (result.get("result") or {}).get("query_id")
|
|
or (result.get("result") or {}).get("id")
|
|
or (result.get("query") or {}).get("query_id")
|
|
or (result.get("query") or {}).get("id")
|
|
)
|
|
status = result.get("status", "unknown")
|
|
|
|
# Log full response for debugging if query_id is missing
|
|
if not query_id:
|
|
logger.explore("No query_id from SQL Lab execute", extra={
|
|
"database_id": db_id,
|
|
"status": status,
|
|
"raw_response_keys": list(result.keys()),
|
|
"raw_response_preview": json.dumps(result)[:2000],
|
|
})
|
|
logger.reason("SQL Lab execute response (no query_id)", extra={
|
|
"database_id": db_id,
|
|
"status": status,
|
|
"response_keys": list(result.keys()),
|
|
"has_result": "result" in result,
|
|
"has_query": "query" in result,
|
|
"raw_preview": json.dumps(result)[:2000],
|
|
})
|
|
|
|
logger.reason("SQL Lab execute response", {
|
|
"query_id": query_id,
|
|
"status": status,
|
|
"response_keys": list(result.keys()),
|
|
})
|
|
|
|
return {
|
|
"query_id": query_id,
|
|
"status": status,
|
|
"raw_response": result,
|
|
"database_id": db_id,
|
|
}
|
|
# endregion execute_sql
|
|
|
|
# region poll_execution_status [TYPE Function]
|
|
# @PURPOSE Poll Superset for SQL execution status until completion or timeout.
|
|
# @PRE query_id is a valid Superset query ID.
|
|
# @POST Returns final execution status dict.
|
|
# @SIDE_EFFECT Makes HTTP GET requests to Superset.
|
|
async def poll_execution_status(
|
|
self,
|
|
query_id: str,
|
|
max_polls: int = 60,
|
|
poll_interval_seconds: float = 2.0,
|
|
) -> dict[str, Any]:
|
|
with belief_scope("SupersetSqlLabExecutor.poll_execution_status"):
|
|
client = self._get_client()
|
|
|
|
logger.reason("Polling SQL execution status", {
|
|
"query_id": query_id,
|
|
"max_polls": max_polls,
|
|
"interval": poll_interval_seconds,
|
|
})
|
|
|
|
for attempt in range(max_polls):
|
|
try:
|
|
response = await client.network.request(
|
|
method="GET",
|
|
endpoint=f"/api/v1/query/{query_id}",
|
|
)
|
|
raw = response if isinstance(response, dict) else {}
|
|
# Superset wraps query data in a nested "result" field
|
|
result = raw.get("result", raw)
|
|
status = result.get("status", "unknown")
|
|
state = result.get("state", result.get("status", ""))
|
|
|
|
# Terminal states
|
|
if state in ("success", "finished", "completed"):
|
|
logger.reason("SQL execution completed", {
|
|
"query_id": query_id,
|
|
"attempt": attempt + 1,
|
|
"rows_affected": result.get("rows"),
|
|
})
|
|
return {
|
|
"query_id": query_id,
|
|
"status": "success",
|
|
"state": state,
|
|
"rows_affected": result.get("rows"),
|
|
"error_message": result.get("error_message"),
|
|
"results": result.get("results"),
|
|
"completed_on": result.get("completed_on"),
|
|
}
|
|
|
|
if state in ("failed", "error", "stopped"):
|
|
error_msg = result.get("error_message", "Unknown error")
|
|
logger.explore("SQL execution failed", {
|
|
"query_id": query_id,
|
|
"state": state,
|
|
"error": error_msg,
|
|
})
|
|
return {
|
|
"query_id": query_id,
|
|
"status": "failed",
|
|
"state": state,
|
|
"error_message": error_msg,
|
|
"results": None,
|
|
}
|
|
|
|
if state in ("pending", "running", "started"):
|
|
await asyncio.sleep(poll_interval_seconds)
|
|
continue
|
|
|
|
# Unknown state — treat as still running
|
|
await asyncio.sleep(poll_interval_seconds)
|
|
|
|
except Exception as e:
|
|
logger.explore("Polling error, retrying", {
|
|
"query_id": query_id,
|
|
"attempt": attempt + 1,
|
|
"error": str(e),
|
|
})
|
|
await asyncio.sleep(poll_interval_seconds)
|
|
continue
|
|
|
|
# Timeout
|
|
logger.explore("SQL execution polling timed out", {
|
|
"query_id": query_id,
|
|
"max_polls": max_polls,
|
|
})
|
|
return {
|
|
"query_id": query_id,
|
|
"status": "timeout",
|
|
"error_message": f"Polling timed out after {max_polls} attempts",
|
|
"results": None,
|
|
}
|
|
# endregion poll_execution_status
|
|
|
|
# region execute_and_poll [TYPE Function]
|
|
# @PURPOSE Execute SQL and wait for completion. One-shot convenience method.
|
|
# @PRE sql is valid SQL.
|
|
# @POST Returns final execution result.
|
|
# @SIDE_EFFECT Makes HTTP calls to Superset API.
|
|
async def execute_and_poll(
|
|
self,
|
|
sql: str,
|
|
database_id: int | None = None,
|
|
max_polls: int = 60,
|
|
poll_interval_seconds: float = 2.0,
|
|
) -> dict[str, Any]:
|
|
with belief_scope("SupersetSqlLabExecutor.execute_and_poll"):
|
|
exec_result = await self.execute_sql(
|
|
sql=sql,
|
|
database_id=database_id,
|
|
run_async=False, # Sync mode — Superset returns result directly
|
|
)
|
|
|
|
status = exec_result.get("status", "unknown")
|
|
|
|
# Sync mode: response already has the final status and data
|
|
if status == "success":
|
|
logger.reason("SQL execution completed (sync)", {
|
|
"query_id": exec_result.get("query_id"),
|
|
})
|
|
return {
|
|
"query_id": exec_result.get("query_id"),
|
|
"status": "success",
|
|
"state": "success",
|
|
"rows_affected": exec_result.get("raw_response", {}).get("query", {}).get("rows"),
|
|
"error_message": None,
|
|
"results": exec_result.get("raw_response", {}).get("data"),
|
|
"completed_on": exec_result.get("raw_response", {}).get("query", {}).get("endDttm"),
|
|
}
|
|
|
|
# For async mode (fallback): poll for completion
|
|
query_id = exec_result.get("query_id")
|
|
if not query_id:
|
|
logger.explore("No query_id from SQL Lab execute", {
|
|
"raw_response": exec_result.get("raw_response"),
|
|
})
|
|
return {
|
|
"status": "failed",
|
|
"error_message": "No query_id returned from SQL Lab",
|
|
"query_id": None,
|
|
}
|
|
|
|
return await self.poll_execution_status(
|
|
query_id=str(query_id),
|
|
max_polls=max_polls,
|
|
poll_interval_seconds=poll_interval_seconds,
|
|
)
|
|
# endregion execute_and_poll
|
|
|
|
# region get_query_results [TYPE Function]
|
|
# @PURPOSE Fetch the results of a completed query.
|
|
# @PRE query_id is a valid Superset query ID.
|
|
# @POST Returns query results if available.
|
|
# @SIDE_EFFECT Makes HTTP GET to Superset.
|
|
async def get_query_results(self, query_id: str) -> dict[str, Any]:
|
|
with belief_scope("SupersetSqlLabExecutor.get_query_results"):
|
|
client = self._get_client()
|
|
try:
|
|
response = await client.network.request(
|
|
method="GET",
|
|
endpoint=f"/api/v1/query/{query_id}/results",
|
|
)
|
|
result = response if isinstance(response, dict) else {}
|
|
return {
|
|
"query_id": query_id,
|
|
"status": "success" if result.get("results") else "no_results",
|
|
"results": result.get("results"),
|
|
}
|
|
except Exception as e:
|
|
logger.explore("Failed to fetch query results", {
|
|
"query_id": query_id,
|
|
"error": str(e),
|
|
})
|
|
return {
|
|
"query_id": query_id,
|
|
"status": "failed",
|
|
"error_message": str(e),
|
|
"results": None,
|
|
}
|
|
# endregion get_query_results
|
|
|
|
|
|
# #endregion SupersetSqlLabExecutor
|
|
# #endregion SupersetSqlLabExecutor
|