pavansai1997 commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4506791620
##########
providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py:
##########
@@ -412,6 +435,84 @@ def collect_table_references(
return TableScan(tables=tables, unverifiable_sources=unverifiable)
+#: Case-folded ``(catalog, schema, table)`` key compared against an
``allowed_tables`` list.
+TableKey = tuple[str | None, str | None, str]
+
+
+def canonical_table_key(catalog: str | None, schema: str | None, table: str)
-> TableKey:
+ """
+ Normalize a ``(catalog, schema, table)`` reference to its case-folded
comparison key.
+
+ Identifiers are case-folded because databases reflect them in their own
case (Snowflake
+ stores unquoted names uppercase but reflects them lowercased), so a
byte-exact match
+ against the user's entries would silently miss.
+ """
+ return (
+ catalog.casefold() if catalog else None,
+ schema.casefold() if schema else None,
+ table.casefold(),
+ )
+
+
+def canonical_allowed_tables(
+ allowed_tables: Iterable[str], default_schema: str | None = None
+) -> frozenset[TableKey]:
+ """
+ Return the comparison keys for ``allowed_tables`` entries (``table`` or
``schema.table``).
+
+ Unqualified entries resolve to ``default_schema`` so that ``"orders"`` and
+ ``"<schema>.orders"`` denote the same table. Entries carry no catalog, so
any
+ catalog-qualified reference (``otherdb.public.orders``) cannot match --
that closes
+ cross-database access a single-connection allow-list cannot describe.
+ """
+ keys: set[TableKey] = set()
+ for entry in allowed_tables:
+ entry_schema, sep, table = entry.rpartition(".")
+ keys.add(canonical_table_key(None, entry_schema if sep else
default_schema, table))
+ return frozenset(keys)
+
+
+class AllowedTablesViolations(NamedTuple):
+ """Result of :func:`find_allowed_tables_violations`; both lists are empty
when the query complies."""
+
+ #: Constructs that cannot be checked against the allow-list (see
:class:`TableScan`).
+ unverifiable_sources: list[str]
+ #: Sorted, de-duplicated names of referenced tables that are not on the
allow-list.
+ disallowed_tables: list[str]
+
+
+def find_allowed_tables_violations(
+ statements: list[exp.Expr],
+ allowed: frozenset[TableKey],
+ *,
+ default_schema: str | None = None,
+ allowed_functions: frozenset[str] = frozenset(),
+) -> AllowedTablesViolations:
+ """
+ Report what in parsed ``statements`` an ``allowed_tables`` list does not
permit.
+
+ Every table the statements reach (resolved scope-correctly by
+ :func:`collect_table_references`) must match a key in ``allowed``, built
with
+ :func:`canonical_allowed_tables`; unqualified references resolve to
``default_schema``.
+ Callers turn a non-empty result into an error worded for their own context.
+
+ :param statements: Parsed statements, e.g. as returned by
:func:`validate_sql`.
+ :param allowed: Comparison keys of the allowed tables.
+ :param default_schema: Schema assumed for references that do not name one.
+ :param allowed_functions: Case-folded names of functions sqlglot cannot
type that are
+ still permitted (see :func:`collect_table_references`).
+ """
+ scan = collect_table_references(statements,
allowed_functions=allowed_functions)
+ disallowed = {
+ ".".join(part for part in (catalog, schema, table) if part)
+ for catalog, schema, table in scan.tables
+ if canonical_table_key(catalog, schema or default_schema, table) not
in allowed
Review Comment:
Added in bd4189206e: `test_query_unqualified_resolves_to_default_schema` is
now parametrized over `public.orders` and `orders`, and both cases fail without
`or default_schema`. `TestCanonicalAllowLists` checks it at the helper level as
well.
--
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]
