kaxil commented on code in PR #74448:
URL: https://github.com/apache/airflow/pull/74448#discussion_r4218073926
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py:
##########
@@ -548,6 +548,42 @@ 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):
Review Comment:
Nothing pins `require_parse`: flipping `check_query`'s `True` to `False`, or
`query`'s `False` to `True`, leaves every test in this file passing. Could you
add one case with `allow_writes=True` and a malformed write like `INSERT INTO
users VALUES (`, asserting `check_query` returns `valid: false` and `query`
still reaches the hook (`_assert_executed`)? That covers both halves of the
behavior described in the PR.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -546,15 +535,36 @@ 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.
+ statements = _parse_sql(sql, dialect=dialect)
Review Comment:
With `allow_writes=True` and no `allowed_tables`, only `check_query` reaches
this branch, and `_parse_sql` defaults to `allow_multiple_statements=False`. So
`INSERT INTO t VALUES (1); DELETE FROM t` comes back `valid: false` ("Multiple
statements detected") while `query` sends the same string to the hook, which is
the mismatch this PR set out to remove and contradicts the new "same statement
rules as `query`" line in `sql.rst`. Passing
`allow_multiple_statements=self._allowed_canonical is None` here would fix it
without touching `query`, since `query` only gets here when the allow-list is
set. If rejecting multi-statement is intended, the docs row should say so
instead.
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py:
##########
@@ -548,6 +548,42 @@ 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
+
+ @pytest.mark.parametrize(
+ ("table", "expected_valid"),
+ [("orders", True), ("secret_salaries", False)],
Review Comment:
The `secret_salaries` case also passes without this change, since the old
`check_query` rejected every INSERT by statement type. Asserting
`"secret_salaries" in data["error"]` (as
`test_check_query_reports_disallowed_table_as_invalid` does) would make it fail
if writes ever stopped going through the allow-list.
--
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]
