Skip to content

Make SQLToolset check_query follow allow_writes like query - #74448

Merged
Lee-W merged 3 commits into
apache:mainfrom
astronomer:common-ai-sql-check-query-allow-writes
Oct 9, 2026
Merged

Lee-W merged 3 commits into
apache:mainfrom
astronomer:common-ai-sql-check-query-allow-writes

Conversation

@Lee-W

@Lee-W Lee-W commented Oct 8, 2026 •

Copy link
Copy Markdown
Member

Why

With allow_writes=True, SQLToolset.query executes write statements, but check_query still reported them as invalid. An agent that validated before running was told not to run a statement it was permitted to run.

What

query and check_query now share _validate_for_execution(sql, *, dialect, require_parse), so check_query applies the same statement and allowed_tables rules as query.

  • query passes require_parse=False, so its behavior is unchanged. With allow_writes=True and no allowed_tables it does not parse, so statements sqlglot cannot parse are still sent to the database.
  • check_query passes require_parse=True: write statements are accepted when allow_writes=True, and still get a syntax check. Without allowed_tables it also accepts multiple statements, as query does; with allowed_tables both reject them.
  • With allow_writes=False (the default), both still reject write statements.
  • docs/toolsets/sql.rst describes the new check_query behavior.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: [Claude] following the guidelines


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

Comment thread providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py Outdated
Comment thread providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
Comment thread providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py Outdated
Lee-W added 3 commits October 9, 2026 05:54
With `allow_writes=True`, `check_query` reported write statements as invalid
even though `query` executes them, so an agent that validated before running
was told not to run a statement it was permitted to run. `check_query` now
applies the same statement and `allowed_tables` rules as `query`.
…` would

With `allow_writes=True` and no `allowed_tables`, `query` sends the string to
the hook unparsed, but `check_query` parsed it with the single-statement
guard and reported `INSERT ...; DELETE ...` as invalid. Only reject multiple
statements when an allow-list is set, which is when `query` parses too.
Cover that `check_query` syntax-checks writes while `query` still sends
them to the hook unparsed, and that a write to a table outside
`allowed_tables` is rejected with that table named.
@Lee-W
Lee-W force-pushed the common-ai-sql-check-query-allow-writes branch from bf64784 to 536e848 Compare October 9, 2026 04:54
@Lee-W
Lee-W merged commit 7f945b4 into apache:main Oct 9, 2026
85 checks passed
@Lee-W
Lee-W deleted the common-ai-sql-check-query-allow-writes branch October 9, 2026 06:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants