feat(agent-superset): extend SupersetClient with agent-critical methods + DDL/DML guard + tests (126 passed)

Ported from mcp-superset research module, integrated into existing async SupersetClient:

New mixins (4 files):
- safety.py: DDL/DML guard (13 keywords), comment/string stripping, viz_type validation
- _sql_lab.py: execute_sql, format_sql, results, estimate, CSV export, query history
- _saved_queries.py: saved queries CRUD (5 methods)
- _audit.py: permissions audit matrix (user × dashboard × dataset × RLS)

Extended mixins (4 files):
- _dashboards_write.py: +create, +update (general), +copy, +publish, +unpublish
- _dashboards_crud.py: +standalone get_dashboard_charts/datasets, +get_dashboard
- _datasets.py: +create, +delete, +duplicate, +refresh_schema, +get_or_create, +export/import
- _databases.py: +update, +test_connection, +schemas/tables/catalogs, +validate_sql, +select_star, +table_metadata

FastAPI proxy (2 files):
- agent_superset.py (284L): SQL Lab + Dashboards/Datasets write endpoints
- agent_superset_explore.py (240L): DB explore + Audit + Saved Queries endpoints

Agent tools (2 files):
- tools.py: +7 LangChain @tools (24 total), intent-routing keywords
- _tool_resolver.py: updated SAFE/GUARDED/FAST_CONFIRM classification sets

Tests (3 files, 126 tests):
- test_superset_safety.py (51): DDL/DML bypass/legitimate/safe scenarios
- test_superset_extended.py (42): all mixin methods with mock AsyncAPIClient
- test_superset_tools.py (33): agent tool registration, intent matching, @tool .ainvoke()
This commit is contained in:
2026-06-30 17:26:51 +03:00
parent a0619ca049
commit 79d0f8f678
18 changed files with 3993 additions and 0 deletions

View File

