This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7f945b45d07 Make `SQLToolset` `check_query` follow `allow_writes` like
`query` (#74448)
7f945b45d07 is described below
commit 7f945b45d0798f1599083b8e8db785835cc28d3d
Author: Wei Lee <[email protected]>
AuthorDate: Fri Oct 9 15:34:10 2026 +0900
Make `SQLToolset` `check_query` follow `allow_writes` like `query` (#74448)
---
providers/common/ai/docs/toolsets/sql.rst | 4 +-
.../airflow/providers/common/ai/toolsets/sql.py | 41 +++++++++-----
.../ai/tests/unit/common/ai/toolsets/test_sql.py | 62 ++++++++++++++++++++++
3 files changed, 92 insertions(+), 15 deletions(-)
diff --git a/providers/common/ai/docs/toolsets/sql.rst
b/providers/common/ai/docs/toolsets/sql.rst
index 51a8d602071..93b16520943 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 a18149ea951..1311c9ee157 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 @@ class SQLToolset(AirflowToolset):
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 <write>.
- 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 @@ class SQLToolset(AirflowToolset):
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 <write>.
+ 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 @@ class SQLToolset(AirflowToolset):
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 53a978cf7e8..dd16a1c1b17 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 @@ class TestSQLToolsetQuery:
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()
