diff --git a/backend/src/agent/_tool_resolver.py b/backend/src/agent/_tool_resolver.py index dea52317..8e29fc10 100644 --- a/backend/src/agent/_tool_resolver.py +++ b/backend/src/agent/_tool_resolver.py @@ -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 diff --git a/backend/src/agent/tools.py b/backend/src/agent/tools.py index a7348363..a1f2b105 100644 --- a/backend/src/agent/tools.py +++ b/backend/src/agent/tools.py @@ -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]) diff --git a/backend/src/api/routes/__init__.py b/backend/src/api/routes/__init__.py index 4cc37b08..5a430c09 100755 --- a/backend/src/api/routes/__init__.py +++ b/backend/src/api/routes/__init__.py @@ -21,6 +21,9 @@ __all__ = [ "admin", "admin_api_keys", + "agent_conversations", + "agent_superset", + "agent_superset_explore", "assistant", "clean_release", "clean_release_v2", diff --git a/backend/src/api/routes/agent_superset.py b/backend/src/api/routes/agent_superset.py new file mode 100644 index 00000000..cbbeb9a1 --- /dev/null +++ b/backend/src/api/routes/agent_superset.py @@ -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 diff --git a/backend/src/api/routes/agent_superset_explore.py b/backend/src/api/routes/agent_superset_explore.py new file mode 100644 index 00000000..92684d5a --- /dev/null +++ b/backend/src/api/routes/agent_superset_explore.py @@ -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 diff --git a/backend/src/app.py b/backend/src/app.py index 02510f84..7ca86e74 100755 --- a/backend/src/app.py +++ b/backend/src/app.py @@ -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) diff --git a/backend/src/core/superset_client/__init__.py b/backend/src/core/superset_client/__init__.py index b70d3d7d..0e2e010e 100644 --- a/backend/src/core/superset_client/__init__.py +++ b/backend/src/core/superset_client/__init__.py @@ -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. diff --git a/backend/src/core/superset_client/_audit.py b/backend/src/core/superset_client/_audit.py new file mode 100644 index 00000000..a9be0c34 --- /dev/null +++ b/backend/src/core/superset_client/_audit.py @@ -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 diff --git a/backend/src/core/superset_client/_dashboards_crud.py b/backend/src/core/superset_client/_dashboards_crud.py index 330be60e..7dfa9220 100644 --- a/backend/src/core/superset_client/_dashboards_crud.py +++ b/backend/src/core/superset_client/_dashboards_crud.py @@ -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 diff --git a/backend/src/core/superset_client/_dashboards_write.py b/backend/src/core/superset_client/_dashboards_write.py index 6536308e..9c25968a 100644 --- a/backend/src/core/superset_client/_dashboards_write.py +++ b/backend/src/core/superset_client/_dashboards_write.py @@ -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 diff --git a/backend/src/core/superset_client/_databases.py b/backend/src/core/superset_client/_databases.py index 0303f7e1..2581e845 100644 --- a/backend/src/core/superset_client/_databases.py +++ b/backend/src/core/superset_client/_databases.py @@ -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 diff --git a/backend/src/core/superset_client/_datasets.py b/backend/src/core/superset_client/_datasets.py index f804d391..1d8ce485 100644 --- a/backend/src/core/superset_client/_datasets.py +++ b/backend/src/core/superset_client/_datasets.py @@ -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 diff --git a/backend/src/core/superset_client/_saved_queries.py b/backend/src/core/superset_client/_saved_queries.py new file mode 100644 index 00000000..2fb14a09 --- /dev/null +++ b/backend/src/core/superset_client/_saved_queries.py @@ -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 diff --git a/backend/src/core/superset_client/_sql_lab.py b/backend/src/core/superset_client/_sql_lab.py new file mode 100644 index 00000000..f714e7c2 --- /dev/null +++ b/backend/src/core/superset_client/_sql_lab.py @@ -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 diff --git a/backend/src/core/superset_client/safety.py b/backend/src/core/superset_client/safety.py new file mode 100644 index 00000000..fd453266 --- /dev/null +++ b/backend/src/core/superset_client/safety.py @@ -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 diff --git a/backend/tests/agent/test_superset_tools.py b/backend/tests/agent/test_superset_tools.py new file mode 100644 index 00000000..fcf6d821 --- /dev/null +++ b/backend/tests/agent/test_superset_tools.py @@ -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 diff --git a/backend/tests/test_core/test_superset_extended.py b/backend/tests/test_core/test_superset_extended.py new file mode 100644 index 00000000..80c6684c --- /dev/null +++ b/backend/tests/test_core/test_superset_extended.py @@ -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 diff --git a/backend/tests/test_core/test_superset_safety.py b/backend/tests/test_core/test_superset_safety.py new file mode 100644 index 00000000..1fc3caa7 --- /dev/null +++ b/backend/tests/test_core/test_superset_safety.py @@ -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