Lee-W commented on code in PR #74487:
URL: https://github.com/apache/airflow/pull/74487#discussion_r4236879323
##########
providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py:
##########
@@ -614,3 +617,113 @@ def test_flags_dynamic_sql_as_unverifiable(self, sql,
dialect):
scan = collect_table_references(parse_sql(sql, dialect=dialect))
assert scan.tables == []
assert scan.unverifiable_sources
+
+
+# sqlglot has no upper bound in pyproject.toml, so an upstream rename, split
or removal of an
+# ``exp`` class would otherwise make the read-only check silently stop
recognising a write.
+# The tests below turn that drift into a CI failure; if one fails after a
sqlglot bump, review
+# ``_DATA_MODIFYING_NODES`` / ``DEFAULT_ALLOWED_TYPES`` in
``utils/sql_validation.py``.
+_DIALECTS = ["postgres", "mysql", "snowflake", "bigquery", "sqlite"]
+
+_WRITE_STATEMENTS = {
+ "insert": "INSERT INTO t (a) VALUES (1)",
+ "update": "UPDATE t SET a = 1",
+ "delete": "DELETE FROM t",
+ "merge": "MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN UPDATE SET
a = 1",
+ "drop": "DROP TABLE t",
+ "truncate": "TRUNCATE TABLE t",
+ "create": "CREATE TABLE t (id INT)",
+ "create_as_select": "CREATE TABLE t AS SELECT 1",
+ "alter": "ALTER TABLE t ADD COLUMN c INT",
+ "copy": "COPY t FROM '/tmp/x'",
+ "grant": "GRANT SELECT ON t TO u",
+ "cte_delete": "WITH d AS (DELETE FROM t RETURNING *) SELECT * FROM d",
+ "cte_insert": "WITH d AS (INSERT INTO t VALUES (1) RETURNING *) SELECT *
FROM d",
+ "select_into": "SELECT * INTO n FROM t",
+}
+
+_READ_STATEMENTS = {
+ "select": "SELECT a FROM t",
+ "cte_select": "WITH c AS (SELECT 1 AS a) SELECT * FROM c",
+ "union_all": "SELECT 1 UNION ALL SELECT 2",
+}
+
+
+class TestSqlglotDriftTripwires:
+ """Fail loudly if a sqlglot upgrade silently changes what the read-only
check recognises."""
+
+ @pytest.mark.parametrize(
+ "node_type",
+ [*DEFAULT_ALLOWED_TYPES, *READ_ONLY_METADATA_TYPES,
*_DATA_MODIFYING_NODES],
+ ids=lambda t: t.__name__,
+ )
+ def test_referenced_node_is_still_a_public_expression_class(self,
node_type):
+ assert isinstance(node_type, type)
+ assert issubclass(node_type, exp.Expression), (
+ f"{node_type!r} is no longer an exp.Expression subclass; review
utils/sql_validation.py"
+ )
+ assert getattr(exp, node_type.__name__, None) is node_type, (
+ f"exp.{node_type.__name__} no longer resolves to the class
sql_validation.py references; "
+ "review utils/sql_validation.py"
+ )
+
+ def test_allowed_types_never_overlap_data_modifying_nodes(self):
+ """A write type slipping into the allow-list would bypass the
statement-type check."""
+ for allowed in DEFAULT_ALLOWED_TYPES:
+ for denied in _DATA_MODIFYING_NODES:
+ assert not issubclass(allowed, denied), f"{allowed.__name__}
is a kind of {denied.__name__}"
+ assert not issubclass(denied, allowed), f"{denied.__name__} is
a kind of {allowed.__name__}"
+
+ @pytest.mark.parametrize(
+ ("sql", "node_type"),
+ [
+ (_WRITE_STATEMENTS["insert"], exp.Insert),
+ (_WRITE_STATEMENTS["update"], exp.Update),
+ (_WRITE_STATEMENTS["delete"], exp.Delete),
+ (_WRITE_STATEMENTS["merge"], exp.Merge),
+ (_WRITE_STATEMENTS["drop"], exp.Drop),
+ (_WRITE_STATEMENTS["truncate"], exp.TruncateTable),
+ (_WRITE_STATEMENTS["create"], exp.Create),
+ (_WRITE_STATEMENTS["alter"], exp.Alter),
+ (_WRITE_STATEMENTS["select_into"], exp.Into),
+ ],
+ ids=["insert", "update", "delete", "merge", "drop", "truncate",
"create", "alter", "select_into"],
Review Comment:
It takes a single `_WRITE_NODE_TYPES` mapping with `ids=list(...)`, and
covers `create_as_select`, `cte_delete` and `cte_insert`.
`copy` and `grant` are left out. They are not in `_DATA_MODIFYING_NODES`, so
only the top-level type check stops them. The expected rejecting guard lives in
`_WRITE_REJECTION`.
A key missing from either mapping raises `KeyError`.
--
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]
