fix(llm): add fetch-models endpoint, fix SQL Lab INSERT (client_id truncation, sync mode, target_column, timestamp normalization)
- 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
This commit is contained in:
@@ -32,6 +32,7 @@ class SupersetSqlLabExecutor:
|
||||
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.
|
||||
@@ -53,15 +54,43 @@ class SupersetSqlLabExecutor:
|
||||
# endregion _get_client
|
||||
|
||||
# region resolve_database_id [TYPE Function]
|
||||
# @PURPOSE: Resolve the target database ID from the environment.
|
||||
# @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.
|
||||
# @SIDE_EFFECT: Fetches databases list from Superset.
|
||||
def resolve_database_id(self, database_name: Optional[str] = None) -> int:
|
||||
# @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"]}
|
||||
query={"columns": ["id", "database_name", "backend"]}
|
||||
)
|
||||
if not databases:
|
||||
raise ValueError("No databases found in Superset environment")
|
||||
@@ -70,9 +99,11 @@ class SupersetSqlLabExecutor:
|
||||
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(
|
||||
@@ -81,13 +112,22 @@ class SupersetSqlLabExecutor:
|
||||
|
||||
# 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.
|
||||
@@ -117,27 +157,75 @@ class SupersetSqlLabExecutor:
|
||||
"runAsync": run_async,
|
||||
"schema": None,
|
||||
"tab": "translation-insert",
|
||||
"client_id": f"trl-{uuid.uuid4().hex[:7]}", # max 11 chars for Superset varchar(11)
|
||||
}
|
||||
|
||||
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}")
|
||||
# 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
|
||||
|
||||
# Parse response
|
||||
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 {}
|
||||
query_id = result.get("query_id") or result.get("id")
|
||||
|
||||
# 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 {
|
||||
@@ -174,7 +262,9 @@ class SupersetSqlLabExecutor:
|
||||
method="GET",
|
||||
endpoint=f"/api/v1/query/{query_id}",
|
||||
)
|
||||
result = response if isinstance(response, dict) else {}
|
||||
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", ""))
|
||||
|
||||
@@ -255,9 +345,27 @@ class SupersetSqlLabExecutor:
|
||||
exec_result = self.execute_sql(
|
||||
sql=sql,
|
||||
database_id=database_id,
|
||||
run_async=True,
|
||||
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", {
|
||||
|
||||
Reference in New Issue
Block a user