This is an automated email from the ASF dual-hosted git repository.
kaxil 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 72502d61652 Keep SQL deep scan on for custom read-only allowed_types
(#76688)
72502d61652 is described below
commit 72502d61652da7b9410e830de4d9b2d867fd7600
Author: Jyun-An Chen <[email protected]>
AuthorDate: Sun Nov 8 10:52:12 2026 +0800
Keep SQL deep scan on for custom read-only allowed_types (#76688)
* Keep SQL deep scan on for custom read-only allowed_types
validate_sql() skipped the data-modifying deep scan whenever allowed_types
was not exactly DEFAULT_ALLOWED_TYPES. A narrower or reordered read-only
list such as (exp.Select,) therefore let statements like
"WITH d AS (DELETE ...) SELECT ..." or "SELECT ... INTO ..." through.
Skip the deep scan only when the allow-list itself admits a data-modifying
type.
* Keep allow-lists that admit writes through a base class working
An allow-list such as (exp.Select, exp.DML) admits an INSERT at the type
check, but the one-way issubclass test left the deep scan on and rejected
the same INSERT as read-only. Treat an allowed type that is a write node or
a base of one as admitting writes. Scanning only the write types the list
does not admit is not an option: MERGE nests Insert and Update, and ALTER
... DROP COLUMN nests Drop, so it would reject writes the list allows.
Also update the denylist comment and stability.rst to the new rule.
---------
Co-authored-by: Kaxil Naik <[email protected]>
---
providers/common/ai/docs/stability.rst | 5 ++-
.../providers/common/ai/utils/sql_validation.py | 27 +++++++-----
.../unit/common/ai/utils/test_sql_validation.py | 50 ++++++++++++++++++----
3 files changed, 61 insertions(+), 21 deletions(-)
diff --git a/providers/common/ai/docs/stability.rst
b/providers/common/ai/docs/stability.rst
index f3e6e415730..733ccedc052 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -98,8 +98,9 @@ Pydantic AI toolsets, but the Pydantic AI class they inherit
from can change.
- Generates SQL from the prompt and returns it without running it. With
``validate_sql=True``, the default, the task fails unless the output is
a single
statement that parses for ``dialect`` and is one of
``allowed_sql_types``
- (``SELECT``, ``UNION``, ``INTERSECT`` and ``EXCEPT`` by default); with
the default
- types, a write nested inside the statement fails it too. A reviewer's
edit under
+ (``SELECT``, ``UNION``, ``INTERSECT`` and ``EXCEPT`` by default); unless
+ ``allowed_sql_types`` admits a write, a write nested inside the
statement fails it
+ too. A reviewer's edit under
``allow_modifications=True`` is checked the same way. ``db_conn_id``
with
``table_names``, ``schema_context`` or ``datasource_config`` give the
model the
schema. ``datasource_config`` reads object storage through the
DataFusion engine of
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py
index 17d0120b15e..225b5a31308 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py
@@ -90,7 +90,7 @@ READ_ONLY_METADATA_TYPES: tuple[type[exp.Expr], ...] = (
# Note: exp.Command is sqlglot's fallback for any syntax it doesn't recognize.
# Including it makes the denylist fail-closed (safer), but may block legitimate
# vendor-specific SQL that sqlglot can't parse. Callers who need such syntax
can
-# provide custom allowed_types to bypass the deep scan entirely.
+# add exp.Command to a custom allowed_types, which also turns the deep scan
off.
_DATA_MODIFYING_NODES: tuple[type[exp.Expr], ...] = (
exp.Insert,
exp.Update,
@@ -433,7 +433,9 @@ def validate_sql(
:param allowed_types: Tuple of sqlglot expression types to permit.
Defaults to ``(Select, Union, Intersect, Except)``. When supplied, the
caller takes full control of the allow-list and
``allow_read_only_metadata``
- is ignored.
+ is ignored. The deep scan for data-modifying nodes still runs unless
the
+ allow-list admits a write, directly (``exp.Insert``) or through a base
class
+ (``exp.DML``).
:param dialect: SQL dialect for parsing (``postgres``, ``mysql``, etc.).
:param allow_multiple_statements: Whether to allow multiple
semicolon-separated
statements. Default ``False``.
@@ -445,18 +447,23 @@ def validate_sql(
:raises SQLSafetyError: If the SQL is empty, contains disallowed statement
types,
or has multiple statements when not permitted.
"""
- # A caller-supplied ``allowed_types`` is an explicit opt-out of the curated
- # read-only defaults (and the data-modifying deep scan). Otherwise we use
the
- # read-only defaults, optionally widened with metadata statements, and keep
- # the deep scan on.
if allowed_types is None:
types: tuple[type[exp.Expr], ...] = DEFAULT_ALLOWED_TYPES
if allow_read_only_metadata:
types = types + READ_ONLY_METADATA_TYPES
- run_data_modifying_scan = True
else:
types = allowed_types
- run_data_modifying_scan = types == DEFAULT_ALLOWED_TYPES
+ # Only an allow-list that admits a write opts out of the deep scan,
whether it names
+ # the write type (``exp.Insert``) or a base of it (``exp.DML``). Comparing
to
+ # DEFAULT_ALLOWED_TYPES would also switch it off for a narrower or
reordered
+ # read-only list such as ``(exp.Select,)``, letting ``WITH d AS (DELETE
...)
+ # SELECT ...`` through. Scanning a list that admits writes would instead
reject the
+ # writes it allows: MERGE nests Insert and Update, ALTER ... DROP COLUMN
nests Drop.
+ run_data_modifying_scan = not any(
+ issubclass(allowed, node) or issubclass(node, allowed)
+ for allowed in types
+ for node in _DATA_MODIFYING_NODES
+ )
parsed = parse_sql(sql, dialect=dialect,
allow_multiple_statements=allow_multiple_statements)
@@ -468,9 +475,7 @@ def validate_sql(
)
# Deep scan: reject data-modifying nodes hidden inside otherwise-allowed
statements
- # (e.g. data-modifying CTEs, SELECT INTO, EXPLAIN <write>). Runs for the
curated
- # read-only allow-list — callers who provide custom allowed_types have
explicitly
- # opted into non-read-only operations.
+ # (e.g. data-modifying CTEs, SELECT INTO, EXPLAIN <write>).
if run_data_modifying_scan:
_check_for_data_modifying_nodes(parsed)
diff --git
a/providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py
b/providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py
index 2fd6d31ce41..5fc3ca114e0 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py
@@ -211,14 +211,48 @@ class TestDataModifyingNodeDetection:
result = validate_sql("SELECT * FROM users WHERE id IN (SELECT user_id
FROM orders)")
assert len(result) == 1
- def test_deep_scan_runs_with_explicit_default_types(self):
- """Deep scan should also block when DEFAULT_ALLOWED_TYPES is passed
explicitly."""
- with pytest.raises(SQLSafetyError, match="Data-modifying operation
'Delete'"):
- validate_sql(
- "WITH del AS (DELETE FROM users RETURNING *) SELECT * FROM
del",
- allowed_types=DEFAULT_ALLOWED_TYPES,
- dialect="postgres",
- )
+ @pytest.mark.parametrize(
+ "allowed_types",
+ [
+ pytest.param(DEFAULT_ALLOWED_TYPES, id="explicit-default"),
+ pytest.param((exp.Select,), id="narrower"),
+ pytest.param(tuple(reversed(DEFAULT_ALLOWED_TYPES)),
id="reordered"),
+ pytest.param((*DEFAULT_ALLOWED_TYPES, exp.Describe),
id="widened-read-only"),
+ ],
+ )
+ @pytest.mark.parametrize(
+ ("sql", "node"),
+ [
+ ("WITH del AS (DELETE FROM users RETURNING *) SELECT * FROM del",
"Delete"),
+ ("SELECT * INTO new_table FROM users", "Into"),
+ ],
+ )
+ def test_deep_scan_runs_with_custom_read_only_types(self, allowed_types,
sql, node):
+ """A custom allow-list with no write type keeps the deep scan on."""
+ with pytest.raises(SQLSafetyError, match=f"Data-modifying operation
'{node}'"):
+ validate_sql(sql, allowed_types=allowed_types, dialect="postgres")
+
+ @pytest.mark.parametrize(
+ ("allowed_types", "sql", "expected_type"),
+ [
+ pytest.param((exp.Select, exp.Insert), "INSERT INTO t SELECT *
FROM s", exp.Insert, id="leaf"),
+ pytest.param((exp.Select, exp.DML), "INSERT INTO t SELECT * FROM
s", exp.Insert, id="base-class"),
+ pytest.param(
+ (exp.Select, exp.Merge),
+ "MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN UPDATE
SET v = s.v "
+ "WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)",
+ exp.Merge,
+ id="merge-nests-insert-and-update",
+ ),
+ pytest.param(
+ (exp.Select, exp.Alter), "ALTER TABLE t DROP COLUMN c",
exp.Alter, id="alter-nests-drop"
+ ),
+ ],
+ )
+ def test_deep_scan_skipped_when_custom_types_admit_writes(self,
allowed_types, sql, expected_type):
+ """An allow-list that admits a write, directly or through a base
class, turns the deep scan off."""
+ (stmt,) = validate_sql(sql, allowed_types=allowed_types,
dialect="postgres")
+ assert isinstance(stmt, expected_type)
class TestReadOnlyMetadata:
