Lee-W commented on code in PR #74487:
URL: https://github.com/apache/airflow/pull/74487#discussion_r4242100748
##########
providers/common/ai/tests/unit/common/ai/utils/test_sql_validation.py:
##########
@@ -614,3 +614,144 @@ 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 or split
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",
+}
+
+# Node each write must parse to, keyed like _WRITE_STATEMENTS. copy and grant
are absent
+# on purpose: they are not in _DATA_MODIFYING_NODES and are only stopped by
the top-level type check.
+_WRITE_NODE_TYPES = {
+ "insert": exp.Insert,
+ "update": exp.Update,
+ "delete": exp.Delete,
+ "merge": exp.Merge,
+ "drop": exp.Drop,
+ "truncate": exp.TruncateTable,
+ "create": exp.Create,
+ "create_as_select": exp.Create,
+ "alter": exp.Alter,
+ "cte_delete": exp.Delete,
+ "cte_insert": exp.Insert,
+ "select_into": exp.Into,
+}
+
+_TOP_LEVEL_GUARD = "Statement type '{}' is not allowed"
+_DEEP_SCAN_GUARD = "Data-modifying operation '{}'"
+
+# Guard that must reject each write: the top-level type check, or the deep
scan for writes
+# nested inside an otherwise-allowed SELECT.
+_WRITE_REJECTION = {
+ "insert": _TOP_LEVEL_GUARD.format("Insert"),
+ "update": _TOP_LEVEL_GUARD.format("Update"),
+ "delete": _TOP_LEVEL_GUARD.format("Delete"),
+ "merge": _TOP_LEVEL_GUARD.format("Merge"),
+ "drop": _TOP_LEVEL_GUARD.format("Drop"),
+ "truncate": _TOP_LEVEL_GUARD.format("TruncateTable"),
+ "create": _TOP_LEVEL_GUARD.format("Create"),
+ "create_as_select": _TOP_LEVEL_GUARD.format("Create"),
+ "alter": _TOP_LEVEL_GUARD.format("Alter"),
+ "copy": _TOP_LEVEL_GUARD.format("Copy"),
+ "grant": _TOP_LEVEL_GUARD.format("Grant"),
+ "cte_delete": _DEEP_SCAN_GUARD.format("Delete"),
+ "cte_insert": _DEEP_SCAN_GUARD.format("Insert"),
+ "select_into": _DEEP_SCAN_GUARD.format("Into"),
+}
+
+_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."""
+
+ 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("dialect", _DIALECTS)
+ @pytest.mark.parametrize(
+ ("key", "node_type"), list(_WRITE_NODE_TYPES.items()),
ids=list(_WRITE_NODE_TYPES)
+ )
+ def test_write_statement_parses_to_a_denied_node(self, key, node_type,
dialect):
+ """sqlglot must keep producing the node class the deny-list names for
each write."""
+ sql = _WRITE_STATEMENTS[key]
+ parsed = parse_sql(sql, dialect=dialect)
+ found = [n for stmt in parsed for n in stmt.walk() if isinstance(n,
node_type)]
+ assert found, f"{sql!r} no longer parses to {node_type.__name__};
review _DATA_MODIFYING_NODES"
+ assert issubclass(node_type, _DATA_MODIFYING_NODES)
+
+ @pytest.mark.parametrize("allow_read_only_metadata", [False, True])
+ @pytest.mark.parametrize("dialect", _DIALECTS)
+ @pytest.mark.parametrize("key", list(_WRITE_STATEMENTS))
+ def test_write_and_ddl_statements_rejected_in_every_dialect(self, key,
dialect, allow_read_only_metadata):
+ with pytest.raises(SQLSafetyError, match=_WRITE_REJECTION[key]):
+ validate_sql(
+ _WRITE_STATEMENTS[key], dialect=dialect,
allow_read_only_metadata=allow_read_only_metadata
+ )
+
+ @pytest.mark.parametrize("allow_read_only_metadata", [False, True])
+ @pytest.mark.parametrize(
+ ("sql", "dialect", "node_name"),
+ [
+ ("DESCRIBE ALTER TABLE t ADD COLUMN c INT", "postgres", "Alter"),
+ ("DESCRIBE ALTER TABLE t ADD COLUMN c INT", "mysql", "Alter"),
+ ("EXPLAIN DELETE FROM t", "mysql", "Delete"),
Review Comment:
Added a `DESCRIBE ALTER SYSTEM SET x=1` row
(`describe_alter_system_postgres`) to the wrapped-write matrix. Removing
`exp.Command` from `_DATA_MODIFYING_NODES` now fails its metadata-mode case.
Dropped `test_explain_wrapped_write_still_blocked` as the
`explain_delete_mysql` row covers it.
--
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]
