pavansai1997 commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4436232040
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -80,13 +83,21 @@ class LLMSQLQueryOperator(LLMOperator):
:param db_conn_id: Connection ID for database schema introspection.
The connection must resolve to a ``DbApiHook``.
:param table_names: Tables to include in the LLM's schema context.
- Used with ``db_conn_id`` for automatic introspection.
+ Used with ``db_conn_id`` for automatic introspection. This does not
restrict
+ the tables the generated SQL may reference; use ``allowed_tables`` for
that.
:param schema_context: Manual schema context string. When provided,
this is used instead of ``db_conn_id`` introspection.
:param validate_sql: Whether to validate generated SQL via AST parsing.
Default ``True`` (safe by default).
:param allowed_sql_types: SQL statement types to allow.
Default: ``(Select, Union, Intersect, Except)``.
+ :param allowed_tables: Tables the generated SQL may reference, as
``table`` or
Review Comment:
Good point — a schema is a natural unit to allow-list (e.g. a reporting
schema whose tables change over
time). Since the matching now lives in one shared helper, I'd suggest
supporting schema.* entries in the
same allowed_tables list, e.g. allowed_tables=["analytics.*",
"public.dim_date"], rather than adding a
separate parameter. It would apply to both LLMSQLQueryOperator and
SQLToolset so their rules stay
identical. To stay fail-closed, an unqualified reference wouldn't match a
schema entry unless a default
schema is known (SQLToolset(schema=...)), since the operator can't see the
database's search path, and a
bare * would be rejected in favour of None. Happy to add this to this PR if
that works for you.
---
Drafted-by: Code Puppy; reviewed by @pavansai1997 before posting
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py:
##########
@@ -759,6 +759,103 @@ def _make_context(ti_id=None):
@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with
Airflow >= 3.1.0"
)
+class TestLLMSQLQueryOperatorAllowedTables:
Review Comment:
Good catch, my mistake. Moved the class above the decorator so the approval
suite keeps its 3.1 skip, and `test_reviewer_modified_sql_is_checked` now has
its own skip, in 80e91e559b.
---
Drafted-by: Code Puppy (no human review before posting)
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -192,10 +216,27 @@ 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 _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)
+ if self.allowed_tables is None:
+ return
+ violations = find_allowed_tables_violations(statements,
canonical_allowed_tables(self.allowed_tables))
Review Comment:
Added `allowed_functions` with the same semantics as `SQLToolset`
(case-insensitive, only consulted when `allowed_tables` is set), so the
existing message is now accurate.
`test_allowed_functions_permits_unrecognized_functions` covers your
`json_build_object` and `age` examples, and `age` without it is pinned as
refused.
---
Drafted-by: Code Puppy (no human review before posting)
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -172,7 +196,7 @@ def execute(self, context: Context) -> str:
sql = self._strip_llm_output(result.output,
dialect=self._resolved_dialect)
if self.validate_sql:
- _validate_sql(sql, allowed_types=self.allowed_sql_types,
dialect=self._resolved_dialect)
+ self._check_sql(sql)
Review Comment:
Done both. With `allowed_tables` set, the prompt lists the tables and asks
for unquoted table names and no comments, and the check runs as an output
validator that raises `ModelRetry`, so `output_retries` covers it. The post-run
check stays as a backstop and the reviewer path still raises `SQLSafetyError`.
`test_refused_sql_is_sent_back_to_the_model` drives a real `Agent` with a
`FunctionModel` through a refusal and a successful retry.
---
Drafted-by: Code Puppy (no human review before posting)
--
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]
