I3eka commented on code in PR #43132: URL: https://github.com/apache/superset/pull/43132#discussion_r4004376020
########## superset/ai/policy.py: ########## @@ -0,0 +1,381 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Guards applied to every tool call before it runs. + +These bound blast radius. They are **not** an authorization layer: a tool that +returns or mutates a specific data-bearing object still has to perform its own +``security_manager.raise_for_access(...)`` check. A policy answers "should this +shape of call be attempted at all", which is a cheaper and coarser question. + +Policies are configured as dotted paths in ``AI_AGENT_TOOL_POLICIES`` so a +deployment can add its own without forking. +""" + +from __future__ import annotations + +import logging +import re +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +logger = logging.getLogger(__name__) + +#: A bare identifier, or dotted parts thereof. Deliberately strict: anything +#: with whitespace, quotes, semicolons or parentheses is rejected rather than +#: escaped, because a tool that needs to escape an identifier is a tool that is +#: building SQL by concatenation. +_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$") + +#: Argument names understood to carry identifiers rather than free text. +_IDENTIFIER_ARGUMENTS = frozenset( + {"table", "table_name", "schema", "schema_name", "catalog", "column", "columns"} +) + + +@dataclass(frozen=True) +class Denial: + """ + A refusal to run a tool call. + + ``reason`` is shown to the model, so it should say what would be acceptable + instead. A model that is told "only read-only SQL is allowed" rewrites its + query; a model that is told "denied" retries the same thing. + """ + + reason: str + + +class ToolPolicy(ABC): + """A pre-execution guard over a single tool call.""" + + #: Identifies the policy in logs. + name: str = "policy" + + @abstractmethod + def check( + self, + tool_name: str, + arguments: dict[str, Any], + ) -> Denial | None: + """ + Inspect a pending call. + + Return ``None`` to allow, or a :class:`Denial` to block. A policy that + does not apply to ``tool_name`` returns ``None``. + """ + + +class ReadOnlySqlPolicy(ToolPolicy): + """ + Refuse anything that is not a read. + + Correctness here rests on Superset's own parser rather than a prefix or + keyword match. A regex over the leading token is defeated by a leading + comment, a CTE that wraps a DML statement, ``EXPLAIN ANALYZE DELETE``, and + multi-statement scripts — all of which the parser handles because the rest + of Superset already depends on it for the same decision. + """ + + name = "read_only_sql" + + #: Tools whose payload is SQL to execute. + sql_tools = frozenset( + { + "execute_sql", + "validate_sql", + "run_scoped_sql", + "create_virtual_dataset", + } + ) + + #: Argument names that may carry the SQL. + sql_arguments = ("sql", "query") + + #: Introspection commands permitted even when the parser cannot model them. + #: + #: Most dialects surface ``EXPLAIN`` and ``SHOW`` as an opaque catch-all + #: node, so a blanket "refuse what we cannot parse" rule would also refuse + #: the schema and query-plan inspection an analysis agent legitimately + #: needs. Enumerating them keeps the default deny for everything else. + read_only_commands = frozenset({"EXPLAIN", "SHOW", "DESCRIBE", "DESC"}) + + def check(self, tool_name: str, arguments: dict[str, Any]) -> Denial | None: + sql_arguments = self._arguments_for_sql_tool(tool_name, arguments) + if sql_arguments is None: + return None + + sql = self._extract_sql(sql_arguments) + if sql is None: + return Denial( + f"{tool_name} requires a 'sql' argument containing the statement " + f"to run." + ) + if not sql.strip(): + return Denial("The 'sql' argument is empty.") + + engine = self._engine(sql_arguments) + + try: + from superset.sql.parse import SQLScript + + script = SQLScript(sql, engine=engine) + except Exception: # pylint: disable=broad-except + # Unparseable SQL cannot be shown to be read-only, so it is + # refused. Logged rather than surfaced: parser errors can quote + # arbitrary query text back to the caller. + logger.info("Refusing unparseable SQL from tool %s", tool_name) + return Denial( + "That SQL could not be parsed. Send a single, syntactically " + "valid read-only statement." + ) + + # Checked per statement so a write cannot ride along behind a read. + for statement in script.statements: + if statement.is_mutating(): + return Denial( + "Only read-only SQL is allowed. Rewrite this as a SELECT — " + "statements that modify data or schema are refused." + ) + + # Anything the parser could not model is refused unless every statement + # is recognisably an introspection command. The mutation check above has + # already run, but it cannot reason about an opaque node on every + # dialect, so this is the fail-closed half of the decision. + if script.has_unparseable_statement and not all( + self._is_read_only_command(statement) for statement in script.statements + ): + return Denial( + "That SQL contains a statement this tool cannot verify as " + "read-only. Send a plain SELECT." + ) + return None + + def _arguments_for_sql_tool( + self, + tool_name: str, + arguments: dict[str, Any], + ) -> dict[str, Any] | None: + """Resolve built-in and namespaced MCP SQL payloads.""" + from superset.ai.mcp.config import split_foreign_tool_name + + parts = split_foreign_tool_name(tool_name) + effective_name = parts[1] if parts is not None else tool_name + if effective_name not in self.sql_tools: + return None + + request = arguments.get("request") + return request if isinstance(request, dict) else arguments + + def _engine(self, arguments: dict[str, Any]) -> str: + """Resolve the parser dialect from the selected Superset database.""" + if engine := arguments.get("engine"): + return str(engine) + + database_id = arguments.get("database_id") + if not isinstance(database_id, int) or isinstance(database_id, bool): + return "" + + try: + from superset.daos.database import DatabaseDAO + + database = DatabaseDAO.find_by_id(database_id) + if database is not None: + return str(database.db_engine_spec.engine or "") + except Exception: # pylint: disable=broad-except + logger.debug("Could not resolve SQL dialect for database %s", database_id) + return "" Review Comment: Fixed in 63c0244d7c: `_engine` no longer accepts a tool/model-supplied override; it resolves the selected database's engine instead. The regression covers both the built-in SQL tool and the namespaced nested virtual-dataset payload, with and without a conflicting `engine`. The two conflicting-dialect cases failed before the change, and all 701 AI unit tests pass after it. The existing SQL parser and mutation checks are unchanged. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
