pavansai1997 commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4500661413
##########
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:
Done in 17fe1a3e7e. When pydantic-ai gives up, `execute` now re-raises the
`SQLSafetyError` from the cause chain, and lets any other
`UnexpectedModelBehavior` through unchanged.
`test_model_that_never_complies_raises_sql_safety_error` covers both a
disallowed table and `DROP TABLE orders`, and
`test_unrelated_model_failure_is_raised_unchanged` pins the pass-through.
##########
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:
Added `test_fenced_allowed_answer_is_accepted_on_first_call`. It fails if
the validator checks the raw output instead of the stripped SQL.
##########
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` now catches `SqlglotError`, so `I can't write that query` goes
back to the model as a parse error. `test_rejects_untokenizable` covers it.
##########
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:
Fixed here rather than in a follow-up, since it applies to `SQLToolset` too.
Anything under `LATERAL`/`APPLY` is now unverifiable unless it is a subquery or
a column generator (`UNNEST`, `EXPLODE`/`FLATTEN`, `INLINE`, `OPENJSON`), whose
arguments are still scanned. Checking only for a typed function would have
missed allow-listed and schema-qualified functions under `LATERAL` (`LATERAL
dblink(...)` with `allowed_functions=["dblink"]`, `LATERAL
pg_catalog.fn(...)`), so those are refused too and covered by tests.
--
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]
