busya/ss-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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:
busya
2026-06-30 17:26:51 +03:00
parent a0619ca049
commit 79d0f8f678
18 changed files with 3993 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
backend/src/agent/_tool_resolver.py
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

297
backend/src/agent/tools.py
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])

3
backend/src/api/routes/__init__.py
View File

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

284
backend/src/api/routes/agent_superset.py Normal file
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

240
backend/src/api/routes/agent_superset_explore.py Normal file
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

4
backend/src/app.py
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)

14
backend/src/core/superset_client/__init__.py
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.

358
backend/src/core/superset_client/_audit.py Normal file
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

76
backend/src/core/superset_client/_dashboards_crud.py
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

182
backend/src/core/superset_client/_dashboards_write.py
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

329
backend/src/core/superset_client/_databases.py
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

210
backend/src/core/superset_client/_datasets.py
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

160
backend/src/core/superset_client/_saved_queries.py Normal file
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

265
backend/src/core/superset_client/_sql_lab.py Normal file
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

107
backend/src/core/superset_client/safety.py Normal file
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

558
backend/tests/agent/test_superset_tools.py Normal file
View File

@@ -0,0 +1,558 @@
# #region Test.Agent.SupersetTools [C:3] [TYPE Module] [SEMANTICS test,agent,tools,superset,sql,audit,dashboard,dataset]
# @BRIEF Integration tests for new Superset agent tools: SQL, explore DB, audit, create/copy dashboard, create dataset, format SQL.
# @RELATION BINDS_TO -> [AgentChat.Tools]
# @RELATION BINDS_TO -> [AgentChat.ToolResolver]
# @TEST_EDGE: safe_sql -> superset_execute_sql with SELECT passes
# @TEST_EDGE: dangerous_sql -> superset_execute_sql with DROP returns error
# @TEST_EDGE: external_fail -> HTTP error returns error string
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import httpx
# ── Helpers ──
def _mock_httpx_response(status_code=200, json_data=None, text=""):
"""Build a mock httpx.Response with given status, JSON, and text."""
resp = MagicMock(spec=httpx.Response)
resp.status_code = status_code
resp.text = text
resp.json = MagicMock(return_value=json_data or {})
return resp
# ═══════════════════════════════════════════════════════════════════
# Tool name registration
# ═══════════════════════════════════════════════════════════════════
class TestToolRegistration:
"""Test that new tools are registered in get_all_tools() and classification sets."""
def test_all_new_tools_registered(self):
"""All 7 new tools appear in get_all_tools()."""
with patch("src.agent.tools.logger", MagicMock()):
from src.agent.tools import get_all_tools
tools = get_all_tools()
tool_names = {t.name for t in tools}
expected = {
"superset_execute_sql",
"superset_explore_database",
"superset_audit_permissions",
"superset_create_dashboard",
"superset_copy_dashboard",
"superset_create_dataset",
"superset_format_sql",
}
assert expected.issubset(tool_names), f"Missing: {expected - tool_names}"
def test_tool_resolver_sets_contain_new_tools(self):
"""Classification sets in _tool_resolver.py contain the new tools."""
with patch("src.agent.tools.logger", MagicMock()):
from src.agent._tool_resolver import (
_SAFE_AGENT_TOOLS,
_GUARDED_AGENT_TOOLS,
_FAST_CONFIRM_TOOLS,
)
assert "superset_execute_sql" in _SAFE_AGENT_TOOLS
assert "superset_explore_database" in _SAFE_AGENT_TOOLS
assert "superset_audit_permissions" in _SAFE_AGENT_TOOLS
assert "superset_format_sql" in _SAFE_AGENT_TOOLS
assert "superset_create_dashboard" in _GUARDED_AGENT_TOOLS
assert "superset_copy_dashboard" in _GUARDED_AGENT_TOOLS
assert "superset_create_dataset" in _GUARDED_AGENT_TOOLS
assert "superset_explore_database" in _FAST_CONFIRM_TOOLS
assert "superset_audit_permissions" in _FAST_CONFIRM_TOOLS
assert "superset_format_sql" in _FAST_CONFIRM_TOOLS
# ═══════════════════════════════════════════════════════════════════
# Intent matching (get_tools_for_query)
# ═══════════════════════════════════════════════════════════════════
class TestIntentMatching:
"""Test get_tools_for_query matches intent keywords for new tools."""
def _get_for_query(self, query):
with patch("src.agent.tools.logger", MagicMock()):
from src.agent.tools import get_tools_for_query
return {t.name for t in get_tools_for_query(query)}
def test_sql_intent(self):
tools = self._get_for_query("run a SQL query on users")
assert "superset_execute_sql" in tools
def test_schema_intent(self):
tools = self._get_for_query("list all schemas in the database")
assert "superset_explore_database" in tools
def test_table_intent(self):
tools = self._get_for_query("show me the tables in public schema")
assert "superset_explore_database" in tools
def test_audit_intent(self):
tools = self._get_for_query("audit user permissions access rights")
assert "superset_audit_permissions" in tools
def test_create_dashboard_intent(self):
tools = self._get_for_query("create a new dashboard called Sales")
assert "superset_create_dashboard" in tools
def test_copy_dashboard_intent(self):
tools = self._get_for_query("copy dashboard 42")
assert "superset_copy_dashboard" in tools
def test_create_dataset_intent(self):
tools = self._get_for_query("create a new dataset for orders table")
assert "superset_create_dataset" in tools
def test_format_sql_intent(self):
# The intent matcher looks for exact phrase "format sql" in text
tools = self._get_for_query("please format sql for me")
assert "superset_format_sql" in tools
def test_no_match_returns_defaults(self):
tools = self._get_for_query("hello how are you")
assert "show_capabilities" in tools
assert "search_dashboards" in tools
# ═══════════════════════════════════════════════════════════════════
# superset_execute_sql
# ═══════════════════════════════════════════════════════════════════
class TestSupersetExecuteSql:
"""Test superset_execute_sql agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_safe_sql_execution(self):
"""Safe SELECT returns API result."""
mock_resp = _mock_httpx_response(200, text='{"status": "success", "data": []}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_execute_sql
result = await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "SELECT 1",
})
assert "success" in result
@pytest.mark.asyncio
async def test_dangerous_sql_returns_error(self):
"""Dangerous SQL returns error from API (guard is server-side)."""
mock_resp = _mock_httpx_response(200, text='{"error": "Dangerous SQL rejected"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_execute_sql
result = await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "DROP TABLE users",
})
assert "Dangerous" in result
@pytest.mark.asyncio
async def test_optional_schema(self):
"""Schema parameter is passed when provided."""
mock_resp = _mock_httpx_response(200, text="{}")
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_execute_sql
await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "SELECT 1",
"query_schema": "public",
})
assert "schema" in mock_post.call_args[1]["params"]
@pytest.mark.asyncio
async def test_http_error(self):
"""HTTP error returns error string."""
mock_resp = _mock_httpx_response(500, text="Internal Server Error")
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_execute_sql
result = await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "SELECT 1",
})
assert "Error 500" in result
# ═══════════════════════════════════════════════════════════════════
# superset_explore_database
# ═══════════════════════════════════════════════════════════════════
class TestSupersetExploreDatabase:
"""Test superset_explore_database agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_list_schemas(self):
"""Action 'schemas' GETs /api/agent/superset/databases/{id}/schemas."""
mock_resp = _mock_httpx_response(200, text='["public", "analytics"]')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "schemas",
})
assert "public" in result
@pytest.mark.asyncio
async def test_list_tables(self):
"""Action 'tables' GETs /api/agent/superset/databases/{id}/tables."""
mock_resp = _mock_httpx_response(200, text="table1, table2")
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "tables",
"schema_name": "public",
})
assert "table" in result
@pytest.mark.asyncio
async def test_table_metadata(self):
"""Action 'table_metadata' GETs metadata endpoint."""
mock_resp = _mock_httpx_response(200, text='{"columns": []}')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "table_metadata",
"table_name": "users",
})
assert "columns" in result
@pytest.mark.asyncio
async def test_select_star(self):
"""Action 'select_star' GETs select_star endpoint."""
mock_resp = _mock_httpx_response(200, text='"SELECT * FROM users"')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "select_star",
"table_name": "users",
})
assert "SELECT" in result
@pytest.mark.asyncio
async def test_tables_missing_schema_returns_error(self):
"""Missing schema_name for 'tables' returns error without HTTP call."""
with patch("httpx.AsyncClient.get") as mock_get:
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "tables",
})
assert "Error" in result
mock_get.assert_not_called()
@pytest.mark.asyncio
async def test_table_metadata_missing_table_name_returns_error(self):
"""Missing table_name for 'table_metadata' returns error without HTTP call."""
with patch("httpx.AsyncClient.get") as mock_get:
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "table_metadata",
})
assert "Error" in result
mock_get.assert_not_called()
@pytest.mark.asyncio
async def test_unknown_action_returns_error(self):
"""Unknown action returns error without HTTP call."""
with patch("httpx.AsyncClient.get") as mock_get:
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "invalid_action",
})
assert "unknown action" in result.lower()
mock_get.assert_not_called()
# ═══════════════════════════════════════════════════════════════════
# superset_audit_permissions
# ═══════════════════════════════════════════════════════════════════
class TestSupersetAuditPermissions:
"""Test superset_audit_permissions agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_basic_audit(self):
"""Permissions audit GETs /api/agent/superset/audit/permissions."""
mock_resp = _mock_httpx_response(200, text='{"users": [], "total_users": 0}')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_audit_permissions
result = await superset_audit_permissions.ainvoke({
"environment_id": "prod",
})
assert "total_users" in result
@pytest.mark.asyncio
async def test_audit_with_filters(self):
"""Username filter and include_admin are passed as query params."""
mock_resp = _mock_httpx_response(200, text="{}")
with patch("httpx.AsyncClient.get") as mock_get:
mock_get.return_value = mock_resp
from src.agent.tools import superset_audit_permissions
await superset_audit_permissions.ainvoke({
"environment_id": "prod",
"username_filter": "alice",
"include_admin": True,
})
call_params = mock_get.call_args[1]["params"]
assert call_params["username_filter"] == "alice"
assert call_params["include_admin"] == "true"
@pytest.mark.asyncio
async def test_audit_without_filters(self):
"""Without filters, no additional params beyond environment_id."""
mock_resp = _mock_httpx_response(200, text="{}")
with patch("httpx.AsyncClient.get") as mock_get:
mock_get.return_value = mock_resp
from src.agent.tools import superset_audit_permissions
await superset_audit_permissions.ainvoke({
"environment_id": "prod",
})
call_params = mock_get.call_args[1]["params"]
assert "username_filter" not in call_params
assert "include_admin" not in call_params
# ═══════════════════════════════════════════════════════════════════
# superset_create_dashboard
# ═══════════════════════════════════════════════════════════════════
class TestSupersetCreateDashboard:
"""Test superset_create_dashboard agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_create_dashboard(self):
"""Creates dashboard via POST to /api/agent/superset/dashboards."""
mock_resp = _mock_httpx_response(201, text='{"id": 1, "dashboard_title": "Sales"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_create_dashboard
result = await superset_create_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_title": "Sales",
"slug": "sales",
"published": True,
})
assert "Sales" in result
@pytest.mark.asyncio
async def test_create_dashboard_minimal(self):
"""Minimal create with just title."""
mock_resp = _mock_httpx_response(201, text='{"id": 2}')
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_create_dashboard
await superset_create_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_title": "Minimal",
})
call_params = mock_post.call_args[1]["params"]
assert call_params["dashboard_title"] == "Minimal"
assert call_params["published"] == "false"
# ═══════════════════════════════════════════════════════════════════
# superset_copy_dashboard
# ═══════════════════════════════════════════════════════════════════
class TestSupersetCopyDashboard:
"""Test superset_copy_dashboard agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_copy_dashboard(self):
"""Copies dashboard via POST to /api/agent/superset/dashboards/{id}/copy."""
mock_resp = _mock_httpx_response(201, text='{"id": 2, "result": "copied"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_copy_dashboard
result = await superset_copy_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_id": 1,
"dashboard_title": "Copy of Sales",
})
assert "copied" in result
@pytest.mark.asyncio
async def test_copy_dashboard_without_title(self):
"""Copy without optional title."""
mock_resp = _mock_httpx_response(201, text="{}")
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_copy_dashboard
await superset_copy_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_id": 1,
})
call_params = mock_post.call_args[1]["params"]
assert "dashboard_title" not in call_params
# ═══════════════════════════════════════════════════════════════════
# superset_create_dataset
# ═══════════════════════════════════════════════════════════════════
class TestSupersetCreateDataset:
"""Test superset_create_dataset agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_create_dataset(self):
"""Creates dataset via POST to /api/agent/superset/datasets."""
mock_resp = _mock_httpx_response(201, text='{"id": 5, "table_name": "orders"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_create_dataset
result = await superset_create_dataset.ainvoke({
"environment_id": "prod",
"table_name": "orders",
"database": 1,
"schema_name": "public",
})
assert "orders" in result
@pytest.mark.asyncio
async def test_create_dataset_minimal(self):
"""Minimal create without schema."""
mock_resp = _mock_httpx_response(201, text='{"id": 6}')
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_create_dataset
await superset_create_dataset.ainvoke({
"environment_id": "prod",
"table_name": "users",
"database": 1,
})
call_params = mock_post.call_args[1]["params"]
assert "schema_name" not in call_params
# ═══════════════════════════════════════════════════════════════════
# superset_format_sql
# ═══════════════════════════════════════════════════════════════════
class TestSupersetFormatSql:
"""Test superset_format_sql agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_format_sql(self):
"""Formats SQL via POST to /api/agent/superset/sqllab/format."""
mock_resp = _mock_httpx_response(200, text="SELECT\n 1\nFROM\n users\n")
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_format_sql
result = await superset_format_sql.ainvoke({
"environment_id": "prod",
"sql": "SELECT 1 FROM users",
})
assert "SELECT" in result
@pytest.mark.asyncio
async def test_format_sql_error(self):
"""HTTP error returns error string."""
mock_resp = _mock_httpx_response(500, text="Internal Error")
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_format_sql
result = await superset_format_sql.ainvoke({
"environment_id": "prod",
"sql": "INVALID",
})
assert "Error 500" in result
# ═══════════════════════════════════════════════════════════════════
# Tool Resolver inference
# ═══════════════════════════════════════════════════════════════════
class TestToolResolverInference:
"""Test _tool_resolver.py inference with new Superset tool patterns."""
def test_infer_from_text_sql(self):
"""SQL-related query infers superset_execute_sql."""
with patch("src.agent.tools.logger", MagicMock()):
from src.agent._tool_resolver import infer_tool_from_text
# The resolver doesn't currently have SQL inference — verify fallback
# but ensure it's extended if needed in the future
result = infer_tool_from_text("sql select query users")
# No specific SQL inference yet — returns None or base tool
# This test documents current behavior; update if resolver is extended
assert result is None or isinstance(result, str)
# #endregion Test.Agent.SupersetTools

651
backend/tests/test_core/test_superset_extended.py Normal file
View File

@@ -0,0 +1,651 @@
# #region Test.Superset.Extended [C:3] [TYPE Module] [SEMANTICS test,superset,sql-lab,dashboard,dataset,database,saved-query]
# @BRIEF Unit tests for SupersetClient extended mixins: SQL Lab, dashboards write, datasets, databases, saved queries, audit.
# @RELATION BINDS_TO -> [SupersetSqlLabMixin]
# @RELATION BINDS_TO -> [SupersetDashboardsWriteMixin]
# @RELATION BINDS_TO -> [SupersetDatasetsMixin]
# @RELATION BINDS_TO -> [SupersetDatabasesMixin]
# @RELATION BINDS_TO -> [SupersetSavedQueriesMixin]
# @RELATION BINDS_TO -> [SupersetAuditMixin]
# @TEST_EDGE: ddl_guard -> execute_sql blocks dangerous SQL via safety.py
# @TEST_EDGE: external_fail -> Exception in API call is raised
# @TEST_EDGE: missing_field -> Optional parameters omitted from payloads
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
# ── Shared fixture: base mixin with mocked client ──
@pytest.fixture
def mock_client():
"""Return a mocked AsyncAPIClient with AsyncMock request method."""
client = MagicMock()
client.request = AsyncMock()
return client
def _setup_mixin(mixin_class, mock_client):
"""Instantiate a mixin class and assign the mocked client."""
m = mixin_class.__new__(mixin_class)
m.client = mock_client
m._validate_query_params = MagicMock(return_value={"page": 0, "page_size": 100})
m._fetch_all_pages = AsyncMock(return_value=[{"id": 1}, {"id": 2}])
return m
# ═══════════════════════════════════════════════════════════════════
# SQL Lab mixin tests
# ═══════════════════════════════════════════════════════════════════
class TestSupersetSqlLabMixin:
"""Test SupersetSqlLabMixin methods."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._sql_lab.belief_scope", MagicMock()), \
patch("src.core.superset_client._sql_lab.app_logger", MagicMock()):
from src.core.superset_client._sql_lab import SupersetSqlLabMixin
return _setup_mixin(SupersetSqlLabMixin, mock_client)
@pytest.mark.asyncio
async def test_execute_sql_safe_query(self, mixin, mock_client):
"""Safe SELECT passes DDL guard and POSTs to /sqllab/execute/."""
mock_client.request.return_value = {"status": "success", "data": []}
with patch("src.core.superset_client._sql_lab.is_dangerous_sql", return_value=(False, None)):
result = await mixin.execute_sql(database_id=1, sql="SELECT 1")
assert result == {"status": "success", "data": []}
mock_client.request.assert_called_once()
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/sqllab/execute/"
assert call_args["data"]["database_id"] == 1
assert call_args["data"]["sql"] == "SELECT 1"
@pytest.mark.asyncio
async def test_execute_sql_dangerous_query_blocked(self, mixin, mock_client):
"""DDL query is blocked by guard, no API call made."""
with patch("src.core.superset_client._sql_lab.is_dangerous_sql", return_value=(True, "DROP")):
result = await mixin.execute_sql(database_id=1, sql="DROP TABLE users")
assert "error" in result
assert "DROP" in result["error"]
assert result["keyword"] == "DROP"
mock_client.request.assert_not_called()
@pytest.mark.asyncio
async def test_execute_sql_with_optional_params(self, mixin, mock_client):
"""Schema, catalog, tab_name, template_params are included when provided."""
mock_client.request.return_value = {"status": "success"}
with patch("src.core.superset_client._sql_lab.is_dangerous_sql", return_value=(False, None)):
await mixin.execute_sql(
database_id=1,
sql="SELECT 1",
schema="public",
catalog="pg_catalog",
tab_name="My Tab",
template_params='{"param": "value"}',
)
data = mock_client.request.call_args[1]["data"]
assert data["schema"] == "public"
assert data["catalog"] == "pg_catalog"
assert data["tab_name"] == "My Tab"
assert data["template_params"] == '{"param": "value"}'
@pytest.mark.asyncio
async def test_execute_sql_raises_on_api_error(self, mixin, mock_client):
"""Exception from client.request is re-raised."""
mock_client.request.side_effect = RuntimeError("API down")
with patch("src.core.superset_client._sql_lab.is_dangerous_sql", return_value=(False, None)):
with pytest.raises(RuntimeError, match="API down"):
await mixin.execute_sql(database_id=1, sql="SELECT 1")
@pytest.mark.asyncio
async def test_format_sql(self, mixin, mock_client):
"""format_sql POSTs to /sqllab/format_sql/ and returns formatted string."""
mock_client.request.return_value = {"result": "SELECT\n 1\nFROM\n users\n"}
result = await mixin.format_sql("SELECT 1 FROM users")
assert "SELECT" in result
assert "users" in result
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/sqllab/format_sql/"
assert call_args["data"]["sql"] == "SELECT 1 FROM users"
@pytest.mark.asyncio
async def test_get_sql_results(self, mixin, mock_client):
"""get_sql_results GETs /sqllab/results/ with results key."""
mock_client.request.return_value = {"data": [{"id": 1}], "columns": ["id"]}
result = await mixin.get_sql_results("abc123")
assert result["data"] == [{"id": 1}]
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/sqllab/results/"
@pytest.mark.asyncio
async def test_estimate_sql_cost(self, mixin, mock_client):
"""estimate_sql_cost POSTs to /sqllab/estimate/."""
mock_client.request.return_value = {"estimated_cost": 100}
result = await mixin.estimate_sql_cost(database_id=1, sql="SELECT * FROM huge_table")
assert result["estimated_cost"] == 100
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/sqllab/estimate/"
@pytest.mark.asyncio
async def test_export_sql_csv(self, mixin, mock_client):
"""export_sql_csv GETs /sqllab/export/{id}/ with raw_response=True."""
mock_response = MagicMock()
mock_response.content = b"csv,data,here"
mock_client.request.return_value = mock_response
result = await mixin.export_sql_csv("client123")
assert result == b"csv,data,here"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/sqllab/export/client123/"
assert call_args.get("raw_response") is True
@pytest.mark.asyncio
async def test_get_query_history(self, mixin):
"""get_query_history delegates to _fetch_all_pages for /query/."""
result = await mixin.get_query_history()
assert result[0] == 2 # count from fixture
assert len(result[1]) == 2
@pytest.mark.asyncio
async def test_get_query(self, mixin, mock_client):
"""get_query GETs /query/{id}."""
mock_client.request.return_value = {"id": 42, "sql": "SELECT 1"}
result = await mixin.get_query(42)
assert result["id"] == 42
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/query/42"
@pytest.mark.asyncio
async def test_stop_query(self, mixin, mock_client):
"""stop_query POSTs to /query/stop with client_id."""
mock_client.request.return_value = {"status": "stopped"}
result = await mixin.stop_query(42)
assert result["status"] == "stopped"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/query/stop"
assert call_args["data"]["client_id"] == "42"
# ═══════════════════════════════════════════════════════════════════
# Dashboards Write mixin tests
# ═══════════════════════════════════════════════════════════════════
class TestSupersetDashboardsWriteMixin:
"""Test SupersetDashboardsWriteMixin new methods."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._dashboards_write.belief_scope", MagicMock()), \
patch("src.core.superset_client._dashboards_write.app_logger", MagicMock()):
from src.core.superset_client._dashboards_write import SupersetDashboardsWriteMixin
return _setup_mixin(SupersetDashboardsWriteMixin, mock_client)
@pytest.mark.asyncio
async def test_create_dashboard(self, mixin, mock_client):
"""create_dashboard POSTs to /dashboard/ with title and published."""
mock_client.request.return_value = {"id": 1, "dashboard_title": "My Dashboard"}
result = await mixin.create_dashboard(dashboard_title="My Dashboard", published=True, slug="my-dash")
assert result["id"] == 1
assert result["dashboard_title"] == "My Dashboard"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/dashboard/"
assert call_args["data"]["dashboard_title"] == "My Dashboard"
assert call_args["data"]["published"] is True
assert call_args["data"]["slug"] == "my-dash"
@pytest.mark.asyncio
async def test_create_dashboard_minimal(self, mixin, mock_client):
"""create_dashboard with only required fields."""
mock_client.request.return_value = {"id": 2}
await mixin.create_dashboard(dashboard_title="Minimal")
assert "css" not in mock_client.request.call_args[1]["data"]
@pytest.mark.asyncio
async def test_update_dashboard(self, mixin, mock_client):
"""update_dashboard PUTs to /dashboard/{id}."""
mock_client.request.return_value = {"id": 1, "dashboard_title": "Updated"}
result = await mixin.update_dashboard(dashboard_id=1, dashboard_title="Updated", published=False)
assert result["dashboard_title"] == "Updated"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "PUT"
assert call_args["endpoint"] == "/dashboard/1"
assert call_args["data"]["dashboard_title"] == "Updated"
assert call_args["data"]["published"] is False
@pytest.mark.asyncio
async def test_copy_dashboard(self, mixin, mock_client):
"""copy_dashboard POSTs to /dashboard/{id}/copy/."""
mock_client.request.return_value = {"id": 2, "dashboard_title": "Copy of Dashboard"}
result = await mixin.copy_dashboard(dashboard_id=1, dashboard_title="Copy of Dashboard")
assert result["id"] == 2
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/dashboard/1/copy/"
assert call_args["data"]["dashboard_title"] == "Copy of Dashboard"
@pytest.mark.asyncio
async def test_publish_dashboard(self, mixin, mock_client):
"""publish_dashboard PUTs {published: True} to /dashboard/{id}."""
mock_client.request.return_value = {"id": 1, "published": True}
result = await mixin.publish_dashboard(1)
assert result["published"] is True
call_args = mock_client.request.call_args[1]
assert call_args["data"]["published"] is True
@pytest.mark.asyncio
async def test_unpublish_dashboard(self, mixin, mock_client):
"""unpublish_dashboard PUTs {published: False} to /dashboard/{id}."""
mock_client.request.return_value = {"id": 1, "published": False}
result = await mixin.unpublish_dashboard(1)
assert result["published"] is False
call_args = mock_client.request.call_args[1]
assert call_args["data"]["published"] is False
# ═══════════════════════════════════════════════════════════════════
# Dashboards CRUD mixin tests (new methods)
# ═══════════════════════════════════════════════════════════════════
class TestSupersetDashboardsCrudExtended:
"""Test SupersetDashboardsCrudMixin new methods."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._dashboards_crud.belief_scope", MagicMock()), \
patch("src.core.superset_client._dashboards_crud.app_logger", MagicMock()):
from src.core.superset_client._dashboards_crud import SupersetDashboardsCrudMixin
return _setup_mixin(SupersetDashboardsCrudMixin, mock_client)
@pytest.mark.asyncio
async def test_get_dashboard_charts(self, mixin, mock_client):
"""get_dashboard_charts GETs /dashboard/{id}/charts."""
mock_client.request.return_value = {"result": [{"id": 10}, {"id": 11}]}
result = await mixin.get_dashboard_charts(1)
assert len(result) == 2
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/dashboard/1/charts"
@pytest.mark.asyncio
async def test_get_dashboard_datasets(self, mixin, mock_client):
"""get_dashboard_datasets GETs /dashboard/{id}/datasets."""
mock_client.request.return_value = {"result": [{"id": 100}]}
result = await mixin.get_dashboard_datasets(1)
assert len(result) == 1
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/dashboard/1/datasets"
@pytest.mark.asyncio
async def test_get_dashboard(self, mixin, mock_client):
"""get_dashboard GETs /dashboard/{ref}."""
mock_client.request.return_value = {"id": 1, "dashboard_title": "Test Dash"}
result = await mixin.get_dashboard(1)
assert result["dashboard_title"] == "Test Dash"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/dashboard/1"
# ═══════════════════════════════════════════════════════════════════
# Datasets mixin tests (new methods)
# ═══════════════════════════════════════════════════════════════════
class TestSupersetDatasetsExtended:
"""Test SupersetDatasetsMixin extended methods."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._datasets.belief_scope", MagicMock()), \
patch("src.core.superset_client._datasets.app_logger", MagicMock()):
from src.core.superset_client._datasets import SupersetDatasetsMixin
return _setup_mixin(SupersetDatasetsMixin, mock_client)
@pytest.mark.asyncio
async def test_create_dataset(self, mixin, mock_client):
"""create_dataset POSTs to /dataset/ with table_name and database."""
mock_client.request.return_value = {"id": 5, "table_name": "new_table"}
result = await mixin.create_dataset(table_name="new_table", database=1, schema_name="public")
assert result["id"] == 5
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/dataset/"
assert call_args["data"]["table_name"] == "new_table"
assert call_args["data"]["database"] == 1
assert call_args["data"]["schema"] == "public"
@pytest.mark.asyncio
async def test_create_dataset_with_sql(self, mixin, mock_client):
"""create_dataset includes sql field for virtual datasets."""
mock_client.request.return_value = {"id": 6}
await mixin.create_dataset(table_name="virtual_table", database=1, sql="SELECT 1")
assert mock_client.request.call_args[1]["data"]["sql"] == "SELECT 1"
@pytest.mark.asyncio
async def test_delete_dataset(self, mixin, mock_client):
"""delete_dataset DELETEs /dataset/{id}."""
mock_client.request.return_value = {"message": "OK"}
result = await mixin.delete_dataset(dataset_id=5)
assert result["message"] == "OK"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "DELETE"
assert call_args["endpoint"] == "/dataset/5"
@pytest.mark.asyncio
async def test_duplicate_dataset(self, mixin, mock_client):
"""duplicate_dataset POSTs to /dataset/duplicate."""
mock_client.request.return_value = {"id": 7}
result = await mixin.duplicate_dataset(base_model_id=5, table_name="table_copy")
assert result["id"] == 7
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/dataset/duplicate"
assert call_args["data"]["base_model_id"] == 5
assert call_args["data"]["table_name"] == "table_copy"
@pytest.mark.asyncio
async def test_refresh_dataset_schema(self, mixin, mock_client):
"""refresh_dataset_schema PUTs to /dataset/{id}/refresh."""
mock_client.request.return_value = {"result": "refreshed"}
result = await mixin.refresh_dataset_schema(dataset_id=5)
assert result["result"] == "refreshed"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "PUT"
assert call_args["endpoint"] == "/dataset/5/refresh"
# ═══════════════════════════════════════════════════════════════════
# Databases mixin tests (new methods)
# ═══════════════════════════════════════════════════════════════════
class TestSupersetDatabasesExtended:
"""Test SupersetDatabasesMixin extended methods."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._databases.belief_scope", MagicMock()), \
patch("src.core.superset_client._databases.app_logger", MagicMock()):
from src.core.superset_client._databases import SupersetDatabasesMixin
return _setup_mixin(SupersetDatabasesMixin, mock_client)
@pytest.mark.asyncio
async def test_get_database_schemas(self, mixin, mock_client):
"""get_database_schemas GETs /database/{id}/schemas/."""
mock_client.request.return_value = {"result": ["public", "analytics"]}
result = await mixin.get_database_schemas(database_id=1)
assert result == ["public", "analytics"]
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/database/1/schemas/"
@pytest.mark.asyncio
async def test_get_database_tables(self, mixin, mock_client):
"""get_database_tables GETs /database/{id}/tables/ with optional schema."""
mock_client.request.return_value = {"result": [{"name": "users"}, {"name": "orders"}]}
result = await mixin.get_database_tables(database_id=1, schema_name="public")
assert len(result) == 2
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/database/1/tables/"
@pytest.mark.asyncio
async def test_get_database_select_star(self, mixin, mock_client):
"""get_database_select_star GETs /database/{id}/select_star/."""
mock_client.request.return_value = {"result": "SELECT * FROM public.users"}
result = await mixin.get_database_select_star(database_id=1, table_name="users", schema_name="public")
assert "SELECT *" in result
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/database/1/select_star/"
@pytest.mark.asyncio
async def test_get_database_table_metadata(self, mixin, mock_client):
"""get_database_table_metadata GETs /database/{id}/table_metadata/."""
mock_client.request.return_value = {
"columns": [{"name": "id", "type": "INTEGER"}],
}
result = await mixin.get_database_table_metadata(database_id=1, table_name="users")
assert "columns" in result
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/database/1/table_metadata/"
@pytest.mark.asyncio
async def test_validate_sql(self, mixin, mock_client):
"""validate_sql POSTs to /database/{id}/validate_sql/."""
mock_client.request.return_value = {"result": "Valid"}
result = await mixin.validate_sql(database_id=1, sql="SELECT 1")
assert result["result"] == "Valid"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/database/1/validate_sql/"
assert call_args["data"]["sql"] == "SELECT 1"
# ═══════════════════════════════════════════════════════════════════
# Saved Queries mixin tests
# ═══════════════════════════════════════════════════════════════════
class TestSupersetSavedQueriesMixin:
"""Test SupersetSavedQueriesMixin CRUD methods."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._saved_queries.belief_scope", MagicMock()), \
patch("src.core.superset_client._saved_queries.app_logger", MagicMock()):
from src.core.superset_client._saved_queries import SupersetSavedQueriesMixin
return _setup_mixin(SupersetSavedQueriesMixin, mock_client)
@pytest.mark.asyncio
async def test_get_saved_queries(self, mixin):
"""get_saved_queries delegates to _fetch_all_pages for /saved_query/."""
result = await mixin.get_saved_queries()
assert result[0] == 2
assert len(result[1]) == 2
@pytest.mark.asyncio
async def test_get_saved_query(self, mixin, mock_client):
"""get_saved_query GETs /saved_query/{id}."""
mock_client.request.return_value = {"id": 1, "label": "My Query"}
result = await mixin.get_saved_query(1)
assert result["label"] == "My Query"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "GET"
assert call_args["endpoint"] == "/saved_query/1"
@pytest.mark.asyncio
async def test_create_saved_query(self, mixin, mock_client):
"""create_saved_query POSTs to /saved_query/."""
mock_client.request.return_value = {"id": 10, "label": "New Query"}
result = await mixin.create_saved_query(
label="New Query", db_id=1, sql="SELECT 1", schema="public", description="test"
)
assert result["id"] == 10
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "POST"
assert call_args["endpoint"] == "/saved_query/"
assert call_args["data"]["label"] == "New Query"
@pytest.mark.asyncio
async def test_update_saved_query(self, mixin, mock_client):
"""update_saved_query PUTs to /saved_query/{id}."""
mock_client.request.return_value = {"id": 1, "label": "Updated Query"}
result = await mixin.update_saved_query(query_id=1, label="Updated Query")
assert result["label"] == "Updated Query"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "PUT"
assert call_args["endpoint"] == "/saved_query/1"
@pytest.mark.asyncio
async def test_delete_saved_query(self, mixin, mock_client):
"""delete_saved_query DELETEs /saved_query/{id}."""
mock_client.request.return_value = {"message": "Deleted"}
result = await mixin.delete_saved_query(1)
assert result["message"] == "Deleted"
call_args = mock_client.request.call_args[1]
assert call_args["method"] == "DELETE"
assert call_args["endpoint"] == "/saved_query/1"
# ═══════════════════════════════════════════════════════════════════
# Audit mixin tests
# ═══════════════════════════════════════════════════════════════════
class TestSupersetAuditMixin:
"""Test SupersetAuditMixin permissions audit and helpers."""
@pytest.fixture
def mixin(self, mock_client):
with patch("src.core.superset_client._audit.belief_scope", MagicMock()), \
patch("src.core.superset_client._audit.app_logger", MagicMock()):
from src.core.superset_client._audit import SupersetAuditMixin
m = SupersetAuditMixin.__new__(SupersetAuditMixin)
m.client = mock_client
m._validate_query_params = MagicMock(return_value={"page": 0, "page_size": 100})
m._fetch_all_pages = AsyncMock(return_value=[])
return m
@pytest.mark.asyncio
async def test_permissions_audit_returns_structure(self, mixin):
"""permissions_audit returns audit dict with expected keys when all data is empty."""
mixin._fetch_users = AsyncMock(return_value=(0, []))
mixin.get_dashboards = AsyncMock(return_value=(0, []))
mixin.get_datasets = AsyncMock(return_value=(0, []))
mixin._build_role_permissions_map = AsyncMock(return_value={})
mixin._find_datasource_permissions = AsyncMock(return_value={})
# Mock remaining direct client.request calls
mixin.client.request.side_effect = [
{"result": []}, # groups
{"result": []}, # RLS
]
result = await mixin.permissions_audit(include_admin=False)
assert "users" in result
assert "total_users" in result
assert result["total_users"] == 0
assert "dashboards_checked" in result
assert "datasets_checked" in result
@pytest.mark.asyncio
async def test_permissions_audit_with_users(self, mixin):
"""permissions_audit handles a single user with no dashboards/datasets."""
mixin._fetch_users = AsyncMock(return_value=(1, [
{"id": 1, "username": "alice", "roles": [], "active": True}
]))
mixin.get_dashboards = AsyncMock(return_value=(0, []))
mixin.get_datasets = AsyncMock(return_value=(0, []))
mixin._build_role_permissions_map = AsyncMock(return_value={})
mixin._find_datasource_permissions = AsyncMock(return_value={})
# groups + RLS
mixin.client.request.side_effect = [
{"result": []}, # groups
{"result": []}, # RLS
]
result = await mixin.permissions_audit(include_admin=False)
assert result["total_users"] == 1
assert len(result["users"]) == 1
assert result["users"][0]["username"] == "alice"
@pytest.mark.asyncio
async def test_permissions_audit_page_size_capped(self, mixin):
"""page_size is capped at 50 in permissions_audit."""
mixin._fetch_users = AsyncMock(return_value=(0, []))
mixin.get_dashboards = AsyncMock(return_value=(0, []))
mixin.get_datasets = AsyncMock(return_value=(0, []))
mixin._build_role_permissions_map = AsyncMock(return_value={})
mixin._find_datasource_permissions = AsyncMock(return_value={})
mixin.client.request.side_effect = [
{"result": []}, # groups
{"result": []}, # RLS
]
result = await mixin.permissions_audit(page=0, page_size=100)
assert result["page_size"] == 50 # capped at 50 by min(page_size, 50)
# Should still work with no users
@pytest.mark.asyncio
async def test_fetch_users(self, mixin):
"""_fetch_users delegates to _fetch_all_pages for security API."""
result = await mixin._fetch_users()
assert result[0] == 0 # mock returns empty list
mixin._fetch_all_pages.assert_awaited_once()
@pytest.mark.asyncio
async def test_build_role_permissions_map(self, mixin, mock_client):
"""_build_role_permissions_map fetches roles and their permissions."""
mock_client.request.side_effect = [
{"result": [{"id": 1, "name": "Admin"}]}, # roles
{"result": [{"id": 100}, {"id": 101}]}, # perms for role 1
]
result = await mixin._build_role_permissions_map()
assert result == {1: {100, 101}}
@pytest.mark.asyncio
async def test_find_datasource_permissions(self, mixin, mock_client):
"""_find_datasource_permissions scans permissions-resources with pagination."""
mock_client.request.return_value = {
"result": [
{
"id": 500,
"permission": {"name": "datasource_access"},
"view_menu": {"name": "my_table (id:42)"},
},
{
"id": 501,
"permission": {"name": "not_datasource_access"},
"view_menu": {"name": "other (id:99)"},
},
]
}
result = await mixin._find_datasource_permissions()
assert result == {42: 500}
# #endregion Test.Superset.Extended

243
backend/tests/test_core/test_superset_safety.py Normal file
View File

@@ -0,0 +1,243 @@
# #region Test.Superset.Safety [C:3] [TYPE Module] [SEMANTICS test,safety,sql,ddl,dml,guard]
# @BRIEF Unit tests for core/superset_client/safety.py — DDL/DML guard and viz_type validation.
# @RELATION BINDS_TO -> [SupersetClient.Safety]
# @TEST_EDGE: empty_sql -> Empty/whitespace SQL passes (no dangerous keywords)
# @TEST_EDGE: comment_bypass -> DDL keywords inside comments are stripped before matching
# @TEST_EDGE: string_literal_bypass -> DDL keywords inside string literals don't trigger false positives
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
import pytest
from src.core.superset_client.safety import (
is_dangerous_sql,
validate_viz_type,
_strip_sql_comments,
_strip_sql_strings,
)
class TestDangerousSqlDetection:
"""Test DDL/DML detection: safe queries pass, dangerous queries blocked."""
# ── Safe queries ──
@pytest.mark.parametrize("sql", [
"SELECT * FROM users",
"select * from users where id = 1",
"SELECT u.name, o.total FROM users u JOIN orders o ON u.id = o.user_id",
"SELECT COUNT(*) FROM users",
" SELECT 1 ",
])
def test_safe_select_queries_pass(self, sql):
"""@POST Safe SELECT queries return (False, None)."""
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
assert keyword is None
# ── Dangerous queries blocked ──
@pytest.mark.parametrize("sql, expected_keyword", [
("DROP TABLE users", "DROP"),
("drop table users", "DROP"),
(" DROP TABLE users ", "DROP"),
("DELETE FROM users", "DELETE"),
("UPDATE users SET name = 'test'", "UPDATE"),
("INSERT INTO users VALUES (1, 'test')", "INSERT"),
("TRUNCATE TABLE users", "TRUNCATE"),
("ALTER TABLE users ADD COLUMN age INT", "ALTER"),
("CREATE TABLE test (id INT)", "CREATE"),
("GRANT SELECT ON users TO bob", "GRANT"),
("REVOKE SELECT ON users FROM bob", "REVOKE"),
("MERGE INTO target t USING source s ON t.id = s.id", "MERGE"),
("COPY (SELECT * FROM users) TO '/tmp/out.csv'", "COPY"),
("DO $$ BEGIN DELETE FROM users; END $$", "DELETE"), # DELETE found before DO since $$-quoted strings are not stripped
("EXECUTE 'DROP TABLE users'", "EXECUTE"),
])
def test_dangerous_queries_blocked(self, sql, expected_keyword):
"""@POST Dangerous DDL/DML returns (True, keyword)."""
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == expected_keyword
# ── Bypass prevention ──
def test_comment_before_drop_blocked(self):
"""Keywords after comments are detected (comments stripped first)."""
sql = "-- This is just a comment\nDROP TABLE users"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == "DROP"
def test_multiline_comment_hides_drop_blocked(self):
"""Multi-line comment /* ... */ before DROP still detected."""
sql = "/* clean up */ DROP TABLE users"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == "DROP"
def test_cte_with_delete_blocked(self):
"""CTE containing DELETE is detected as dangerous (whole word anywhere)."""
sql = "WITH deleted AS (DELETE FROM temp RETURNING *) SELECT * FROM deleted"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == "DELETE"
def test_parenthesized_delete_blocked(self):
"""Parentheses around DELETE don't hide it from detection."""
sql = "SELECT * FROM (DELETE FROM users) t"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == "DELETE"
def test_explain_prefix_does_not_hide_drop(self):
"""EXPLAIN before DROP still detects DROP as whole word."""
sql = "EXPLAIN DROP TABLE users"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == "DROP"
# ── False positive prevention ──
def test_string_literal_containing_drop_not_flagged(self):
"""DROP inside a string literal is not flagged as dangerous."""
sql = "SELECT 'DROP TABLE users' AS warning"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
def test_comment_containing_drop_not_flagged(self):
"""DROP inside a comment alone should not be flagged."""
sql = "-- DROP TABLE users"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
def test_update_date_identifier_not_flagged(self):
"""update_date column name does NOT trigger UPDATE keyword (whole-word match)."""
sql = "SELECT update_date FROM users"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
def test_deleted_at_identifier_not_flagged(self):
"""deleted_at column name does NOT trigger DELETE keyword (whole-word match)."""
sql = "SELECT deleted_at FROM users WHERE deleted_at IS NULL"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
def test_dropped_at_identifier_not_flagged(self):
"""dropped_at column name does NOT trigger DROP keyword (whole-word match)."""
sql = "SELECT dropped_at FROM events"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
# ── Edge cases ──
@pytest.mark.parametrize("sql", [
"",
" ",
"\n",
"\t",
" \n \t ",
])
def test_empty_or_whitespace_sql_passes(self, sql):
"""Empty or whitespace-only SQL returns (False, None)."""
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
assert keyword is None
def test_case_insensitive_detection(self):
"""Mixed-case dangerous keywords are detected."""
sql = "Drop Table users"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is True
assert keyword == "DROP"
def test_keywords_embedded_in_other_words_not_flagged(self):
"""DROP appearing as substring of another word is NOT flagged."""
sql = "SELECT dropbox_path FROM settings"
dangerous, keyword = is_dangerous_sql(sql)
assert dangerous is False
class TestStripSqlComments:
"""Unit tests for the _strip_sql_comments helper."""
def test_strip_single_line_comment(self):
result = _strip_sql_comments("SELECT 1 -- comment")
assert "comment" not in result
assert "SELECT 1" in result
def test_strip_multiline_comment(self):
result = _strip_sql_comments("SELECT 1 /* inline */, 2")
assert "inline" not in result
assert "SELECT 1" in result
assert ", 2" in result
def test_strip_multiline_comment_multiline(self):
result = _strip_sql_comments("SELECT /* line1\nline2 */ 1")
assert "line1" not in result
assert "line2" not in result
assert "1" in result
def test_no_comments_unchanged(self):
result = _strip_sql_comments("SELECT 1 FROM users")
assert result.strip() == "SELECT 1 FROM users"
class TestStripSqlStrings:
"""Unit tests for the _strip_sql_strings helper."""
def test_strip_simple_string(self):
result = _strip_sql_strings("SELECT 'hello world' AS greeting")
assert "hello world" not in result
assert "SELECT" in result
def test_strip_string_with_escaped_quotes(self):
result = _strip_sql_strings("SELECT 'it''s ok' AS msg")
assert "it's ok" not in result
assert "SELECT" in result
def test_no_strings_unchanged(self):
result = _strip_sql_strings("SELECT id FROM users")
assert result.strip() == "SELECT id FROM users"
class TestValidateVizType:
"""Unit tests for validate_viz_type."""
def test_valid_viz_type_passes(self):
deprecated, reason = validate_viz_type("table")
assert deprecated is False
assert reason is None
def test_valid_bar_chart_passes(self):
deprecated, reason = validate_viz_type("echarts_timeseries_bar")
assert deprecated is False
assert reason is None
def test_deprecated_pivot_table_rejected(self):
deprecated, reason = validate_viz_type("pivot_table_v2")
assert deprecated is True
assert "pivot_table_v3" in reason
def test_deprecated_filter_box_rejected(self):
deprecated, reason = validate_viz_type("filter_box")
assert deprecated is True
assert "native dashboard filters" in reason
def test_deprecated_world_map_rejected(self):
deprecated, reason = validate_viz_type("world_map")
assert deprecated is True
def test_empty_viz_type_passes(self):
deprecated, reason = validate_viz_type("")
assert deprecated is False
assert reason is None
def test_unknown_viz_type_passes(self):
deprecated, reason = validate_viz_type("non_existent_viz")
assert deprecated is False
assert reason is None
# #endregion Test.Superset.Safety