@@ -23,6 +23,11 @@ _SAFE_AGENT_TOOLS = {
"list_llm_providers",
"get_llm_status",
"list_maintenance_events",
# NEW: read-only Superset tools
"superset_execute_sql",
"superset_explore_database",
"superset_audit_permissions",
"superset_format_sql",
}
_GUARDED_AGENT_TOOLS = {
"create_branch",
@@ -33,6 +38,10 @@ _GUARDED_AGENT_TOOLS = {
"run_llm_documentation",
"start_maintenance",
"end_maintenance",
# NEW: guarded Superset write operations
"superset_create_dashboard",
"superset_copy_dashboard",
"superset_create_dataset",
}
_DANGEROUS_AGENT_TOOLS = {
"deploy_dashboard",
@@ -44,6 +53,9 @@ _FAST_CONFIRM_TOOLS = {
"list_llm_providers",
"get_llm_status",
"list_maintenance_events",
"superset_explore_database",
"superset_audit_permissions",
"superset_format_sql",
}
# #endregion AgentChat.ToolResolver.Sets

View File

@@ -759,6 +759,272 @@ async def end_maintenance(
# #endregion AgentChat.Tools.EndMaintenance
# ═══════════════════════════════════════════════════════════════════
# NEW: Agent-critical Superset operations (Phase 4)
# ═══════════════════════════════════════════════════════════════════
# #region AgentChat.Tools.SupersetSql.Schema [C:1] [TYPE Class] [SEMANTICS agent-chat,tools,schema,superset,sql]
class ExecuteSupersetSqlInput(BaseModel):
environment_id: str = Field(description="Target Superset environment ID")
database_id: int = Field(description="Database connection ID")
sql: str = Field(description="SQL SELECT query to execute")
query_schema: str | None = Field(default=None, description="Optional schema for the query")
# #endregion AgentChat.Tools.SupersetSql.Schema
# #region AgentChat.Tools.SupersetSql [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,superset,sql]
# @ingroup AgentChat
# @BRIEF Execute a read-only SQL query in Superset SQL Lab (DDL/DML blocked).
# @PRE User authenticated via dual-identity JWT. Database exists in target environment.
# @POST Returns query results dict or error with blocked keyword.
# @SIDE_EFFECT HTTP POST to FastAPI /api/agent/superset/sqllab/execute.
@tool(args_schema=ExecuteSupersetSqlInput)
async def superset_execute_sql(environment_id: str, database_id: int, sql: str, query_schema: str | None = None) -> str:
"""Execute a read-only SQL query in Superset SQL Lab (SELECT only, DDL/DML blocked)."""
logger.reason("Execute Superset SQL", payload={"environment_id": environment_id, "database_id": database_id},
extra={"src": "AgentChat.Tools.SupersetSql"})
params: dict[str, Any] = {
"environment_id": environment_id,
"database_id": str(database_id),
"sql": sql,
}
if query_schema:
params["schema"] = query_schema
resp = await _post("/api/agent/superset/sqllab/execute", params=params)
result = _api_result(resp)
logger.reflect("Superset SQL executed", payload={"status": resp.status_code},
extra={"src": "AgentChat.Tools.SupersetSql"})
return result
# #endregion AgentChat.Tools.SupersetSql
# #region AgentChat.Tools.SupersetDatabases.Schema [C:1] [TYPE Class] [SEMANTICS agent-chat,tools,schema,superset,database]
class ListDatabaseSchemasInput(BaseModel):
environment_id: str = Field(description="Target Superset environment ID")
database_id: int = Field(description="Database connection ID")
schema_name: str | None = Field(default=None, description="Optional schema to list tables for")
table_name: str | None = Field(default=None, description="Optional table to get metadata/select_star for")
action: str = Field(default="schemas", description="Action: schemas, tables, table_metadata, select_star, validate_sql")
# #endregion AgentChat.Tools.SupersetDatabases.Schema
# #region AgentChat.Tools.SupersetDatabases [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,superset,database,explore]
# @ingroup AgentChat
# @BRIEF Explore database structure: list schemas, tables, get table metadata, generate SELECT *.
# @PRE User authenticated via dual-identity JWT. Database exists.
# @POST Returns schema/table/metadata info.
# @SIDE_EFFECT HTTP GET/POST to FastAPI /api/agent/superset/databases/{id}/...
@tool(args_schema=ListDatabaseSchemasInput)
async def superset_explore_database(
environment_id: str,
database_id: int,
schema_name: str | None = None,
table_name: str | None = None,
action: str = "schemas",
) -> str:
"""Explore database: list schemas, tables, get table metadata, or generate SELECT * template."""
logger.reason("Explore Superset database", payload={"environment_id": environment_id, "database_id": database_id, "action": action},
extra={"src": "AgentChat.Tools.SupersetDatabases"})
params: dict[str, Any] = {"environment_id": environment_id}
if action == "schemas":
path = f"/api/agent/superset/databases/{database_id}/schemas"
elif action == "tables":
if not schema_name:
return "Error: schema_name is required for tables action."
path = f"/api/agent/superset/databases/{database_id}/tables"
params["schema_name"] = schema_name
elif action == "table_metadata":
if not table_name:
return "Error: table_name is required for table_metadata action."
path = f"/api/agent/superset/databases/{database_id}/table_metadata"
params["table_name"] = table_name
if schema_name:
params["schema_name"] = schema_name
elif action == "select_star":
if not table_name:
return "Error: table_name is required for select_star action."
path = f"/api/agent/superset/databases/{database_id}/select_star"
params["table_name"] = table_name
if schema_name:
params["schema_name"] = schema_name
else:
return f"Error: unknown action '{action}'. Use: schemas, tables, table_metadata, select_star."
resp = await _get(path, params=params)
result = _api_result(resp)
logger.reflect("Database explored", payload={"status": resp.status_code, "action": action},
extra={"src": "AgentChat.Tools.SupersetDatabases"})
return result
# #endregion AgentChat.Tools.SupersetDatabases
# #region AgentChat.Tools.SupersetAudit.Schema [C:1] [TYPE Class] [SEMANTICS agent-chat,tools,schema,superset,audit]
class AuditPermissionsInput(BaseModel):
environment_id: str = Field(description="Target Superset environment ID")
username_filter: str | None = Field(default=None, description="Optional filter by username")
include_admin: bool = Field(default=False, description="Include Admin users in the audit")
# #endregion AgentChat.Tools.SupersetAudit.Schema
# #region AgentChat.Tools.SupersetAudit [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,superset,audit,permissions]
# @ingroup AgentChat
# @BRIEF Audit access rights: user × dashboard × dataset × RLS permission matrix.
# @PRE User authenticated via dual-identity JWT.
# @POST Returns audit report with per-user dashboard/dataset access and RLS regions.
# @SIDE_EFFECT HTTP GET to FastAPI /api/agent/superset/audit/permissions.
@tool(args_schema=AuditPermissionsInput)
async def superset_audit_permissions(
environment_id: str,
username_filter: str | None = None,
include_admin: bool = False,
) -> str:
"""Audit access rights: user × dashboard × dataset × RLS permission matrix."""
logger.reason("Audit Superset permissions", payload={"environment_id": environment_id},
extra={"src": "AgentChat.Tools.SupersetAudit"})
params: dict[str, Any] = {"environment_id": environment_id}
if username_filter:
params["username_filter"] = username_filter
if include_admin:
params["include_admin"] = "true"
resp = await _get("/api/agent/superset/audit/permissions", params=params)
result = _api_result(resp)
logger.reflect("Permissions audit completed", payload={"status": resp.status_code},
extra={"src": "AgentChat.Tools.SupersetAudit"})
return result
# #endregion AgentChat.Tools.SupersetAudit
# #region AgentChat.Tools.SupersetCreateDashboard.Schema [C:1] [TYPE Class] [SEMANTICS agent-chat,tools,schema,superset,dashboard,create]
class CreateSupersetDashboardInput(BaseModel):
environment_id: str = Field(description="Target Superset environment ID")
dashboard_title: str = Field(description="Dashboard title")
slug: str | None = Field(default=None, description="Optional URL slug")
published: bool = Field(default=False, description="Whether the dashboard is published")
# #endregion AgentChat.Tools.SupersetCreateDashboard.Schema
# #region AgentChat.Tools.SupersetCreateDashboard [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,superset,dashboard,create]
# @ingroup AgentChat
# @BRIEF Create a new dashboard in Superset.
# @PRE User authenticated via dual-identity JWT.
# @POST Returns created dashboard dict.
# @SIDE_EFFECT HTTP POST to FastAPI /api/agent/superset/dashboards.
@tool(args_schema=CreateSupersetDashboardInput)
async def superset_create_dashboard(
environment_id: str,
dashboard_title: str,
slug: str | None = None,
published: bool = False,
) -> str:
"""Create a new dashboard in Superset."""
logger.reason("Create Superset dashboard", payload={"title": dashboard_title, "environment_id": environment_id},
extra={"src": "AgentChat.Tools.SupersetCreateDashboard"})
params: dict[str, Any] = {
"environment_id": environment_id,
"dashboard_title": dashboard_title,
"published": "true" if published else "false",
}
if slug:
params["slug"] = slug
resp = await _post("/api/agent/superset/dashboards", params=params)
result = _api_result(resp, {200, 201})
logger.reflect("Dashboard created", payload={"status": resp.status_code},
extra={"src": "AgentChat.Tools.SupersetCreateDashboard"})
return result
# #endregion AgentChat.Tools.SupersetCreateDashboard
# #region AgentChat.Tools.SupersetCopyDashboard.Schema [C:1] [TYPE Class] [SEMANTICS agent-chat,tools,schema,superset,dashboard,copy]
class CopySupersetDashboardInput(BaseModel):
environment_id: str = Field(description="Target Superset environment ID")
dashboard_id: int = Field(description="Source dashboard ID to copy")
dashboard_title: str | None = Field(default=None, description="Optional title for the copy")
# #endregion AgentChat.Tools.SupersetCopyDashboard.Schema
# #region AgentChat.Tools.SupersetCopyDashboard [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,superset,dashboard,copy]
# @ingroup AgentChat
# @BRIEF Deep-copy a Superset dashboard including all charts.
# @PRE User authenticated via dual-identity JWT.
# @POST Returns copy result dict.
# @SIDE_EFFECT HTTP POST to FastAPI /api/agent/superset/dashboards/{id}/copy.
@tool(args_schema=CopySupersetDashboardInput)
async def superset_copy_dashboard(
environment_id: str,
dashboard_id: int,
dashboard_title: str | None = None,
) -> str:
"""Deep-copy a dashboard including all charts."""
logger.reason("Copy Superset dashboard", payload={"dashboard_id": dashboard_id, "environment_id": environment_id},
extra={"src": "AgentChat.Tools.SupersetCopyDashboard"})
params: dict[str, Any] = {"environment_id": environment_id}
if dashboard_title:
params["dashboard_title"] = dashboard_title
resp = await _post(f"/api/agent/superset/dashboards/{dashboard_id}/copy", params=params)
result = _api_result(resp, {200, 201})
logger.reflect("Dashboard copied", payload={"status": resp.status_code, "source_id": dashboard_id},
extra={"src": "AgentChat.Tools.SupersetCopyDashboard"})
return result
# #endregion AgentChat.Tools.SupersetCopyDashboard
# #region AgentChat.Tools.SupersetCreateDataset.Schema [C:1] [TYPE Class] [SEMANTICS agent-chat,tools,schema,superset,dataset,create]
class CreateSupersetDatasetInput(BaseModel):
environment_id: str = Field(description="Target Superset environment ID")
table_name: str = Field(description="Table or view name")
database: int = Field(description="Database connection ID")
schema_name: str | None = Field(default=None, description="Optional schema name")
# #endregion AgentChat.Tools.SupersetCreateDataset.Schema
# #region AgentChat.Tools.SupersetCreateDataset [C:3] [TYPE Function] [SEMANTICS agent-chat,tools,superset,dataset,create]
# @ingroup AgentChat
# @BRIEF Create a new dataset in Superset.
# @PRE User authenticated via dual-identity JWT.
# @POST Returns created dataset dict.
# @SIDE_EFFECT HTTP POST to FastAPI /api/agent/superset/datasets.
@tool(args_schema=CreateSupersetDatasetInput)
async def superset_create_dataset(
environment_id: str,
table_name: str,
database: int,
schema_name: str | None = None,
) -> str:
"""Create a new dataset in Superset."""
logger.reason("Create Superset dataset", payload={"table_name": table_name, "environment_id": environment_id},
extra={"src": "AgentChat.Tools.SupersetCreateDataset"})
params: dict[str, Any] = {
"environment_id": environment_id,
"table_name": table_name,
"database": str(database),
}
if schema_name:
params["schema_name"] = schema_name
resp = await _post("/api/agent/superset/datasets", params=params)
result = _api_result(resp, {200, 201})
logger.reflect("Dataset created", payload={"status": resp.status_code},
extra={"src": "AgentChat.Tools.SupersetCreateDataset"})
return result
# #endregion AgentChat.Tools.SupersetCreateDataset
# #region AgentChat.Tools.SupersetFormatSql [C:2] [TYPE Function] [SEMANTICS agent-chat,tools,superset,sql,format]
# @ingroup AgentChat
# @BRIEF Format/pretty-print a SQL query using Superset SQL Lab.
# @SIDE_EFFECT HTTP POST to FastAPI /api/agent/superset/sqllab/format.
@tool
async def superset_format_sql(environment_id: str, sql: str) -> str:
"""Format/pretty-print a SQL query using Superset SQL Lab."""
logger.reason("Format SQL", extra={"src": "AgentChat.Tools.SupersetFormatSql"})
resp = await _post("/api/agent/superset/sqllab/format", params={
"environment_id": environment_id,
"sql": sql,
})
result = _api_result(resp)
logger.reflect("SQL formatted", extra={"src": "AgentChat.Tools.SupersetFormatSql"})
return result
# #endregion AgentChat.Tools.SupersetFormatSql
# ═══════════════════════════════════════════════════════════════════
# Tool registry
# ═══════════════════════════════════════════════════════════════════
@@ -785,6 +1051,14 @@ def get_all_tools() -> list:
list_maintenance_events,
start_maintenance,
end_maintenance,
# NEW: Superset direct operations
superset_execute_sql,
superset_explore_database,
superset_audit_permissions,
superset_create_dashboard,
superset_copy_dashboard,
superset_create_dataset,
superset_format_sql,
]
# #endregion AgentChat.Tools.GetAll
@@ -842,6 +1116,29 @@ def get_tools_for_query(query: str, *, prefetch_available: bool = False) -> list
if any(word in text for word in ["maintenance", "обслуж", "баннер"]):
matched_intent = True
selected.extend([list_maintenance_events, start_maintenance, end_maintenance])
# NEW: Superset direct tools intent matching
if any(word in text for word in ["sql", "запрос", "select", "query"]):
matched_intent = True
selected.append(superset_execute_sql)
if any(word in text for word in ["форматировать sql", "format sql", "формат sql"]):
matched_intent = True
selected.append(superset_format_sql)
if any(word in text for word in ["схем", "schema", "таблиц", "table", "колонк", "column",
"select star", "метаданные", "metadata"]):
matched_intent = True
selected.append(superset_explore_database)
if any(word in text for word in ["аудит", "audit", "прав", "permission", "доступ", "access"]):
matched_intent = True
selected.append(superset_audit_permissions)
if any(word in text for word in ["создать дашборд", "create dashboard", "новый дашборд", "new dashboard"]):
matched_intent = True
selected.append(superset_create_dashboard)
if any(word in text for word in ["копировать дашборд", "copy dashboard", "дублировать дашборд"]):
matched_intent = True
selected.append(superset_copy_dashboard)
if any(word in text for word in ["создать датасет", "create dataset", "новый датасет", "new dataset"]):
matched_intent = True
selected.append(superset_create_dataset)
if len(selected) == 1 and not matched_intent:
selected.extend([search_dashboards, get_health_summary, list_environments, get_task_status])

View File

@@ -21,6 +21,9 @@
__all__ = [
"admin",
"admin_api_keys",
"agent_conversations",
"agent_superset",
"agent_superset_explore",
"assistant",
"clean_release",
"clean_release_v2",

View File

@@ -0,0 +1,284 @@
# #region AgentSupersetRoutes [C:4] [TYPE Module] [SEMANTICS api,agent,superset,sql,dashboard,dataset,crud]
# @defgroup Api Agent Superset proxy routes for the Gradio agent (write/mutate endpoints).
# @BRIEF FastAPI endpoints proxying SupersetClient write/mutate operations for the agent chat.
# @LAYER API
# @RELATION DEPENDS_ON -> [SupersetClient]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API через SupersetClient.
# @RATIONALE The Gradio agent container has no direct SupersetClient — it reaches Superset
# through these proxy endpoints. Each endpoint resolves the target environment and delegates
# to the appropriate SupersetClient method. Dual-identity JWT (service + user) is enforced
# per FR-007/FR-019.
# @INVARIANT Module stays under 400 lines per INV_7. Read-only explore/audit endpoints live in
# agent_superset_explore.py.
from fastapi import APIRouter, HTTPException, Query
from src.core.config_models import Environment
from src.core.superset_client import SupersetClient
router = APIRouter(prefix="/api/agent/superset", tags=["Agent Superset"])
# ── Helper: resolve SupersetClient for an environment ───────────
async def _get_superset_client(environment_id: str) -> SupersetClient:
"""Resolve a SupersetClient for the given environment_id."""
from ..api.routes.environments import _get_environment_by_id as _resolve_env
env_data = await _resolve_env(environment_id)
if not env_data:
raise HTTPException(status_code=404, detail=f"Environment '{environment_id}' not found.")
env = Environment(**env_data)
return SupersetClient(env)
# ═══════════════════════════════════════════════════════════════════
# SQL Lab
# ═══════════════════════════════════════════════════════════════════
# #region AgentSuperset.SqlExecute [C:3] [TYPE Endpoint] [SEMANTICS api,agent,sql,execute]
# @ingroup Api
# @BRIEF Execute a read-only SQL query in Superset SQL Lab (DDL/DML blocked).
# @SIDE_EFFECT HTTP POST to Superset /api/v1/sqllab/execute/.
# @RELATION CALLS -> [SupersetSqlLab.ExecuteSql]
@router.post("/sqllab/execute")
async def agent_sqllab_execute(
environment_id: str = Query(..., description="Target Superset environment ID"),
database_id: int = Query(..., description="Database connection ID"),
sql: str = Query(..., description="SQL query (SELECT only)"),
schema: str | None = Query(None),
catalog: str | None = Query(None),
tab_name: str | None = Query(None),
template_params: str | None = Query(None),
) -> dict:
"""Execute a read-only SQL query in Superset SQL Lab."""
client = await _get_superset_client(environment_id)
try:
return await client.execute_sql(
database_id=database_id,
sql=sql,
schema=schema,
catalog=catalog,
tab_name=tab_name,
template_params=template_params,
)
finally:
await client.aclose()
# #endregion AgentSuperset.SqlExecute
# #region AgentSuperset.SqlFormat [C:2] [TYPE Endpoint] [SEMANTICS api,agent,sql,format]
# @ingroup Api
# @BRIEF Format/pretty-print a SQL query using Superset SQL Lab.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/sqllab/format_sql/.
# @RELATION CALLS -> [SupersetSqlLab.FormatSql]
@router.post("/sqllab/format")
async def agent_sqllab_format(
environment_id: str = Query(...),
sql: str = Query(..., description="SQL to format"),
) -> dict:
"""Format/pretty-print a SQL query."""
client = await _get_superset_client(environment_id)
try:
formatted = await client.format_sql(sql)
return {"result": formatted}
finally:
await client.aclose()
# #endregion AgentSuperset.SqlFormat
# #region AgentSuperset.SqlEstimate [C:2] [TYPE Endpoint] [SEMANTICS api,agent,sql,estimate]
# @ingroup Api
# @BRIEF Estimate the cost of a SQL query via Superset SQL Lab.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/sqllab/estimate/.
# @RELATION CALLS -> [SupersetSqlLab.EstimateCost]
@router.post("/sqllab/estimate")
async def agent_sqllab_estimate(
environment_id: str = Query(...),
database_id: int = Query(...),
sql: str = Query(...),
schema: str | None = Query(None),
) -> dict:
"""Estimate the cost of a SQL query."""
client = await _get_superset_client(environment_id)
try:
return await client.estimate_sql_cost(database_id=database_id, sql=sql, schema=schema)
finally:
await client.aclose()
# #endregion AgentSuperset.SqlEstimate
# ═══════════════════════════════════════════════════════════════════
# Dashboards — Write (HITL-guarded)
# ═══════════════════════════════════════════════════════════════════
# #region AgentSuperset.DashboardCreate [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dashboard,create]
# @ingroup Api
# @BRIEF Create a new dashboard in Superset.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/dashboard/.
# @RELATION CALLS -> [SupersetDashboardsWriteMixin.create_dashboard]
@router.post("/dashboards")
async def agent_dashboard_create(
environment_id: str = Query(...),
dashboard_title: str = Query(...),
slug: str | None = Query(None),
published: bool = Query(False),
json_metadata: str | None = Query(None),
css: str | None = Query(None),
position_json: str | None = Query(None),
) -> dict:
"""Create a new dashboard in Superset."""
client = await _get_superset_client(environment_id)
try:
return await client.create_dashboard(
dashboard_title=dashboard_title,
slug=slug,
published=published,
json_metadata=json_metadata,
css=css,
position_json=position_json,
)
finally:
await client.aclose()
# #endregion AgentSuperset.DashboardCreate
# #region AgentSuperset.DashboardCopy [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dashboard,copy]
# @ingroup Api
# @BRIEF Deep-copy a dashboard including all charts.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/dashboard/{id}/copy/.
# @RELATION CALLS -> [SupersetDashboardsWriteMixin.copy_dashboard]
@router.post("/dashboards/{dashboard_id}/copy")
async def agent_dashboard_copy(
dashboard_id: int,
environment_id: str = Query(...),
dashboard_title: str | None = Query(None),
) -> dict:
"""Deep-copy a dashboard including all charts."""
client = await _get_superset_client(environment_id)
try:
return await client.copy_dashboard(dashboard_id=dashboard_id, dashboard_title=dashboard_title)
finally:
await client.aclose()
# #endregion AgentSuperset.DashboardCopy
# #region AgentSuperset.DashboardUpdate [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dashboard,update]
# @ingroup Api
# @BRIEF Update an existing dashboard's properties.
# @SIDE_EFFECT HTTP PUT to Superset /api/v1/dashboard/{id}.
# @RELATION CALLS -> [SupersetDashboardsWriteMixin.update_dashboard]
@router.put("/dashboards/{dashboard_id}")
async def agent_dashboard_update(
dashboard_id: int,
environment_id: str = Query(...),
dashboard_title: str | None = Query(None),
slug: str | None = Query(None),
published: bool | None = Query(None),
css: str | None = Query(None),
json_metadata: str | None = Query(None),
) -> dict:
"""Update an existing dashboard's properties."""
client = await _get_superset_client(environment_id)
try:
return await client.update_dashboard(
dashboard_id=dashboard_id,
dashboard_title=dashboard_title,
slug=slug,
published=published,
css=css,
json_metadata=json_metadata,
)
finally:
await client.aclose()
# #endregion AgentSuperset.DashboardUpdate
# ═══════════════════════════════════════════════════════════════════
# Datasets — Write (HITL-guarded)
# ═══════════════════════════════════════════════════════════════════
# #region AgentSuperset.DatasetCreate [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dataset,create]
# @ingroup Api
# @BRIEF Create a new dataset in Superset.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/dataset/.
# @RELATION CALLS -> [SupersetClient.CreateDataset]
@router.post("/datasets")
async def agent_dataset_create(
environment_id: str = Query(...),
table_name: str = Query(...),
database: int = Query(...),
schema_name: str | None = Query(None),
sql: str | None = Query(None),
) -> dict:
"""Create a new dataset in Superset."""
client = await _get_superset_client(environment_id)
try:
return await client.create_dataset(
table_name=table_name,
database=database,
schema_name=schema_name,
sql=sql,
)
finally:
await client.aclose()
# #endregion AgentSuperset.DatasetCreate
# #region AgentSuperset.DatasetDelete [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dataset,delete]
# @ingroup Api
# @BRIEF Delete a dataset from Superset by ID.
# @SIDE_EFFECT HTTP DELETE to Superset /api/v1/dataset/{id}.
# @RELATION CALLS -> [SupersetClient.DeleteDataset]
@router.delete("/datasets/{dataset_id}")
async def agent_dataset_delete(
dataset_id: int,
environment_id: str = Query(...),
) -> dict:
"""Delete a dataset from Superset."""
client = await _get_superset_client(environment_id)
try:
return await client.delete_dataset(dataset_id=dataset_id)
finally:
await client.aclose()
# #endregion AgentSuperset.DatasetDelete
# #region AgentSuperset.DatasetDuplicate [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dataset,duplicate]
# @ingroup Api
# @BRIEF Duplicate a dataset including columns and metrics.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/dataset/duplicate.
# @RELATION CALLS -> [SupersetClient.DuplicateDataset]
@router.post("/datasets/{dataset_id}/duplicate")
async def agent_dataset_duplicate(
dataset_id: int,
environment_id: str = Query(...),
table_name: str = Query(...),
) -> dict:
"""Duplicate a dataset including columns and metrics."""
client = await _get_superset_client(environment_id)
try:
return await client.duplicate_dataset(base_model_id=dataset_id, table_name=table_name)
finally:
await client.aclose()
# #endregion AgentSuperset.DatasetDuplicate
# #region AgentSuperset.DatasetRefresh [C:3] [TYPE Endpoint] [SEMANTICS api,agent,dataset,refresh]
# @ingroup Api
# @BRIEF Rescan columns and types for a dataset from its source database.
# @SIDE_EFFECT HTTP PUT to Superset /api/v1/dataset/{id}/refresh.
# @RELATION CALLS -> [SupersetClient.RefreshDatasetSchema]
@router.post("/datasets/{dataset_id}/refresh")
async def agent_dataset_refresh(
dataset_id: int,
environment_id: str = Query(...),
) -> dict:
"""Rescan columns and types for a dataset from its source database."""
client = await _get_superset_client(environment_id)
try:
return await client.refresh_dataset_schema(dataset_id=dataset_id)
finally:
await client.aclose()
# #endregion AgentSuperset.DatasetRefresh
# #endregion AgentSupersetRoutes

View File

@@ -0,0 +1,240 @@
# #region AgentSupersetExploreRoutes [C:3] [TYPE Module] [SEMANTICS api,agent,superset,database,explore,audit]
# @defgroup Api Agent Superset read-only proxy routes (database exploration, audit, saved queries).
# @BRIEF FastAPI endpoints proxying SupersetClient read-only operations for the agent chat.
# @LAYER API
# @RELATION DEPENDS_ON -> [SupersetClient]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API через SupersetClient.
# @RATIONALE Split from agent_superset.py to satisfy INV_7 (module < 400 lines).
# Write/mutate endpoints remain in agent_superset.py.
from fastapi import APIRouter, HTTPException, Query
from src.core.config_models import Environment
from src.core.superset_client import SupersetClient
router = APIRouter(prefix="/api/agent/superset", tags=["Agent Superset"])
async def _get_superset_client(environment_id: str) -> SupersetClient:
"""Resolve a SupersetClient for the given environment_id."""
from ..api.routes.environments import _get_environment_by_id as _resolve_env
env_data = await _resolve_env(environment_id)
if not env_data:
raise HTTPException(status_code=404, detail=f"Environment '{environment_id}' not found.")
env = Environment(**env_data)
return SupersetClient(env)
# ═══════════════════════════════════════════════════════════════════
# Databases — Read-only (schema/table exploration)
# ═══════════════════════════════════════════════════════════════════
# #region AgentSuperset.DatabaseSchemas [C:2] [TYPE Endpoint] [SEMANTICS api,agent,database,schemas]
# @ingroup Api
# @BRIEF List all schemas for a database.
# @SIDE_EFFECT HTTP GET to Superset /api/v1/database/{id}/schemas/.
# @RELATION CALLS -> [SupersetDatabasesMixin.get_database_schemas]
@router.get("/databases/{database_id}/schemas")
async def agent_database_schemas(
database_id: int,
environment_id: str = Query(...),
) -> list:
"""List all schemas for a database."""
client = await _get_superset_client(environment_id)
try:
return await client.get_database_schemas(database_id=database_id)
finally:
await client.aclose()
# #endregion AgentSuperset.DatabaseSchemas
# #region AgentSuperset.DatabaseTables [C:2] [TYPE Endpoint] [SEMANTICS api,agent,database,tables]
# @ingroup Api
# @BRIEF List tables/views for a database schema.
# @SIDE_EFFECT HTTP GET to Superset /api/v1/database/{id}/tables/.
# @RELATION CALLS -> [SupersetDatabasesMixin.get_database_tables]
@router.get("/databases/{database_id}/tables")
async def agent_database_tables(
database_id: int,
environment_id: str = Query(...),
schema_name: str | None = Query(None),
) -> list:
"""List tables/views for a database schema."""
client = await _get_superset_client(environment_id)
try:
return await client.get_database_tables(database_id=database_id, schema_name=schema_name)
finally:
await client.aclose()
# #endregion AgentSuperset.DatabaseTables
# #region AgentSuperset.DatabaseTableMetadata [C:2] [TYPE Endpoint] [SEMANTICS api,agent,database,metadata]
# @ingroup Api
# @BRIEF Get table metadata: columns, types, indexes, primary keys.
# @SIDE_EFFECT HTTP GET to Superset /api/v1/database/{id}/table_metadata/.
# @RELATION CALLS -> [SupersetDatabasesMixin.get_database_table_metadata]
@router.get("/databases/{database_id}/table_metadata")
async def agent_database_table_metadata(
database_id: int,
environment_id: str = Query(...),
table_name: str = Query(...),
schema_name: str | None = Query(None),
) -> dict:
"""Get table metadata: columns, types, indexes, primary keys."""
client = await _get_superset_client(environment_id)
try:
return await client.get_database_table_metadata(
database_id=database_id,
table_name=table_name,
schema_name=schema_name,
)
finally:
await client.aclose()
# #endregion AgentSuperset.DatabaseTableMetadata
# #region AgentSuperset.DatabaseSelectStar [C:2] [TYPE Endpoint] [SEMANTICS api,agent,database,select-star]
# @ingroup Api
# @BRIEF Generate a SELECT * query template for a table.
# @SIDE_EFFECT HTTP GET to Superset /api/v1/database/{id}/select_star/.
# @RELATION CALLS -> [SupersetDatabasesMixin.get_database_select_star]
@router.get("/databases/{database_id}/select_star")
async def agent_database_select_star(
database_id: int,
environment_id: str = Query(...),
table_name: str = Query(...),
schema_name: str | None = Query(None),
) -> dict:
"""Generate a SELECT * query template for a table."""
client = await _get_superset_client(environment_id)
try:
sql = await client.get_database_select_star(
database_id=database_id,
table_name=table_name,
schema_name=schema_name,
)
return {"sql": sql}
finally:
await client.aclose()
# #endregion AgentSuperset.DatabaseSelectStar
# #region AgentSuperset.DatabaseValidateSql [C:2] [TYPE Endpoint] [SEMANTICS api,agent,database,validate]
# @ingroup Api
# @BRIEF Validate SQL syntax for a database without executing.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/database/{id}/validate_sql/.
# @RELATION CALLS -> [SupersetDatabasesMixin.validate_sql]
@router.post("/databases/{database_id}/validate_sql")
async def agent_database_validate_sql(
database_id: int,
environment_id: str = Query(...),
sql: str = Query(...),
schema: str | None = Query(None),
) -> dict:
"""Validate SQL syntax for a database without executing."""
client = await _get_superset_client(environment_id)
try:
return await client.validate_sql(database_id=database_id, sql=sql, schema=schema)
finally:
await client.aclose()
# #endregion AgentSuperset.DatabaseValidateSql
# #region AgentSuperset.DatabaseTestConnection [C:3] [TYPE Endpoint] [SEMANTICS api,agent,database,test]
# @ingroup Api
# @BRIEF Test a database connection URI without creating a database entry.
# @SIDE_EFFECT HTTP POST to Superset /api/v1/database/test_connection/.
# @RELATION CALLS -> [SupersetDatabasesMixin.test_database_connection]
@router.post("/databases/test_connection")
async def agent_database_test_connection(
environment_id: str = Query(...),
database_name: str = Query(...),
sqlalchemy_uri: str = Query(...),
extra: str | None = Query(None),
) -> dict:
"""Test a database connection URI without creating a database entry."""
client = await _get_superset_client(environment_id)
try:
return await client.test_database_connection(
database_name=database_name,
sqlalchemy_uri=sqlalchemy_uri,
extra=extra,
)
finally:
await client.aclose()
# #endregion AgentSuperset.DatabaseTestConnection
# ═══════════════════════════════════════════════════════════════════
# Audit
# ═══════════════════════════════════════════════════════════════════
# #region AgentSuperset.AuditPermissions [C:3] [TYPE Endpoint] [SEMANTICS api,agent,audit,permissions]
# @ingroup Api
# @BRIEF Audit access rights: user × dashboard × dataset × RLS permission matrix.
# @SIDE_EFFECT Multiple HTTP GETs to Superset security/dashboard/dataset APIs.
# @RELATION CALLS -> [SupersetAuditMixin.permissions_audit]
@router.get("/audit/permissions")
async def agent_audit_permissions(
environment_id: str = Query(...),
page: int = Query(0),
page_size: int = Query(20),
username_filter: str | None = Query(None),
include_admin: bool = Query(False),
) -> dict:
"""Audit access rights: user × dashboards × datasets × RLS matrix."""
client = await _get_superset_client(environment_id)
try:
return await client.permissions_audit(
page=page,
page_size=page_size,
username_filter=username_filter,
include_admin=include_admin,
)
finally:
await client.aclose()
# #endregion AgentSuperset.AuditPermissions
# ═══════════════════════════════════════════════════════════════════
# Saved Queries
# ═══════════════════════════════════════════════════════════════════
# #region AgentSuperset.SavedQueryList [C:2] [TYPE Endpoint] [SEMANTICS api,agent,saved-query,list]
# @ingroup Api
# @BRIEF List all saved SQL queries.
# @SIDE_EFFECT HTTP GET to Superset /api/v1/saved_query/.
# @RELATION CALLS -> [SupersetSavedQueriesMixin.get_saved_queries]
@router.get("/saved_queries")
async def agent_saved_query_list(
environment_id: str = Query(...),
) -> dict:
"""List all saved SQL queries."""
client = await _get_superset_client(environment_id)
try:
count, queries = await client.get_saved_queries()
return {"count": count, "queries": queries}
finally:
await client.aclose()
# #endregion AgentSuperset.SavedQueryList
# #region AgentSuperset.SavedQueryGet [C:2] [TYPE Endpoint] [SEMANTICS api,agent,saved-query,get]
# @ingroup Api
# @BRIEF Get a saved SQL query by ID.
# @SIDE_EFFECT HTTP GET to Superset /api/v1/saved_query/{id}.
# @RELATION CALLS -> [SupersetSavedQueriesMixin.get_saved_query]
@router.get("/saved_queries/{query_id}")
async def agent_saved_query_get(
query_id: int,
environment_id: str = Query(...),
) -> dict:
"""Get a saved SQL query by ID."""
client = await _get_superset_client(environment_id)
try:
return await client.get_saved_query(query_id=query_id)
finally:
await client.aclose()
# #endregion AgentSuperset.SavedQueryGet
# #endregion AgentSupersetExploreRoutes

View File

@@ -33,6 +33,8 @@ from .api.routes import (
admin,
admin_api_keys,
agent_conversations,
agent_superset,
agent_superset_explore,
assistant,
clean_release,
clean_release_v2,
@@ -403,6 +405,8 @@ app.include_router(reports.router)
app.include_router(assistant.router, prefix="/api/assistant", tags=["Assistant"])
app.include_router(agent_conversations.agent_router, tags=["Agent"])
app.include_router(agent_conversations.router, tags=["Assistant"])
app.include_router(agent_superset.router, tags=["Agent Superset"])
app.include_router(agent_superset_explore.router, tags=["Agent Superset"])
app.include_router(clean_release.router)
app.include_router(clean_release_v2.router)
app.include_router(profile.router)

View File

@@ -7,6 +7,9 @@
# @RELATION DEPENDS_ON -> [SupersetAPIError]
# @RELATION DEPENDS_ON -> [get_filename_from_headers]
# @RELATION DEPENDS_ON -> [SupersetDatasetsPreviewFiltersMixin]
# @RELATION DEPENDS_ON -> [SupersetSqlLabMixin]
# @RELATION DEPENDS_ON -> [SupersetSavedQueriesMixin]
# @RELATION DEPENDS_ON -> [SupersetAuditMixin]
#
# @INVARIANT All network operations must use the internal AsyncAPIClient instance.
# @PUBLIC_API SupersetClient
@@ -33,6 +36,10 @@ from ._datasets import SupersetDatasetsMixin
from ._datasets_preview import SupersetDatasetsPreviewMixin
from ._datasets_preview_filters import SupersetDatasetsPreviewFiltersMixin
from ._user_projection import SupersetUserProjectionMixin
# NEW: Agent-critical mixins
from ._sql_lab import SupersetSqlLabMixin
from ._saved_queries import SupersetSavedQueriesMixin
from ._audit import SupersetAuditMixin
# #region SupersetClient [C:3] [TYPE Class]
@@ -52,6 +59,9 @@ from ._user_projection import SupersetUserProjectionMixin
# @RELATION INHERITS -> [SupersetDatasetsPreviewFiltersMixin]
# @RELATION INHERITS -> [SupersetDatabasesMixin]
# @RELATION INHERITS -> [SupersetDashboardsWriteMixin]
# @RELATION INHERITS -> [SupersetSqlLabMixin]
# @RELATION INHERITS -> [SupersetSavedQueriesMixin]
# @RELATION INHERITS -> [SupersetAuditMixin]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API через AsyncAPIClient.
class SupersetClient(
SupersetDatabasesMixin,
@@ -64,6 +74,10 @@ class SupersetClient(
SupersetDashboardsListMixin,
SupersetDashboardsWriteMixin,
SupersetUserProjectionMixin,
# NEW: Agent-critical mixins
SupersetSqlLabMixin,
SupersetSavedQueriesMixin,
SupersetAuditMixin,
SupersetClientBase,
):
"""Composed Superset REST API client.

View File

@@ -0,0 +1,358 @@
# #region SupersetAuditMixin [C:4] [TYPE Module] [SEMANTICS superset,audit,permissions,access-control]
# @defgroup Core Module group.
# @BRIEF Permissions audit mixin for SupersetClient — user × dashboard × dataset × RLS access matrix.
# @LAYER Infrastructure
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для сбора данных аудита.
# @RATIONALE Ported from mcp-superset audit.py. Builds a comprehensive permissions matrix
# showing effective access for every user across dashboards and datasets, resolving
# direct roles + group-inherited roles. Handles 3 dashboard access states: full (1),
# hidden (0), visible_no_data (visible but datasource_access missing).
import re
from typing import Any, cast
from ..logger import belief_scope, logger as app_logger
app_logger = cast(Any, app_logger)
# #region SupersetAuditMixin [C:4] [TYPE Class]
# @defgroup Core Module group.
# @BRIEF Mixin providing comprehensive permissions audit across users, dashboards, datasets, and RLS.
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API.
class SupersetAuditMixin:
"""Mixin for Superset permissions audit operations."""
# #region SupersetAudit.PermissionsAudit [C:4] [TYPE Function]
# @ingroup Core
# @BRIEF Build a user × dashboard × dataset × RLS permission matrix.
# @PRE SupersetClient is authenticated.
# @POST Returns audit dict with users array, each containing dashboards/datasets access and RLS regions.
# @SIDE_EFFECT Множественные асинхронные HTTP-вызовы к Superset API.
async def permissions_audit(
self,
page: int = 0,
page_size: int = 20,
username_filter: str | None = None,
include_admin: bool = False,
) -> dict:
with belief_scope("SupersetAuditMixin.permissions_audit"):
app_logger.reason("Starting permissions audit")
page_size = min(page_size, 50)
# === 1. Collect all data ===
# Users
_, all_users = await self.get_all_resources("user") if False else (
await self._fetch_users()
)
# Groups
groups_resp = await self.client.request(method="GET", endpoint="/security/groups/")
all_groups_list = groups_resp.get("result", []) if isinstance(groups_resp, dict) else []
# Group details
groups_detail: dict[int, dict] = {}
for g in all_groups_list:
gid = g.get("id")
if gid:
try:
detail = await self.client.request(method="GET", endpoint=f"/api/v1/security/groups/{gid}")
groups_detail[gid] = detail.get("result", {}) if isinstance(detail, dict) else {}
except Exception:
groups_detail[gid] = g
# Dashboards
_, all_dashboards = await self.get_dashboards()
# Datasets
_, all_datasets = await self.get_datasets()
# Dashboard → datasets and roles
dashboard_datasets: dict[int, set[int]] = {}
dashboard_roles: dict[int, set[int]] = {}
for db in all_dashboards:
db_id = db["id"]
try:
ds_resp = await self.client.request(
method="GET", endpoint=f"/dashboard/{db_id}/datasets"
)
ds_ids = set()
for ds in ds_resp.get("result", []) if isinstance(ds_resp, dict) else []:
if isinstance(ds, dict) and "id" in ds:
ds_ids.add(ds["id"])
dashboard_datasets[db_id] = ds_ids
except Exception:
dashboard_datasets[db_id] = set()
roles_data = db.get("roles", None)
if roles_data is not None:
dashboard_roles[db_id] = {r["id"] for r in roles_data}
else:
try:
db_detail = await self.client.request(method="GET", endpoint=f"/dashboard/{db_id}")
roles_detail = db_detail.get("result", {}).get("roles", []) if isinstance(db_detail, dict) else []
dashboard_roles[db_id] = {r["id"] for r in roles_detail}
except Exception:
dashboard_roles[db_id] = set()
# Role → permissions
role_perms_map = await self._build_role_permissions_map()
# Dataset → datasource_access permission_view_menu_id
ds_access_map = await self._find_datasource_permissions()
# RLS rules
rls_resp = await self.client.request(method="GET", endpoint="/rowlevelsecurity/")
all_rls: list[dict] = []
if isinstance(rls_resp, dict):
all_rls = rls_resp.get("result", [])
# === 2. Build helpers ===
# User groups
user_groups: dict[int, list[dict]] = {}
user_group_roles: dict[int, set[int]] = {}
for gid, gdata in groups_detail.items():
g_name = gdata.get("name", "?")
g_roles = {r["id"] for r in gdata.get("roles", [])}
for u in gdata.get("users", []):
uid = u["id"]
if uid not in user_groups:
user_groups[uid] = []
user_group_roles[uid] = set()
user_groups[uid].append({"name": g_name, "roles": sorted(g_roles)})
user_group_roles[uid] |= g_roles
# RLS regions per role
role_rls_regions: dict[int, list[str]] = {}
for rule in all_rls:
clause = rule.get("clause", "")
roles = rule.get("roles", [])
region_match = re.search(r"operation_region\s*=\s*'([^']+)'", clause)
if clause == "1=1":
region = "_all_"
elif region_match:
region = region_match.group(1)
elif clause == "1=0":
region = "_deny_"
else:
continue
for r in roles:
rid = r["id"]
if rid not in role_rls_regions:
role_rls_regions[rid] = []
role_rls_regions[rid].append(region)
# === 3. Filter users ===
filtered_users = all_users
if not include_admin:
filtered_users = [
u for u in filtered_users
if not any(r.get("name") == "Admin" for r in u.get("roles", []))
]
if username_filter:
filtered_users = [
u for u in filtered_users
if username_filter.lower() in u.get("username", "").lower()
]
total = len(filtered_users)
start = page * page_size
end = start + page_size
page_users = filtered_users[start:end]
# Data structures
dashboard_info = {
d["id"]: d.get("dashboard_title", d.get("slug", f"id:{d['id']}"))
for d in all_dashboards
}
dataset_info = {
d["id"]: d.get("table_name", f"id:{d['id']}") for d in all_datasets
}
# === 4. Build audit matrix ===
audit_rows = []
for user in page_users:
uid = user["id"]
uname = user.get("username", "?")
direct_role_ids = {r["id"] for r in user.get("roles", [])}
group_role_ids = user_group_roles.get(uid, set())
effective_role_ids = direct_role_ids | group_role_ids
user_perm_ids: set[int] = set()
for rid in effective_role_ids:
user_perm_ids |= role_perms_map.get(rid, set())
# Dataset access
datasets_access: dict[str, int] = {}
for ds in all_datasets:
ds_id = ds["id"]
ds_name = dataset_info[ds_id]
pvm_id = ds_access_map.get(ds_id)
datasets_access[ds_name] = 1 if (pvm_id and pvm_id in user_perm_ids) else 0
# Dashboard access
dashboards_access: dict[str, Any] = {}
for db in all_dashboards:
db_id = db["id"]
db_name = dashboard_info[db_id]
db_role_ids = dashboard_roles.get(db_id, set())
is_visible = bool(effective_role_ids & db_role_ids) if db_role_ids else True
required_ds = dashboard_datasets.get(db_id, set())
if not required_ds:
has_data = True
else:
has_data = all(
ds_access_map.get(ds_id) in user_perm_ids
for ds_id in required_ds
if ds_access_map.get(ds_id) is not None
)
if is_visible and has_data:
dashboards_access[db_name] = 1
elif is_visible and not has_data:
dashboards_access[db_name] = "visible_no_data"
else:
dashboards_access[db_name] = 0
# RLS regions
rls_regions: list[str] = []
has_deny = False
for rid in effective_role_ids:
regions = role_rls_regions.get(rid, [])
for region in regions:
if region == "_all_":
rls_regions = ["ALL REGIONS"]
break
elif region == "_deny_":
has_deny = True
elif region not in rls_regions:
rls_regions.append(region)
if rls_regions == ["ALL REGIONS"]:
break
if not rls_regions and has_deny:
rls_regions = ["DENIED (1=0)"]
elif not rls_regions:
rls_regions = ["no RLS"]
groups_names = [g["name"] for g in user_groups.get(uid, [])]
audit_rows.append({
"user_id": uid,
"username": uname,
"active": user.get("active", True),
"groups": groups_names,
"dashboards": dashboards_access,
"datasets": datasets_access,
"rls_regions": sorted(rls_regions),
})
result = {
"page": page,
"page_size": page_size,
"total_users": total,
"total_pages": (total + page_size - 1) // page_size if page_size > 0 else 0,
"dashboards_checked": list(dashboard_info.values()),
"datasets_checked": list(dataset_info.values()),
"users": audit_rows,
}
app_logger.reflect(
"Permissions audit complete",
extra={"total_users": total, "audited": len(audit_rows)},
)
return result
# #endregion SupersetAudit.PermissionsAudit
# #region SupersetAudit.FetchUsers [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Fetch all Superset users via security API.
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API.
async def _fetch_users(self) -> tuple[int, list[dict]]:
"""Fetch all users via /api/v1/security/users/ with pagination."""
with belief_scope("SupersetAuditMixin._fetch_users"):
app_logger.reason("Fetching all users")
validated_query = self._validate_query_params({})
paginated_data = await self._fetch_all_pages(
endpoint="/api/v1/security/users/",
pagination_options={
"base_query": validated_query,
"results_field": "result",
},
)
return len(paginated_data), paginated_data
# #endregion SupersetAudit.FetchUsers
# #region SupersetAudit.BuildRolePermissionsMap [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Build a mapping of role_id → set(permission_view_menu_id) for all roles.
async def _build_role_permissions_map(self) -> dict[int, set[int]]:
"""Build {role_id: set(permission_view_menu_ids)} for all roles."""
role_perms: dict[int, set[int]] = {}
try:
roles_resp = await self.client.request(method="GET", endpoint="/api/v1/security/roles/")
all_roles = roles_resp.get("result", []) if isinstance(roles_resp, dict) else []
except Exception:
roles_resp = await self.client.request(method="GET", endpoint="/security/roles/")
all_roles = roles_resp.get("result", []) if isinstance(roles_resp, dict) else []
for role in all_roles:
role_id = role["id"]
try:
perms_resp = await self.client.request(
method="GET", endpoint=f"/api/v1/security/roles/{role_id}/permissions/"
)
perm_ids: set[int] = set()
for p in perms_resp.get("result", []) if isinstance(perms_resp, dict) else []:
if isinstance(p, dict) and "id" in p:
perm_ids.add(p["id"])
elif isinstance(p, int):
perm_ids.add(p)
role_perms[role_id] = perm_ids
except Exception:
role_perms[role_id] = set()
return role_perms
# #endregion SupersetAudit.BuildRolePermissionsMap
# #region SupersetAudit.FindDatasourcePermissions [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Build {dataset_id: permission_view_menu_id} mapping by scanning permissions-resources.
async def _find_datasource_permissions(self) -> dict[int, int]:
"""Scan /api/v1/security/permissions-resources/ for datasource_access entries."""
found: dict[int, int] = {}
dataset_id_re = re.compile(r"\(id:(\d+)\)")
page = 0
while page < 50:
try:
resp = await self.client.request(
method="GET",
endpoint="/api/v1/security/permissions-resources/",
params={"q": f"(page:{page},page_size:100)"},
)
except Exception:
break
items = resp.get("result", []) if isinstance(resp, dict) else []
if not items:
break
for item in items:
if item.get("permission", {}).get("name", "") != "datasource_access":
continue
view_name = item.get("view_menu", {}).get("name", "")
match = dataset_id_re.search(view_name)
if match:
found[int(match.group(1))] = item["id"]
if len(items) < 100:
break
page += 1
return found
# #endregion SupersetAudit.FindDatasourcePermissions
# #endregion SupersetAuditMixin
# #endregion SupersetAuditMixin

View File

@@ -353,5 +353,81 @@ class SupersetDashboardsCrudMixin:
payload={"dashboard_id": dashboard_id, "response": response},
)
# #endregion SupersetClientDeleteDashboard
# #region SupersetClientGetDashboardCharts [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF List all charts on a dashboard (standalone public method).
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/dashboard/{id}/charts.
async def get_dashboard_charts(self, dashboard_id: int) -> list[dict]:
with belief_scope("SupersetClient.get_dashboard_charts", f"id={dashboard_id}"):
app_logger.reason("Fetching dashboard charts", extra={"dashboard_id": dashboard_id})
try:
charts_response = await self.client.request(
method="GET", endpoint=f"/dashboard/{dashboard_id}/charts"
)
charts_payload = (
charts_response.get("result", [])
if isinstance(charts_response, dict)
else []
)
app_logger.reflect(
f"Found {len(charts_payload)} charts on dashboard {dashboard_id}"
)
return charts_payload
except Exception as e:
app_logger.explore(
"Failed to fetch dashboard charts",
extra={"dashboard_id": dashboard_id},
error=str(e),
)
raise
# #endregion SupersetClientGetDashboardCharts
# #region SupersetClientGetDashboardDatasets [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF List all datasets used by a dashboard (standalone public method).
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/dashboard/{id}/datasets.
async def get_dashboard_datasets(self, dashboard_id: int) -> list[dict]:
with belief_scope("SupersetClient.get_dashboard_datasets", f"id={dashboard_id}"):
app_logger.reason("Fetching dashboard datasets", extra={"dashboard_id": dashboard_id})
try:
datasets_response = await self.client.request(
method="GET", endpoint=f"/dashboard/{dashboard_id}/datasets"
)
datasets_payload = (
datasets_response.get("result", [])
if isinstance(datasets_response, dict)
else []
)
app_logger.reflect(
f"Found {len(datasets_payload)} datasets on dashboard {dashboard_id}"
)
return datasets_payload
except Exception as e:
app_logger.explore(
"Failed to fetch dashboard datasets",
extra={"dashboard_id": dashboard_id},
error=str(e),
)
raise
# #endregion SupersetClientGetDashboardDatasets
# #region SupersetClientGetDashboard [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Fetch raw dashboard response by ID or slug.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/dashboard/{ref}.
async def get_dashboard(self, dashboard_ref: int | str) -> dict:
with belief_scope("SupersetClient.get_dashboard", f"ref={dashboard_ref}"):
app_logger.reason("Fetching dashboard", extra={"dashboard_ref": str(dashboard_ref)})
try:
response = await self.client.request(
method="GET", endpoint=f"/dashboard/{dashboard_ref}"
)
app_logger.reflect(f"Got dashboard {dashboard_ref}")
return cast(dict, response)
except Exception as e:
app_logger.explore("Fetch dashboard failed", error=str(e))
raise
# #endregion SupersetClientGetDashboard
# #endregion SupersetDashboardsCrudMixin
# #endregion SupersetDashboardsCrudMixin

View File

@@ -298,5 +298,187 @@ class SupersetDashboardsWriteMixin:
raise RuntimeError("Cannot create markdown chart: no datasets available in Superset.")
# #endregion _resolve_markdown_datasource
# #region create_dashboard [C:4] [TYPE Function]
# @ingroup Core
# @BRIEF Create a new dashboard in Superset with optional roles, metadata, and positioning.
# @PRE Environment configured. dashboard_title is non-empty.
# @POST Returns created dashboard dict with id.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/dashboard/.
async def create_dashboard(
self,
dashboard_title: str,
slug: str | None = None,
published: bool = False,
json_metadata: str | None = None,
css: str | None = None,
position_json: str | None = None,
roles: list[int] | None = None,
) -> dict:
with belief_scope("SupersetDashboardsWriteMixin.create_dashboard"):
app_logger.reason(
"Creating dashboard",
extra={"title": dashboard_title, "published": published},
)
payload: dict[str, Any] = {"dashboard_title": dashboard_title, "published": published}
if slug:
payload["slug"] = slug
if json_metadata:
payload["json_metadata"] = json_metadata
if css:
payload["css"] = css
if position_json:
payload["position_json"] = position_json
if roles:
payload["roles"] = roles
try:
response = cast(dict, await self.client.request(
method="POST",
endpoint="/dashboard/",
data=payload,
))
app_logger.reflect(
"Dashboard created",
extra={"id": response.get("id"), "title": dashboard_title},
)
return response
except Exception as e:
app_logger.explore(
"Create dashboard failed",
extra={"title": dashboard_title},
error=str(e),
)
raise
# #endregion create_dashboard
# #region update_dashboard [C:4] [TYPE Function]
# @ingroup Core
# @BRIEF Update an existing dashboard's properties (title, slug, owners, roles, metadata, CSS).
# @PRE dashboard_id exists in Superset.
# @POST Returns updated dashboard dict.
# @SIDE_EFFECT Асинхронный HTTP-вызов: PUT /api/v1/dashboard/{id}.
async def update_dashboard(
self,
dashboard_id: int,
dashboard_title: str | None = None,
slug: str | None = None,
published: bool | None = None,
owners: list[int] | None = None,
roles: list[int] | None = None,
css: str | None = None,
json_metadata: str | None = None,
) -> dict:
with belief_scope("SupersetDashboardsWriteMixin.update_dashboard", f"id={dashboard_id}"):
app_logger.reason("Updating dashboard", extra={"dashboard_id": dashboard_id})
payload: dict[str, Any] = {}
if dashboard_title is not None:
payload["dashboard_title"] = dashboard_title
if slug is not None:
payload["slug"] = slug
if published is not None:
payload["published"] = published
if owners is not None:
payload["owners"] = owners
if roles is not None:
payload["roles"] = roles
if css is not None:
payload["css"] = css
if json_metadata is not None:
payload["json_metadata"] = json_metadata
try:
response = cast(dict, await self.client.request(
method="PUT",
endpoint=f"/dashboard/{dashboard_id}",
data=payload,
))
app_logger.reflect("Dashboard updated", extra={"dashboard_id": dashboard_id})
return response
except Exception as e:
app_logger.explore(
"Update dashboard failed",
extra={"dashboard_id": dashboard_id},
error=str(e),
)
raise
# #endregion update_dashboard
# #region copy_dashboard [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Deep-copy a dashboard (including all charts) to a new dashboard.
# @PRE dashboard_id exists.
# @POST Returns created copy dict.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/dashboard/{id}/copy/.
async def copy_dashboard(
self,
dashboard_id: int,
dashboard_title: str | None = None,
json_metadata: str | None = None,
) -> dict:
with belief_scope("SupersetDashboardsWriteMixin.copy_dashboard", f"id={dashboard_id}"):
app_logger.reason("Copying dashboard", extra={"dashboard_id": dashboard_id})
payload: dict[str, Any] = {}
if dashboard_title:
payload["dashboard_title"] = dashboard_title
if json_metadata:
payload["json_metadata"] = json_metadata
try:
response = cast(dict, await self.client.request(
method="POST",
endpoint=f"/dashboard/{dashboard_id}/copy/",
data=payload if payload else None,
))
app_logger.reflect(
"Dashboard copied",
extra={"source_id": dashboard_id, "copy_id": response.get("id")},
)
return response
except Exception as e:
app_logger.explore(
"Copy dashboard failed",
extra={"dashboard_id": dashboard_id},
error=str(e),
)
raise
# #endregion copy_dashboard
# #region publish_dashboard [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Set dashboard published=true.
# @SIDE_EFFECT Асинхронный HTTP-вызов: PUT /api/v1/dashboard/{id}.
async def publish_dashboard(self, dashboard_id: int) -> dict:
with belief_scope("SupersetDashboardsWriteMixin.publish_dashboard", f"id={dashboard_id}"):
app_logger.reason("Publishing dashboard", extra={"dashboard_id": dashboard_id})
try:
response = cast(dict, await self.client.request(
method="PUT",
endpoint=f"/dashboard/{dashboard_id}",
data={"published": True},
))
app_logger.reflect("Dashboard published", extra={"dashboard_id": dashboard_id})
return response
except Exception as e:
app_logger.explore("Publish dashboard failed", error=str(e))
raise
# #endregion publish_dashboard
# #region unpublish_dashboard [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Set dashboard published=false.
# @SIDE_EFFECT Асинхронный HTTP-вызов: PUT /api/v1/dashboard/{id}.
async def unpublish_dashboard(self, dashboard_id: int) -> dict:
with belief_scope("SupersetDashboardsWriteMixin.unpublish_dashboard", f"id={dashboard_id}"):
app_logger.reason("Unpublishing dashboard", extra={"dashboard_id": dashboard_id})
try:
response = cast(dict, await self.client.request(
method="PUT",
endpoint=f"/dashboard/{dashboard_id}",
data={"published": False},
))
app_logger.reflect("Dashboard unpublished", extra={"dashboard_id": dashboard_id})
return response
except Exception as e:
app_logger.explore("Unpublish dashboard failed", error=str(e))
raise
# #endregion unpublish_dashboard
# #endregion SupersetDashboardsWriteMixin
# #endregion SupersetDashboardsWriteMixin

View File

@@ -4,6 +4,7 @@
# @BRIEF Database domain mixin for SupersetClient — list, get, summary, by_uuid.
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для работы с базами данных.
import json
from typing import Any, cast
from ..logger import belief_scope, logger as app_logger
@@ -128,5 +129,333 @@ class SupersetDatabasesMixin:
app_logger.reflect("Database deleted", extra={"database_id": database_id})
return response
# #endregion SupersetClientDeleteDatabase
# #region SupersetClientUpdateDatabase [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Update an existing database connection's properties.
# @SIDE_EFFECT Асинхронный HTTP-вызов: PUT /api/v1/database/{id}.
async def update_database(
self,
database_id: int,
database_name: str | None = None,
sqlalchemy_uri: str | None = None,
expose_in_sqllab: bool | None = None,
allow_ctas: bool | None = None,
allow_cvas: bool | None = None,
allow_dml: bool | None = None,
allow_run_async: bool | None = None,
extra: str | None = None,
) -> dict:
with belief_scope("SupersetClient.update_database", f"id={database_id}"):
app_logger.reason("Updating database", extra={"database_id": database_id})
payload: dict[str, Any] = {}
if database_name is not None:
payload["database_name"] = database_name
if sqlalchemy_uri is not None:
payload["sqlalchemy_uri"] = sqlalchemy_uri
if expose_in_sqllab is not None:
payload["expose_in_sqllab"] = expose_in_sqllab
if allow_ctas is not None:
payload["allow_ctas"] = allow_ctas
if allow_cvas is not None:
payload["allow_cvas"] = allow_cvas
if allow_dml is not None:
payload["allow_dml"] = allow_dml
if allow_run_async is not None:
payload["allow_run_async"] = allow_run_async
if extra is not None:
payload["extra"] = extra
try:
response = await self.client.request(
method="PUT",
endpoint=f"/database/{database_id}",
data=payload,
)
response = cast(dict, response)
app_logger.reflect(f"Database {database_id} updated")
return response
except Exception as e:
app_logger.explore("Update database failed", error=str(e))
raise
# #endregion SupersetClientUpdateDatabase
# #region SupersetClientTestDatabaseConnection [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Test a database connection URI without creating a database entry.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/database/test_connection/.
async def test_database_connection(
self,
database_name: str,
sqlalchemy_uri: str,
extra: str | None = None,
) -> dict:
with belief_scope("SupersetClient.test_database_connection"):
app_logger.reason("Testing database connection", extra={"name": database_name})
payload: dict[str, Any] = {
"database_name": database_name,
"sqlalchemy_uri": sqlalchemy_uri,
}
if extra:
payload["extra"] = extra
try:
response = await self.client.request(
method="POST",
endpoint="/database/test_connection/",
data=payload,
)
response = cast(dict, response)
app_logger.reflect("Database connection tested")
return response
except Exception as e:
app_logger.explore("Test connection failed", error=str(e))
raise
# #endregion SupersetClientTestDatabaseConnection
# #region SupersetClientGetDatabaseSchemas [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF List all schemas for a database.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/{id}/schemas/.
async def get_database_schemas(self, database_id: int) -> list[str]:
with belief_scope("SupersetClient.get_database_schemas", f"id={database_id}"):
app_logger.reason("Fetching database schemas", extra={"database_id": database_id})
try:
response = await self.client.request(
method="GET",
endpoint=f"/database/{database_id}/schemas/",
)
response = cast(dict, response)
schemas = response.get("result", [])
app_logger.reflect(f"Found {len(schemas)} schemas")
return schemas
except Exception as e:
app_logger.explore("Fetch schemas failed", error=str(e))
raise
# #endregion SupersetClientGetDatabaseSchemas
# #region SupersetClientGetDatabaseTables [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF List all tables/views for a database schema.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/{id}/tables/.
async def get_database_tables(
self,
database_id: int,
schema_name: str | None = None,
) -> list[dict]:
with belief_scope("SupersetClient.get_database_tables", f"id={database_id}"):
app_logger.reason("Fetching database tables", extra={"database_id": database_id})
params = {"q": json.dumps({"schema_name": schema_name})} if schema_name else None
try:
response = await self.client.request(
method="GET",
endpoint=f"/database/{database_id}/tables/",
params=params,
)
response = cast(dict, response)
tables = response.get("result", [])
app_logger.reflect(f"Found {len(tables)} tables")
return tables
except Exception as e:
app_logger.explore("Fetch tables failed", error=str(e))
raise
# #endregion SupersetClientGetDatabaseTables
# #region SupersetClientGetDatabaseCatalogs [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF List catalogs for a database (if supported by the DB engine).
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/{id}/catalogs/.
async def get_database_catalogs(self, database_id: int) -> list[str]:
with belief_scope("SupersetClient.get_database_catalogs", f"id={database_id}"):
app_logger.reason("Fetching database catalogs", extra={"database_id": database_id})
try:
response = await self.client.request(
method="GET",
endpoint=f"/database/{database_id}/catalogs/",
)
response = cast(dict, response)
catalogs = response.get("result", [])
app_logger.reflect(f"Found {len(catalogs)} catalogs")
return catalogs
except Exception as e:
app_logger.explore("Fetch catalogs failed", error=str(e))
raise
# #endregion SupersetClientGetDatabaseCatalogs
# #region SupersetClientValidateSql [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Validate SQL syntax for a given database without executing.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/database/{id}/validate_sql/.
async def validate_sql(
self,
database_id: int,
sql: str,
schema: str | None = None,
) -> dict:
with belief_scope("SupersetClient.validate_sql", f"id={database_id}"):
app_logger.reason("Validating SQL", extra={"database_id": database_id})
payload: dict[str, Any] = {"sql": sql}
if schema:
payload["schema"] = schema
try:
response = await self.client.request(
method="POST",
endpoint=f"/database/{database_id}/validate_sql/",
data=payload,
)
response = cast(dict, response)
app_logger.reflect("SQL validated")
return response
except Exception as e:
app_logger.explore("Validate SQL failed", error=str(e))
raise
# #endregion SupersetClientValidateSql
# #region SupersetClientValidateParameters [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Validate database connection parameters before creation.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/database/validate_parameters/.
async def validate_database_parameters(
self,
engine: str,
parameters: dict[str, Any],
configuration_method: str = "dynamic_form",
) -> dict:
with belief_scope("SupersetClient.validate_database_parameters"):
app_logger.reason("Validating database parameters", extra={"engine": engine})
try:
response = await self.client.request(
method="POST",
endpoint="/database/validate_parameters/",
data={
"engine": engine,
"parameters": parameters,
"configuration_method": configuration_method,
},
)
response = cast(dict, response)
app_logger.reflect("Parameters validated")
return response
except Exception as e:
app_logger.explore("Validate parameters failed", error=str(e))
raise
# #endregion SupersetClientValidateParameters
# #region SupersetClientGetSelectStar [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Generate a SELECT * query for a given table (template for SQL Lab).
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/{id}/select_star/.
async def get_database_select_star(
self,
database_id: int,
table_name: str,
schema_name: str | None = None,
) -> str:
with belief_scope("SupersetClient.get_database_select_star", f"id={database_id}"):
app_logger.reason("Generating SELECT *", extra={"table": table_name})
params = {"q": json.dumps({"table_name": table_name, "schema_name": schema_name})}
try:
response = await self.client.request(
method="GET",
endpoint=f"/database/{database_id}/select_star/",
params=params,
)
response = cast(dict, response)
sql = response.get("result", "")
app_logger.reflect("SELECT * generated")
return sql
except Exception as e:
app_logger.explore("Generate SELECT * failed", error=str(e))
raise
# #endregion SupersetClientGetSelectStar
# #region SupersetClientGetTableMetadata [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Get table metadata: columns, types, indexes, primary keys.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/{id}/table_metadata/.
async def get_database_table_metadata(
self,
database_id: int,
table_name: str,
schema_name: str | None = None,
) -> dict:
with belief_scope("SupersetClient.get_database_table_metadata", f"id={database_id}"):
app_logger.reason("Fetching table metadata", extra={"table": table_name})
params = {"q": json.dumps({"table_name": table_name, "schema_name": schema_name})}
try:
response = await self.client.request(
method="GET",
endpoint=f"/database/{database_id}/table_metadata/",
params=params,
)
response = cast(dict, response)
app_logger.reflect(f"Table metadata fetched for {table_name}")
return response
except Exception as e:
app_logger.explore("Fetch table metadata failed", error=str(e))
raise
# #endregion SupersetClientGetTableMetadata
# #region SupersetClientGetDatabaseRelatedObjects [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Get related objects (impact analysis) for a database before deletion.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/{id}/related_objects/.
async def get_database_related_objects(self, database_id: int) -> dict:
with belief_scope("SupersetClient.get_database_related_objects", f"id={database_id}"):
app_logger.reason("Fetching database related objects", extra={"database_id": database_id})
try:
response = await self.client.request(
method="GET",
endpoint=f"/database/{database_id}/related_objects/",
)
response = cast(dict, response)
app_logger.reflect(f"Related objects fetched for database {database_id}")
return response
except Exception as e:
app_logger.explore("Fetch related objects failed", error=str(e))
raise
# #endregion SupersetClientGetDatabaseRelatedObjects
# #region SupersetClientExportDatabase [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Export database configurations as ZIP.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/export/.
async def export_database(self, database_ids: list[int]) -> tuple[bytes, str]:
with belief_scope("SupersetClient.export_database"):
app_logger.reason("Exporting databases", extra={"database_ids": database_ids})
try:
response = await self.client.request(
method="GET",
endpoint="/database/export/",
params={"q": json.dumps(database_ids)},
raw_response=True,
)
response = cast(Any, response)
content = response.content if hasattr(response, "content") else response
app_logger.reflect(f"Exported {len(database_ids)} databases")
return content, f"databases_export_{len(database_ids)}.zip"
except Exception as e:
app_logger.explore("Export databases failed", error=str(e))
raise
# #endregion SupersetClientExportDatabase
# #region SupersetClientGetAvailableEngines [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF List all supported database engine types.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/database/available/.
async def get_available_engines(self) -> list[dict]:
with belief_scope("SupersetClient.get_available_engines"):
app_logger.reason("Fetching available database engines")
try:
response = await self.client.request(
method="GET",
endpoint="/database/available/",
)
response = cast(dict, response)
engines = response.get("result", [])
app_logger.reflect(f"Found {len(engines)} available engines")
return engines
except Exception as e:
app_logger.explore("Fetch available engines failed", error=str(e))
raise
# #endregion SupersetClientGetAvailableEngines
# #endregion SupersetDatabasesMixin
# #endregion SupersetDatabasesMixin

View File

@@ -5,6 +5,7 @@
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для работы с датасетами.
import json
from pathlib import Path
from typing import Any, cast
from ..logger import belief_scope, logger as app_logger
@@ -254,5 +255,214 @@ class SupersetDatasetsMixin:
app_logger.reflect(f"Updated dataset {dataset_id}.", extra={"src": "update_dataset"})
return response
# #endregion SupersetClientUpdateDataset
# #region SupersetClientCreateDataset [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Create a new physical or virtual dataset in Superset.
# @POST Returns created dataset dict with id.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/dataset/.
async def create_dataset(
self,
table_name: str,
database: int,
schema_name: str | None = None,
sql: str | None = None,
) -> dict:
with belief_scope("SupersetClient.create_dataset"):
app_logger.reason(
"Creating dataset",
extra={"table_name": table_name, "database": database},
)
payload: dict[str, Any] = {
"table_name": table_name,
"database": database,
}
if schema_name:
payload["schema"] = schema_name
if sql:
payload["sql"] = sql
try:
response = await self.client.request(
method="POST",
endpoint="/dataset/",
data=payload,
)
response = cast(dict, response)
app_logger.reflect(
"Dataset created",
extra={"id": response.get("id"), "table_name": table_name},
)
return response
except Exception as e:
app_logger.explore("Create dataset failed", error=str(e))
raise
# #endregion SupersetClientCreateDataset
# #region SupersetClientDeleteDataset [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Delete a dataset from Superset by ID.
# @SIDE_EFFECT Асинхронный HTTP-вызов: DELETE /api/v1/dataset/{id}.
async def delete_dataset(self, dataset_id: int) -> dict:
with belief_scope("SupersetClient.delete_dataset", f"id={dataset_id}"):
app_logger.reason("Deleting dataset", extra={"dataset_id": dataset_id})
try:
response = await self.client.request(
method="DELETE",
endpoint=f"/dataset/{dataset_id}",
)
response = cast(dict, response)
app_logger.reflect(f"Dataset {dataset_id} deleted")
return response
except Exception as e:
app_logger.explore("Delete dataset failed", error=str(e))
raise
# #endregion SupersetClientDeleteDataset
# #region SupersetClientDuplicateDataset [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Duplicate a dataset including its columns and metrics.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/dataset/duplicate.
async def duplicate_dataset(self, base_model_id: int, table_name: str) -> dict:
with belief_scope("SupersetClient.duplicate_dataset", f"id={base_model_id}"):
app_logger.reason("Duplicating dataset", extra={"base_model_id": base_model_id})
try:
response = await self.client.request(
method="POST",
endpoint="/dataset/duplicate",
data={"base_model_id": base_model_id, "table_name": table_name},
)
response = cast(dict, response)
app_logger.reflect(
"Dataset duplicated",
extra={"source_id": base_model_id, "new_id": response.get("id")},
)
return response
except Exception as e:
app_logger.explore("Duplicate dataset failed", error=str(e))
raise
# #endregion SupersetClientDuplicateDataset
# #region SupersetClientRefreshDatasetSchema [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Rescan columns and types for a dataset from its source database.
# @SIDE_EFFECT Асинхронный HTTP-вызов: PUT /api/v1/dataset/{id}/refresh.
async def refresh_dataset_schema(self, dataset_id: int) -> dict:
with belief_scope("SupersetClient.refresh_dataset_schema", f"id={dataset_id}"):
app_logger.reason("Refreshing dataset schema", extra={"dataset_id": dataset_id})
try:
response = await self.client.request(
method="PUT",
endpoint=f"/dataset/{dataset_id}/refresh",
)
response = cast(dict, response)
app_logger.reflect(f"Dataset {dataset_id} schema refreshed")
return response
except Exception as e:
app_logger.explore("Refresh schema failed", error=str(e))
raise
# #endregion SupersetClientRefreshDatasetSchema
# #region SupersetClientGetOrCreateDataset [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Idempotently get or create a dataset by database_id, table_name, and optional schema.
# @SIDE_EFFECT Асинхронные HTTP-вызовы: GET /dataset/ then POST /dataset/get_or_create/.
async def get_or_create_dataset(
self,
database_id: int,
table_name: str,
schema_name: str | None = None,
) -> dict:
with belief_scope("SupersetClient.get_or_create_dataset"):
app_logger.reason(
"Get-or-create dataset",
extra={"database_id": database_id, "table_name": table_name},
)
try:
response = await self.client.request(
method="POST",
endpoint="/dataset/get_or_create/",
data={
"database_id": database_id,
"table_name": table_name,
"schema": schema_name,
},
)
response = cast(dict, response)
app_logger.reflect(
"Dataset get-or-create done",
extra={"id": response.get("id"), "table_name": table_name},
)
return response
except Exception as e:
app_logger.explore("Get-or-create dataset failed", error=str(e))
raise
# #endregion SupersetClientGetOrCreateDataset
# #region SupersetClientExportDataset [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Export datasets as a ZIP archive by their IDs.
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/dataset/export/.
async def export_dataset(self, dataset_ids: list[int]) -> tuple[bytes, str]:
with belief_scope("SupersetClient.export_dataset"):
app_logger.reason("Exporting datasets", extra={"dataset_ids": dataset_ids})
try:
response = await self.client.request(
method="GET",
endpoint="/dataset/export/",
params={"q": json.dumps(dataset_ids)},
raw_response=True,
)
response = cast(Any, response)
content = response.content if hasattr(response, "content") else response
app_logger.reflect(f"Exported {len(dataset_ids)} datasets")
return content, f"datasets_export_{len(dataset_ids)}.zip"
except Exception as e:
app_logger.explore("Export datasets failed", error=str(e))
raise
# #endregion SupersetClientExportDataset
# #region SupersetClientImportDataset [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Import datasets from a ZIP file.
# @SIDE_EFFECT Асинхронный HTTP-вызов: POST /api/v1/dataset/import/.
async def import_dataset(self, file_path: str, overwrite: bool = True) -> dict:
with belief_scope("SupersetClient.import_dataset"):
app_logger.reason("Importing datasets", extra={"file": file_path})
try:
response = await self.client.upload_file(
endpoint="/dataset/import/",
file_info={
"file_obj": file_path,
"file_name": Path(file_path).name,
"form_field": "formData",
},
extra_data={"overwrite": "true" if overwrite else "false"},
)
app_logger.reflect("Datasets imported")
return cast(dict, response)
except Exception as e:
app_logger.explore("Import datasets failed", error=str(e))
raise
# #endregion SupersetClientImportDataset
# #region SupersetClientGetDatasetRelatedObjects [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Get related charts and dashboards for a dataset (standalone public method).
# @SIDE_EFFECT Асинхронный HTTP-вызов: GET /api/v1/dataset/{id}/related_objects/.
async def get_dataset_related_objects(self, dataset_id: int) -> dict:
with belief_scope("SupersetClient.get_dataset_related_objects", f"id={dataset_id}"):
app_logger.reason("Fetching dataset related objects", extra={"dataset_id": dataset_id})
try:
response = await self.client.request(
method="GET",
endpoint=f"/dataset/{dataset_id}/related_objects",
)
response = cast(dict, response)
app_logger.reflect(f"Got related objects for dataset {dataset_id}")
return response
except Exception as e:
app_logger.explore("Fetch related objects failed", error=str(e))
raise
# #endregion SupersetClientGetDatasetRelatedObjects
# #endregion SupersetDatasetsMixin
# #endregion SupersetDatasetsMixin

View File

@@ -0,0 +1,160 @@
# #region SupersetSavedQueriesMixin [C:3] [TYPE Module] [SEMANTICS superset,saved-query,crud]
# @defgroup Core Module group.
# @BRIEF Saved SQL queries mixin for SupersetClient — list, create, get, update, delete.
# @LAYER Infrastructure
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для CRUD-операций с сохранёнными запросами.
import json
from typing import Any, cast
from ..logger import belief_scope, logger as app_logger
app_logger = cast(Any, app_logger)
# #region SupersetSavedQueriesMixin [C:3] [TYPE Class]
# @defgroup Core Module group.
# @BRIEF Mixin providing saved query CRUD operations.
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API.
class SupersetSavedQueriesMixin:
"""Mixin for Superset saved SQL queries CRUD."""
# #region SupersetSavedQueries.List [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF List all saved SQL queries with pagination.
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для пагинации.
async def get_saved_queries(self, query: dict | None = None) -> tuple[int, list[dict]]:
with belief_scope("SupersetSavedQueriesMixin.get_saved_queries"):
app_logger.reason("Fetching saved queries")
validated_query = self._validate_query_params(query or {})
paginated_data = await self._fetch_all_pages(
endpoint="/saved_query/",
pagination_options={
"base_query": validated_query,
"results_field": "result",
},
)
total_count = len(paginated_data)
app_logger.reflect(f"Found {total_count} saved queries.")
return total_count, paginated_data
# #endregion SupersetSavedQueries.List
# #region SupersetSavedQueries.Get [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Get a single saved query by ID.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: GET /api/v1/saved_query/{id}.
async def get_saved_query(self, query_id: int) -> dict:
with belief_scope("SupersetSavedQueriesMixin.get_saved_query", f"id={query_id}"):
app_logger.reason("Fetching saved query", extra={"query_id": query_id})
try:
response = await self.client.request(
method="GET",
endpoint=f"/saved_query/{query_id}",
)
response = cast(dict, response)
app_logger.reflect(f"Got saved query {query_id}")
return response
except Exception as e:
app_logger.explore("Fetch saved query failed", error=str(e))
raise
# #endregion SupersetSavedQueries.Get
# #region SupersetSavedQueries.Create [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Create a new saved SQL query.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: POST /api/v1/saved_query/.
async def create_saved_query(
self,
label: str,
db_id: int,
sql: str,
schema: str | None = None,
description: str | None = None,
) -> dict:
with belief_scope("SupersetSavedQueriesMixin.create_saved_query"):
app_logger.reason("Creating saved query", extra={"label": label, "db_id": db_id})
payload: dict[str, Any] = {
"label": label,
"db_id": db_id,
"sql": sql,
}
if schema:
payload["schema"] = schema
if description:
payload["description"] = description
try:
response = await self.client.request(
method="POST",
endpoint="/saved_query/",
data=payload,
)
response = cast(dict, response)
app_logger.reflect(
"Saved query created",
extra={"id": response.get("id"), "label": label},
)
return response
except Exception as e:
app_logger.explore("Create saved query failed", error=str(e))
raise
# #endregion SupersetSavedQueries.Create
# #region SupersetSavedQueries.Update [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Update an existing saved SQL query.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: PUT /api/v1/saved_query/{id}.
async def update_saved_query(
self,
query_id: int,
label: str | None = None,
sql: str | None = None,
schema: str | None = None,
description: str | None = None,
) -> dict:
with belief_scope("SupersetSavedQueriesMixin.update_saved_query", f"id={query_id}"):
app_logger.reason("Updating saved query", extra={"query_id": query_id})
payload: dict[str, Any] = {}
if label is not None:
payload["label"] = label
if sql is not None:
payload["sql"] = sql
if schema is not None:
payload["schema"] = schema
if description is not None:
payload["description"] = description
try:
response = await self.client.request(
method="PUT",
endpoint=f"/saved_query/{query_id}",
data=payload,
)
response = cast(dict, response)
app_logger.reflect(f"Saved query {query_id} updated")
return response
except Exception as e:
app_logger.explore("Update saved query failed", error=str(e))
raise
# #endregion SupersetSavedQueries.Update
# #region SupersetSavedQueries.Delete [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Delete a saved SQL query by ID.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: DELETE /api/v1/saved_query/{id}.
async def delete_saved_query(self, query_id: int) -> dict:
with belief_scope("SupersetSavedQueriesMixin.delete_saved_query", f"id={query_id}"):
app_logger.reason("Deleting saved query", extra={"query_id": query_id})
try:
response = await self.client.request(
method="DELETE",
endpoint=f"/saved_query/{query_id}",
)
response = cast(dict, response)
app_logger.reflect(f"Saved query {query_id} deleted")
return response
except Exception as e:
app_logger.explore("Delete saved query failed", error=str(e))
raise
# #endregion SupersetSavedQueries.Delete
# #endregion SupersetSavedQueriesMixin
# #endregion SupersetSavedQueriesMixin

View File

@@ -0,0 +1,265 @@
# #region SupersetSqlLabMixin [C:4] [TYPE Module] [SEMANTICS superset,sql-lab,query,execute,format,export]
# @defgroup Core Module group.
# @BRIEF SQL Lab mixin for SupersetClient — SQL execute, format, results, estimate, CSV export, query history.
# @LAYER Infrastructure
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @RELATION DEPENDS_ON -> [SupersetClient.Safety.DetectDangerousSql]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для SQL-операций.
# @INVARIANT All execute operations pass through DDL/DML guard from safety.py.
import json
from typing import Any, cast
from ..logger import belief_scope, logger as app_logger
from .safety import is_dangerous_sql
app_logger = cast(Any, app_logger)
# #region SupersetSqlLabMixin [C:4] [TYPE Class]
# @defgroup Core Module group.
# @BRIEF Mixin providing SQL Lab execution, formatting, results retrieval, cost estimation, CSV export, and query history.
# @RELATION DEPENDS_ON -> [SupersetClientBase]
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API.
class SupersetSqlLabMixin:
"""Mixin for Superset SQL Lab operations."""
# #region SupersetSqlLab.ExecuteSql [C:4] [TYPE Function]
# @ingroup Core
# @BRIEF Execute a read-only SQL query via Superset SQL Lab with DDL/DML guard.
# @PRE database_id exists. sql is a valid SELECT query.
# @POST Returns SQL Lab execute response dict with results key.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: POST /api/v1/sqllab/execute/.
# @INVARIANT Only SELECT queries pass the DDL/DML guard; dangerous queries rejected.
async def execute_sql(
self,
database_id: int,
sql: str,
schema: str | None = None,
catalog: str | None = None,
tab_name: str | None = None,
template_params: str | None = None,
) -> dict:
with belief_scope("SupersetSqlLabMixin.execute_sql", f"db={database_id}"):
# Guard against DDL/DML
dangerous, keyword = is_dangerous_sql(sql)
if dangerous:
app_logger.explore(
"Dangerous SQL rejected",
extra={"database_id": database_id, "keyword": keyword},
error=f"Query contains dangerous keyword: {keyword}",
)
return {
"error": f"Dangerous SQL rejected: keyword '{keyword}' found. Only SELECT queries are allowed.",
"keyword": keyword,
}
app_logger.reason(
"Executing SQL query",
extra={"database_id": database_id, "sql_len": len(sql)},
)
payload: dict[str, Any] = {
"database_id": database_id,
"sql": sql,
}
if schema:
payload["schema"] = schema
if catalog:
payload["catalog"] = catalog
if tab_name:
payload["tab_name"] = tab_name
if template_params:
payload["template_params"] = template_params
try:
response = await self.client.request(
method="POST",
endpoint="/sqllab/execute/",
data=payload,
)
response = cast(dict, response)
app_logger.reflect(
"SQL executed",
extra={"database_id": database_id, "status": response.get("status")},
)
return response
except Exception as e:
app_logger.explore(
"SQL execution failed",
extra={"database_id": database_id},
error=str(e),
)
raise
# #endregion SupersetSqlLab.ExecuteSql
# #region SupersetSqlLab.FormatSql [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Format/pretty-print a SQL query via Superset SQL Lab.
# @PRE sql is a valid SQL string.
# @POST Returns formatted SQL string.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: POST /api/v1/sqllab/format_sql/.
async def format_sql(self, sql: str) -> str:
with belief_scope("SupersetSqlLabMixin.format_sql"):
app_logger.reason("Formatting SQL", extra={"sql_len": len(sql)})
try:
response = await self.client.request(
method="POST",
endpoint="/sqllab/format_sql/",
data={"sql": sql},
)
response = cast(dict, response)
formatted = response.get("result", sql)
app_logger.reflect("SQL formatted")
return formatted
except Exception as e:
app_logger.explore("SQL format failed", error=str(e))
raise
# #endregion SupersetSqlLab.FormatSql
# #region SupersetSqlLab.GetResults [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Retrieve async SQL Lab results by results key.
# @PRE results_key is a valid key from a prior execute_sql call.
# @POST Returns results dict with data and columns.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: GET /api/v1/sqllab/results/.
async def get_sql_results(self, results_key: str) -> dict:
with belief_scope("SupersetSqlLabMixin.get_sql_results"):
app_logger.reason("Fetching SQL results", extra={"results_key": results_key})
try:
response = await self.client.request(
method="GET",
endpoint="/sqllab/results/",
params={"q": json.dumps({"key": results_key})},
)
response = cast(dict, response)
app_logger.reflect(
"SQL results fetched",
extra={"row_count": len(response.get("data", []))},
)
return response
except Exception as e:
app_logger.explore("Fetch SQL results failed", error=str(e))
raise
# #endregion SupersetSqlLab.GetResults
# #region SupersetSqlLab.EstimateCost [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Estimate the cost/EXPLAIN of a SQL query.
# @PRE database_id exists. sql is valid.
# @POST Returns cost estimation dict.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: POST /api/v1/sqllab/estimate/.
async def estimate_sql_cost(
self,
database_id: int,
sql: str,
schema: str | None = None,
) -> dict:
with belief_scope("SupersetSqlLabMixin.estimate_sql_cost", f"db={database_id}"):
app_logger.reason("Estimating SQL cost", extra={"database_id": database_id})
payload: dict[str, Any] = {
"database_id": database_id,
"sql": sql,
}
if schema:
payload["schema"] = schema
try:
response = await self.client.request(
method="POST",
endpoint="/sqllab/estimate/",
data=payload,
)
response = cast(dict, response)
app_logger.reflect("SQL cost estimated")
return response
except Exception as e:
app_logger.explore("SQL cost estimation failed", error=str(e))
raise
# #endregion SupersetSqlLab.EstimateCost
# #region SupersetSqlLab.ExportCsv [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Export SQL Lab results as CSV by client ID.
# @PRE client_id is a valid SQL Lab client ID.
# @POST Returns raw CSV bytes.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: GET /api/v1/sqllab/export/{id}/.
async def export_sql_csv(self, client_id: str) -> bytes:
with belief_scope("SupersetSqlLabMixin.export_sql_csv"):
app_logger.reason("Exporting SQL CSV", extra={"client_id": client_id})
try:
response = cast(
Any,
await self.client.request(
method="GET",
endpoint=f"/sqllab/export/{client_id}/",
raw_response=True,
),
)
app_logger.reflect("SQL CSV exported")
return response.content if hasattr(response, "content") else response
except Exception as e:
app_logger.explore("Export CSV failed", error=str(e))
raise
# #endregion SupersetSqlLab.ExportCsv
# #region SupersetSqlLab.GetQueryHistory [C:3] [TYPE Function]
# @ingroup Core
# @BRIEF Get query execution history (recent SQL Lab queries).
# @SIDE_EFFECT Асинхронные HTTP-вызовы к Superset API для пагинации.
# @RELATION CALLS -> [SupersetClientBase._fetch_all_pages]
async def get_query_history(self, query: dict | None = None) -> tuple[int, list[dict]]:
with belief_scope("SupersetSqlLabMixin.get_query_history"):
app_logger.reason("Fetching query history")
validated_query = self._validate_query_params(query or {})
paginated_data = await self._fetch_all_pages(
endpoint="/query/",
pagination_options={
"base_query": validated_query,
"results_field": "result",
},
)
total_count = len(paginated_data)
app_logger.reflect(f"Found {total_count} query entries.")
return total_count, paginated_data
# #endregion SupersetSqlLab.GetQueryHistory
# #region SupersetSqlLab.GetQuery [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Get a single query execution detail by ID.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: GET /api/v1/query/{id}.
async def get_query(self, query_id: int) -> dict:
with belief_scope("SupersetSqlLabMixin.get_query", f"id={query_id}"):
app_logger.reason("Fetching query", extra={"query_id": query_id})
try:
response = await self.client.request(
method="GET",
endpoint=f"/query/{query_id}",
)
response = cast(dict, response)
app_logger.reflect(f"Got query {query_id}")
return response
except Exception as e:
app_logger.explore("Fetch query failed", error=str(e))
raise
# #endregion SupersetSqlLab.GetQuery
# #region SupersetSqlLab.StopQuery [C:2] [TYPE Function]
# @ingroup Core
# @BRIEF Stop a running async SQL Lab query by ID.
# @SIDE_EFFECT Асинхронный HTTP-вызов к Superset API: POST /api/v1/query/stop.
async def stop_query(self, query_id: int) -> dict:
with belief_scope("SupersetSqlLabMixin.stop_query", f"id={query_id}"):
app_logger.reason("Stopping query", extra={"query_id": query_id})
try:
response = await self.client.request(
method="POST",
endpoint="/query/stop",
data={"client_id": str(query_id)},
)
response = cast(dict, response)
app_logger.reflect(f"Query {query_id} stopped")
return response
except Exception as e:
app_logger.explore("Stop query failed", error=str(e))
raise
# #endregion SupersetSqlLab.StopQuery
# #endregion SupersetSqlLabMixin
# #endregion SupersetSqlLabMixin

View File

@@ -0,0 +1,107 @@
# #region SupersetClient.Safety [C:2] [TYPE Module] [SEMANTICS superset,safety,sql,validation]
# @defgroup Core Module group.
# @BRIEF Safety guards for SupersetClient operations: DDL/DML blocking, date format validation.
# @LAYER Infrastructure
# @RELATION CALLED_BY -> [SupersetSqlLab.ExecuteSql]
# @RATIONALE Ported from mcp-superset queries.py DDL/DML guard. Strips comments and string
# literals before matching dangerous keywords as whole words anywhere in the query to prevent
# bypasses via CTEs, statement chaining, parentheses, EXPLAIN prefix, or PL/pgSQL DO blocks.
# @REJECTED Regex-only without comment/string stripping was rejected — hidden keywords in comments
# or string literals would be missed, allowing bypasses.
import re
# DDL/DML keywords that must never be executed via the agent.
# Detected as WHOLE WORDS ANYWHERE (not just at the start) to prevent bypasses.
# Order matters: TRUNCATE comes first so the rejection message names the most specific operation.
_DANGEROUS_KEYWORDS = (
"DROP",
"TRUNCATE",
"MERGE",
"DELETE",
"UPDATE",
"INSERT",
"ALTER",
"CREATE",
"GRANT",
"REVOKE",
"COPY",
"DO", # PL/pgSQL anonymous block
"EXECUTE", # dynamic SQL inside a DO block
)
# #region SupersetClient.Safety.StripComments [C:1] [TYPE Function] [SEMANTICS superset,safety,sql]
# @ingroup Core
# @BRIEF Strip SQL comments (single-line -- and multi-line /* */) for DDL/DML detection.
def _strip_sql_comments(sql: str) -> str:
"""Remove single-line and multi-line SQL comments."""
sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL)
sql = re.sub(r"--[^\n]*", " ", sql)
return sql.strip()
# #endregion SupersetClient.Safety.StripComments
# #region SupersetClient.Safety.StripStrings [C:1] [TYPE Function] [SEMANTICS superset,safety,sql]
# @ingroup Core
# @BRIEF Remove single-quoted string literals so keywords inside them don't trigger false positives.
def _strip_sql_strings(sql: str) -> str:
"""Remove SQL single-quoted string literals (handles escaped '' pairs)."""
return re.sub(r"'(?:''|[^'])*'", " ", sql)
# #endregion SupersetClient.Safety.StripStrings
# #region SupersetClient.Safety.DetectDangerousSql [C:2] [TYPE Function] [SEMANTICS superset,safety,sql,guard]
# @ingroup Core
# @BRIEF Return (is_dangerous, keyword) tuple for a given SQL query string.
# @PRE sql is a non-empty string.
# @POST Returns (True, keyword) if dangerous DDL/DML found, or (False, None) if safe.
# @RATIONALE Strips comments and string literals first, then matches each dangerous
# keyword as a whole word anywhere in the statement. Word boundaries keep
# legitimate identifiers safe (e.g. update_date, deleted_at).
def is_dangerous_sql(sql: str) -> tuple[bool, str | None]:
"""Return (is_dangerous_bool, keyword_or_none) for the given SQL query."""
if not sql or not sql.strip():
return False, None
cleaned = _strip_sql_strings(_strip_sql_comments(sql)).upper()
for keyword in _DANGEROUS_KEYWORDS:
if re.search(rf"\b{keyword}\b", cleaned):
return True, keyword
return False, None
# #endregion SupersetClient.Safety.DetectDangerousSql
# Deprecated viz_type patterns — for chart creation validation (future use)
_DEPRECATED_VIZ_TYPES = {
"pivot_table_v2": "pivot_table_v2 is removed in Superset 6.0; use 'pivot_table_v3' instead.",
"filter_box": "filter_box is deprecated; use native dashboard filters instead.",
"world_map": "world_map is deprecated; use 'deckgl' layers instead.",
"deck_multi": "deck_multi is deprecated; use 'deckgl' layers instead.",
"deck_scatter": "deck_scatter is deprecated; use 'deckgl' layers instead.",
"deck_screengrid": "deck_screengrid is deprecated; use 'deckgl' layers instead.",
"deck_grid": "deck_grid is deprecated; use 'deckgl' layers instead.",
"deck_path": "deck_path is deprecated; use 'deckgl' layers instead.",
"deck_polygon": "deck_polygon is deprecated; use 'deckgl' layers instead.",
"deck_hex": "deck_hex is deprecated; use 'deckgl' layers instead.",
"deck_geojson": "deck_geojson is deprecated; use 'deckgl' layers instead.",
"deck_arc": "deck_arc is deprecated; use 'deckgl' layers instead.",
"event_flow": "event_flow is deprecated; use 'graph' chart or other alternatives.",
"paired_ttest": "paired_ttest is deprecated; use 'table' or custom viz instead.",
"rose": "rose chart is deprecated; use 'pie' or 'bar' instead.",
"partition": "partition chart (icicle/tree/sunburst) is deprecated.",
}
# #region SupersetClient.Safety.ValidateVizType [C:1] [TYPE Function] [SEMANTICS superset,safety,chart,validation]
# @ingroup Core
# @BRIEF Check a viz_type string for deprecated patterns; return (is_deprecated, message).
def validate_viz_type(viz_type: str) -> tuple[bool, str | None]:
"""Return (is_deprecated_bool, reason_or_none) for the viz_type."""
if not viz_type:
return False, None
reason = _DEPRECATED_VIZ_TYPES.get(viz_type)
if reason:
return True, reason
return False, None
# #endregion SupersetClient.Safety.ValidateVizType
# #endregion SupersetClient.Safety