busya
10803a5019
- Add POST /api/llm/providers/fetch-models route with LLMClient.fetch_models() - Add target_column to TranslationJob model/schema/service/orchestrator - Fix SQL Lab execute: truncate client_id to 11 chars (varchar(11)) - Switch SQL Lab to sync mode (runAsync: false) — no Celery workers - Fix polling: unwrap nested result from Superset query API - Fix ClickHouse timestamp: normalize float timestamps to YYYY-MM-DD
422 lines
18 KiB
Python
422 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: 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
|
|
self._database_backend: Optional[str] = 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.
|
|
def resolve_database_id(
|
|
self,
|
|
database_name: Optional[str] = None,
|
|
target_database_id: Optional[str] = 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 = 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")
|
|
logger.reason("Resolved database ID by target_database_id", {
|
|
"target_database_id": target_database_id,
|
|
"database_id": self._database_id,
|
|
"backend": self._database_backend,
|
|
})
|
|
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 = 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")
|
|
logger.reason("Resolved database ID by name", {
|
|
"database_name": database_name,
|
|
"database_id": self._database_id,
|
|
"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")
|
|
logger.reason("Using default database", {
|
|
"database_name": databases[0].get("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) -> Optional[str]:
|
|
return self._database_backend
|
|
# endregion get_database_backend
|
|
|
|
# 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",
|
|
"client_id": f"trl-{uuid.uuid4().hex[:7]}", # max 11 chars for Superset varchar(11)
|
|
}
|
|
|
|
# Try multiple SQL Lab endpoints — some Superset versions use /sqllab/execute/,
|
|
# others use /sql_lab/execute/ (with underscore)
|
|
candidate_endpoints = ["/api/v1/sqllab/execute/", "/api/v1/sql_lab/execute/"]
|
|
response = None
|
|
last_error = None
|
|
for endpoint in candidate_endpoints:
|
|
try:
|
|
response = client.network.request(
|
|
method="POST",
|
|
endpoint=endpoint,
|
|
data=json.dumps(payload),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
if isinstance(response, dict) and response:
|
|
logger.reason(f"SQL Lab endpoint succeeded", extra={
|
|
"endpoint": endpoint,
|
|
})
|
|
break
|
|
except Exception as ep_err:
|
|
last_error = ep_err
|
|
logger.explore(f"SQL Lab endpoint failed, trying next", extra={
|
|
"error": str(ep_err),
|
|
"endpoint": endpoint,
|
|
})
|
|
response = None
|
|
|
|
if response is None:
|
|
error_msg = f"All SQL Lab endpoints failed. Last error: {last_error}"
|
|
logger.explore("SQL Lab execute failed", extra={"error": error_msg})
|
|
raise ValueError(f"Superset SQL Lab execute failed: {error_msg}")
|
|
|
|
# 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.
|
|
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}",
|
|
)
|
|
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"):
|
|
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=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 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
|