kaxil commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4504571325
##########
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:
`None` still means no restriction, so a Dag that builds the list at parse
time, like `allowed_tables=Variable.get("llm_sql_tables", default=None,
deserialize_json=True)`, runs with no allow-list when the Variable is missing.
That's the case the message below says it prevents, and `SQLToolset` uses an
`_UNSET` sentinel to reject an explicit `None` for exactly this. Could the
operator do the same?
##########
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:
Nothing checks the shape of an entry. `allowed_tables="orders"` is iterated
by character, so the prompt lists `o, r, d, e, r, s` and `SELECT * FROM o`
passes. `"db.sales.orders"` becomes schema `db.sales`, so it can never match,
and the refusal reads `not in allowed_tables: db.sales.orders. Allowed tables:
db.sales.orders.` Since both callers go through here now, could it reject a
`str` and entries with more than one dot?
##########
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:
With neither `dialect` nor `db_conn_id` set, as in the `schema_context`-only
examples on the docs page, this parses with no dialect, so the check can read
the SQL differently from the database. `SELECT 'a\' AS x, ' , (SELECT pw FROM
secret) AS b -- '` followed by a `FROM orders` line passes against
`["orders"]`, and MySQL, where the backslash escapes the quote, returns
`secret.pw`. No dialect also refuses builtins like `now()` and `date_part()`
that it types under `postgres`. Should `allowed_tables` require a resolved
dialect?
##########
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:
On SQLite, `SELECT * FROM orders WHERE id IN secret` passes against
`["orders"]`: sqlglot parses `secret` as a column, but SQLite reads `IN <name>`
as a table. `SQLToolset` already has this, so a follow-up is fine, but should
an `IN` with a bare name on the right count as unverifiable?
##########
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:
Nothing pins `or default_schema`: the suite stays green without it, and then
`SQLToolset(allowed_tables=["orders"], schema="public")` refuses a plain
`SELECT * FROM orders`. `test_query_unqualified_resolves_to_default_schema`
only queries `public.orders`. Could you add the unqualified case?
--
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]
