Make SQLToolset check_query follow allow_writes like query

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
 

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 <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,15 +535,39 @@ 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 <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.
         dialect: str | None = None
         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)})

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()