diff --git a/providers/common/ai/docs/toolsets/sql.rst b/providers/common/ai/docs/toolsets/sql.rst index 51a8d60207131..93b165209433a 100644 --- a/providers/common/ai/docs/toolsets/sql.rst +++ b/providers/common/ai/docs/toolsets/sql.rst @@ -36,7 +36,9 @@ Curated toolset wrapping - Executes a SQL query and returns bounded, columnar JSON (see :ref:`bounded-query-results`) * - ``check_query`` - - Validates SQL syntax without executing it + - Validates SQL syntax without executing it, with the same statement and + ``allowed_tables`` rules as ``query``, so it accepts writes only when + ``allow_writes=True`` .. code-block:: python diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py index a18149ea951ca..1311c9ee157ee 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py @@ -514,18 +514,7 @@ def _dialect_for_validation(self) -> str | None: def _query(self, sql: str) -> str: hook = self._get_db_hook() dialect = self._dialect_for_validation() - statements: list[Any] | None = None - if not self._allow_writes: - # allow_read_only_metadata lets agents inspect schemas with DESCRIBE/SHOW - # (a common first move) instead of hard-failing; the deep scan still - # rejects any data-modifying statement, including EXPLAIN . - statements = _validate_sql(sql, dialect=dialect, allow_read_only_metadata=True) - elif self._allowed_canonical is not None: - # Writes are allowed but tables are restricted: parse anyway so the - # allow-list still governs which tables a write may touch. - statements = _parse_sql(sql, dialect=dialect) - if statements is not None: - self._enforce_allowed_tables(statements) + self._validate_for_execution(sql, dialect=dialect, require_parse=False) # One row beyond the cap, so "there is more" is knowable without fetching the # rest. strip_sql_string mirrors what get_records did for the hooks that @@ -546,6 +535,31 @@ def _query(self, sql: str) -> str: total_rows=fetch.total_rows, ) + def _validate_for_execution(self, sql: str, *, dialect: str | None, require_parse: bool) -> None: + """ + Apply the checks ``query`` runs before executing ``sql``, raising when it would be refused. + + ``check_query`` shares this so it never reports valid for a statement ``query`` refuses. + ``require_parse`` makes ``check_query`` syntax-check writes too; ``query`` leaves a write + unparsed unless ``allowed_tables`` needs the AST, so a statement sqlglot cannot parse + still reaches the database. + """ + statements: list[Any] | None = None + if not self._allow_writes: + # allow_read_only_metadata lets agents inspect schemas with DESCRIBE/SHOW + # (a common first move) instead of hard-failing; the deep scan still + # rejects any data-modifying statement, including EXPLAIN . + statements = _validate_sql(sql, dialect=dialect, allow_read_only_metadata=True) + elif require_parse or self._allowed_canonical is not None: + # Writes are allowed, so only parse: the allow-list, when set, still + # governs which tables a write may touch. Without one, ``query`` sends the + # string to the hook unparsed, multi-statement included, so don't reject it here. + statements = _parse_sql( + sql, dialect=dialect, allow_multiple_statements=self._allowed_canonical is None + ) + if statements is not None: + self._enforce_allowed_tables(statements) + def _check_query(self, sql: str) -> str: # Resolve the dialect best-effort: if the connection can't be reached we # still syntax-check dialect-agnostically rather than reporting invalid. @@ -553,8 +567,7 @@ def _check_query(self, sql: str) -> str: with suppress(Exception): dialect = self._dialect_for_validation() try: - statements = _validate_sql(sql, dialect=dialect, allow_read_only_metadata=True) - self._enforce_allowed_tables(statements) + self._validate_for_execution(sql, dialect=dialect, require_parse=True) return dumps_masked({"valid": True}) except Exception as e: return dumps_masked({"valid": False, "error": str(e)}) diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py index 53a978cf7e8db..dd16a1c1b17d3 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py @@ -548,6 +548,68 @@ def test_a_row_carrying_the_connection_password_reaches_the_model_masked(self, r class TestSQLToolsetCheckQuery: + @pytest.mark.parametrize( + ("allow_writes", "expected_valid"), + [(False, False), (True, True)], + ) + def test_write_statement_validity_follows_allow_writes(self, allow_writes, expected_valid): + ts = SQLToolset("pg_default", allow_writes=allow_writes) + ts._hook = _make_mock_db_hook() + + result = asyncio.run( + ts.call_tool( + "check_query", + {"sql": "INSERT INTO users VALUES (3, 'Eve')"}, + ctx=MagicMock(), + tool=MagicMock(), + ) + ) + assert json.loads(result)["valid"] is expected_valid + + def test_write_to_allowed_table_is_valid(self): + ts = SQLToolset("pg_default", allowed_tables=["orders"], allow_writes=True) + ts._hook = _make_mock_db_hook() + + assert _run_check(ts, "INSERT INTO orders (id) VALUES (1)")["valid"] is True + + def test_write_to_disallowed_table_is_invalid_and_names_the_table(self): + ts = SQLToolset("pg_default", allowed_tables=["orders"], allow_writes=True) + ts._hook = _make_mock_db_hook() + + data = _run_check(ts, "INSERT INTO secret_salaries (id) VALUES (1)") + + assert data["valid"] is False + assert "secret_salaries" in data["error"] + + def test_malformed_write_is_invalid_for_check_query_but_still_reaches_the_hook_for_query(self): + """check_query syntax-checks writes; query without an allow-list leaves them unparsed.""" + sql = "INSERT INTO users VALUES (" + ts = SQLToolset("pg_default", allow_writes=True) + ts._hook = _make_mock_db_hook(records=[], last_description=None) + + assert _run_check(ts, sql)["valid"] is False + + _run_query(ts, sql) + _assert_executed(ts._hook, sql) + + @pytest.mark.parametrize( + ("toolset_kwargs", "expected_valid"), + [({}, True), ({"allowed_tables": ["users"]}, False)], + ) + def test_multiple_statements_follow_query_when_writes_allowed(self, toolset_kwargs, expected_valid): + ts = SQLToolset("pg_default", allow_writes=True, **toolset_kwargs) + ts._hook = _make_mock_db_hook() + + result = asyncio.run( + ts.call_tool( + "check_query", + {"sql": "INSERT INTO users VALUES (1, 'a'); DELETE FROM users"}, + ctx=MagicMock(), + tool=MagicMock(), + ) + ) + assert json.loads(result)["valid"] is expected_valid + def test_valid_select(self): ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook()