314 lines
13 KiB
Python
314 lines
13 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: Infra
|
|
# @RELATION DEPENDS_ON -> [SupersetClient]
|
|
# @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.
|
|
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
from datetime import datetime, timezone
|
|
import json
|
|
import time
|
|
import uuid
|
|
|
|
from ...core.logger import logger, belief_scope
|
|
from ...core.config_manager import ConfigManager
|
|
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: Optional[SupersetClient] = None
|
|
self._database_id: Optional[int] = 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.
|
|
# @PRE: database_name or database_id should be known.
|
|
# @POST: Returns database_id integer or raises ValueError.
|
|
# @SIDE_EFFECT: Fetches databases list from Superset.
|
|
def resolve_database_id(self, database_name: Optional[str] = None) -> int:
|
|
with belief_scope("SupersetSqlLabExecutor.resolve_database_id"):
|
|
client = self._get_client()
|
|
_, databases = client.get_databases(
|
|
query={"columns": ["id", "database_name"]}
|
|
)
|
|
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"]
|
|
logger.reason("Resolved database ID by name", {
|
|
"database_name": database_name,
|
|
"database_id": self._database_id,
|
|
})
|
|
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"]
|
|
logger.reason("Using default database", {
|
|
"database_name": databases[0].get("database_name"),
|
|
"database_id": self._database_id,
|
|
})
|
|
return self._database_id
|
|
# endregion resolve_database_id
|
|
|
|
# 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/.
|
|
def execute_sql(
|
|
self,
|
|
sql: str,
|
|
database_id: Optional[int] = 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 = 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",
|
|
}
|
|
|
|
try:
|
|
response = client.network.request(
|
|
method="POST",
|
|
endpoint="/api/v1/sqllab/execute/",
|
|
data=json.dumps(payload),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
except Exception as e:
|
|
logger.explore("SQL Lab execute failed", {"error": str(e)})
|
|
raise ValueError(f"Superset SQL Lab execute failed: {e}")
|
|
|
|
# Parse response
|
|
result = response if isinstance(response, dict) else {}
|
|
query_id = result.get("query_id") or result.get("id")
|
|
status = result.get("status", "unknown")
|
|
|
|
logger.reason("SQL Lab execute response", {
|
|
"query_id": query_id,
|
|
"status": status,
|
|
})
|
|
|
|
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.
|
|
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 = client.network.request(
|
|
method="GET",
|
|
endpoint=f"/api/v1/query/{query_id}",
|
|
)
|
|
result = response if isinstance(response, dict) else {}
|
|
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"):
|
|
time.sleep(poll_interval_seconds)
|
|
continue
|
|
|
|
# Unknown state — treat as still running
|
|
time.sleep(poll_interval_seconds)
|
|
|
|
except Exception as e:
|
|
logger.explore("Polling error, retrying", {
|
|
"query_id": query_id,
|
|
"attempt": attempt + 1,
|
|
"error": str(e),
|
|
})
|
|
time.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.
|
|
def execute_and_poll(
|
|
self,
|
|
sql: str,
|
|
database_id: Optional[int] = None,
|
|
max_polls: int = 60,
|
|
poll_interval_seconds: float = 2.0,
|
|
) -> Dict[str, Any]:
|
|
with belief_scope("SupersetSqlLabExecutor.execute_and_poll"):
|
|
exec_result = self.execute_sql(
|
|
sql=sql,
|
|
database_id=database_id,
|
|
run_async=True,
|
|
)
|
|
|
|
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 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.
|
|
def get_query_results(self, query_id: str) -> Dict[str, Any]:
|
|
with belief_scope("SupersetSqlLabExecutor.get_query_results"):
|
|
client = self._get_client()
|
|
try:
|
|
response = 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
|