kaxil commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4467723006
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -192,10 +231,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 != 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
Review Comment:
Once `output_retries` runs out, this surfaces as pydantic-ai's
`UnexpectedModelBehavior: Exceeded maximum output retries`, with the
`SQLSafetyError` only in the cause chain, so the documented `SQLSafetyError`
(and any retry policy or callback keyed on it) no longer applies. It covers the
statement-type check too: a generated `DROP TABLE orders` fails as
`SQLSafetyError` without `allowed_tables` and as `UnexpectedModelBehavior` with
it. Could `execute` re-raise the `SQLSafetyError` from the cause chain, with a
`FunctionModel` test where the model never complies?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -192,10 +231,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 != 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:
Review Comment:
`parse_sql` only converts `ParseError`, so an answer with an unbalanced
quote, such as a plain-text `I can't write that query`, raises sqlglot's
`TokenError` here and fails the task instead of going back to the model. Could
`parse_sql` catch `SqlglotError` instead?
##########
providers/common/ai/docs/operators/llm_sql.rst:
##########
@@ -178,6 +178,49 @@ By default, the operator validates generated SQL using an
allowlist approach:
You can disable validation with ``validate_sql=False`` or customize the allowed
statement types with ``allowed_sql_types``.
+Restricting the tables a query may reference
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``table_names`` only selects the schemas shown to the LLM; it does not stop the
+generated SQL from reading other tables. Set ``allowed_tables`` to reject any
query
+that references a table outside the list:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
+ :language: python
+ :start-after: [START howto_operator_llm_sql_allowed_tables]
+ :end-before: [END howto_operator_llm_sql_allowed_tables]
+
+Every table the query reaches is checked, including tables in subqueries and
joins,
+while a CTE defined in the query is not mistaken for a table. Entries are
``table`` or
+``schema.table`` and match case-insensitively. The operator has no default
schema, so
+an unqualified entry only matches unqualified references: ``orders`` does not
match
+``sales.orders``. Catalog-qualified references never match.
+
+SQL that cannot be checked against the list is refused as well. In practice
that
+mostly means:
+
+- SQL comments, such as a trailing ``-- per customer``;
+- quoted identifiers, such as ``FROM "orders"``, because a quoted name is
+ case-sensitive on the database while the list is matched case-insensitively;
+- functions sqlglot does not recognize, such as Postgres' ``age`` or
+ ``json_build_object``. List trusted ones in ``allowed_functions``;
+- table-valued functions such as ``dblink``.
Review Comment:
On DuckDB, `SELECT * FROM LATERAL read_parquet('/tmp/private.parquet')` gets
through: sqlglot wraps the typed `ReadParquet` in a `Lateral`, so the scanner
finds no table and nothing it can't check, while the same query without
`LATERAL` is refused. `collect_table_references` already has this gap for
`SQLToolset`, so a follow-up is fine, but should a `Lateral` over a function
count as unverifiable?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -192,10 +231,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 != 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))
Review Comment:
Nothing pins the `_strip_llm_output` call here: with
`self._check_sql(output)` instead the suite stays green, but every fenced
answer would use up its retries and fail. Could the `FunctionModel` test also
cover a fenced, allowed answer that's accepted on the first call?
--
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]
