pavansai1997 commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4506791115
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -133,12 +167,25 @@ def __init__(
raise TypeError(
"LLMSQLQueryOperator does not accept 'output_type'; it always
returns the generated SQL as str."
)
+ if allowed_tables is not None:
Review Comment:
Done in bd4189206e. The operator now uses the same `_UNSET` sentinel as
`SQLToolset`, so an explicit `None` is rejected like `[]`, and omitting the
argument is the only way to allow every table. The `none` case in
`test_invalid_configuration_is_rejected_at_construction` covers the
missing-Variable case.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -193,10 +250,39 @@ def execute_complete(
"""Resume after human review, re-validating if the reviewer modified
the SQL."""
output = self._complete_review(context, generated_output, event,
decision)
if output is not None and output != generated_output:
- _validate_sql(output, allowed_types=self.allowed_sql_types,
dialect=self._resolved_dialect)
+ self._check_sql(output)
self._clear_usage_budget(context)
return output
+ def _retry_unsafe_sql(self, output: str) -> str:
+ """Send a refusal back to the model so ``output_retries`` lets it
rewrite the query."""
+ try:
+ self._check_sql(self._strip_llm_output(output,
dialect=self._resolved_dialect))
+ except SQLSafetyError as e:
+ raise ModelRetry(str(e)) from e
+ return output
+
+ def _check_sql(self, sql: str) -> None:
+ """Validate the statement types of ``sql`` and, when set, the tables
it may reference."""
+ statements = _validate_sql(sql, allowed_types=self.allowed_sql_types,
dialect=self._resolved_dialect)
Review Comment:
Yes, done in bd4189206e. With `allowed_tables` set, construction fails when
neither `dialect` nor `db_conn_id` is set, and `_check_sql` refuses to run when
no dialect resolves (an unknown `dialect` or hook `dialect_name`). That also
covers reviewer-edited SQL when Human-in-the-Loop review resumes, and `execute`
runs the same check before the model is called.
`test_dialect_decides_where_a_backslash_string_ends` pins your example: `mysql`
parses the subquery on `secret` and no dialect does not, and the operator
refuses the query under `mysql`.
`SQLToolset` can still parse with no dialect when the hook reports `default`
or an unknown name. Since that is a released toolset, I'd rather fix it in a
follow-up, unless you want it here.
##########
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(".")
Review Comment:
Done in bd4189206e. `canonical_allowed_tables` now raises `TypeError` for a
string or a non-string entry, and `ValueError` for anything other than `table`
or `schema.table`, including empty parts. `SQLToolset` now passes the argument
as given rather than the `frozenset` it builds first, which had already split a
string into characters, so the check applies to both callers.
`allowed_functions` had the same problem, so both now share a
`canonical_allowed_functions` helper that rejects a string too.
##########
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)
Review Comment:
Fixed here rather than in a follow-up. Anything in the `field` argument of
`exp.In` is now unverifiable: that is where sqlglot puts `IN secret`, `IN
main.secret` and `IN json_each(...)`, while lists, subqueries and `UNNEST` use
other arguments, so `IN (1, 2)` and `IN (SELECT ...)` are checked as before. It
applies to `SQLToolset` too, and the tests cover both sides.
--
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]
