kaxil commented on code in PR #76230:
URL: https://github.com/apache/airflow/pull/76230#discussion_r4437632821
##########
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:
No `allowed_functions` is passed here, so every function sqlglot can't type
is refused. On Postgres with `allowed_tables=["orders"]`, both `SELECT
json_build_object('id', id) FROM orders` and `SELECT age(created_at) FROM
orders` fail, and the error tells the user to "permit them via
allowed_functions", which this operator doesn't accept (passing it fails as an
invalid argument). Should the operator take `allowed_functions` like
`SQLToolset` does? If not, this path's message shouldn't point at it.
##########
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:
With `allowed_tables` set, ordinary model output fails the task here: a
trailing `-- per customer` is refused as an inline comment and `FROM "orders"`
as a quoted identifier, though both only read `orders`. `SQLToolset` hands that
refusal back to the model as a tool error so it can rewrite the query, but this
check runs after the only model call, and `_build_system_prompt` never mentions
the allowed tables or these rules. Could the prompt list the allowed tables and
ask for no comments and unquoted table names when `allowed_tables` is set?
Running the check as an output validator that raises `ModelRetry` would also
let `output_retries` cover it.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -124,6 +135,7 @@ def __init__(
schema_context: str | None = None,
validate_sql: bool = True,
allowed_sql_types: tuple[type[exp.Expr], ...] = DEFAULT_ALLOWED_TYPES,
+ allowed_tables: list[str] | None = None,
Review Comment:
`table_names` is a template field but `allowed_tables` isn't, so a Dag that
templates `table_names` can't keep the allow-list in step with it. Is a static
list intended here since it's a safety control, or should it join
`template_fields`?
##########
providers/common/ai/docs/operators/llm_sql.rst:
##########
@@ -178,6 +178,38 @@ 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:
+
+.. code-block:: python
+
+ LLMSQLQueryOperator(
+ task_id="generate_sql",
+ prompt="Total order value per customer last month",
+ llm_conn_id="pydanticai_default",
+ db_conn_id="postgres_default",
+ table_names=["orders"],
+ allowed_tables=["orders"],
+ )
+
+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. Constructs that
cannot
+be checked against the list, such as table-valued functions like ``dblink``,
are
+rejected too. The rules match ``SQLToolset(allowed_tables=...)``: entries are
Review Comment:
These aren't quite the same rules. `SQLToolset` resolves unqualified entries
through its `schema`, so with `schema="public"` the entry `orders` matches
`public.orders`, while the operator refuses it. `SQLToolset` also has
`allowed_functions` and the operator doesn't. And the refusals users will
actually hit are comments, quoted table names and functions sqlglot doesn't
recognise, rather than `dblink`, so it'd help to list those here.
##########
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:
This class landed between the existing `skipif(not AIRFLOW_V_3_1_PLUS)` and
`TestLLMSQLQueryOperatorApproval`, so the decorator now applies here and the
approval suite lost it. The compat matrix runs common.ai tests on 2.11.1 and
3.0.6, where the approval tests hit `require_approval=True needs Airflow 3.1+`
and patch `airflow.sdk.execution_time.hitl`. Could you move this class above
the decorator, and give `test_reviewer_modified_sql_is_checked` its own 3.1
skip since it's the only one going through the review path?
##########
providers/common/ai/docs/operators/llm_sql.rst:
##########
@@ -178,6 +178,38 @@ 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:
+
+.. code-block:: python
Review Comment:
The other examples on this page come from `example_llm_sql.py` via
`exampleinclude`, so the Dag import tests parse them. Could this one live there
too?
##########
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:
+ @staticmethod
+ def _execute(mock_hook_cls, make_mock_run_result, sql, allowed_tables,
**kwargs):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ sql, make_mock_run_result
+ )
+ op = LLMSQLQueryOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ schema_context="Table: orders",
+ allowed_tables=allowed_tables,
+ **kwargs,
+ )
+ return op.execute(context=MagicMock())
+
+ @pytest.mark.parametrize(
+ ("sql", "allowed_tables"),
+ [
+ pytest.param("SELECT customer_id, SUM(amount) FROM orders GROUP BY
1", ["orders"], id="allowed"),
+ pytest.param("SELECT * FROM ORDERS", ["orders"],
id="case-insensitive"),
+ pytest.param("SELECT * FROM sales.orders", ["sales.orders"],
id="schema-qualified"),
+ pytest.param(
+ "WITH salaries AS (SELECT * FROM orders) SELECT * FROM
salaries",
+ ["orders"],
+ id="cte-named-like-another-table",
+ ),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_sql_within_allowed_tables_is_returned(
+ self, mock_hook_cls, make_mock_run_result, sql, allowed_tables
+ ):
+ assert self._execute(mock_hook_cls, make_mock_run_result, sql,
allowed_tables) == sql
+
+ @pytest.mark.parametrize(
+ ("sql", "disallowed"),
+ [
+ pytest.param("SELECT name, salary FROM salaries", "salaries",
id="other-table"),
+ pytest.param(
+ "SELECT * FROM orders WHERE customer_id IN (SELECT id FROM
customers)",
+ "customers",
+ id="subquery",
+ ),
+ pytest.param(
+ "SELECT * FROM orders o JOIN salaries s ON o.customer_id =
s.id", "salaries", id="join"
+ ),
+ pytest.param("SELECT * FROM otherdb.public.orders",
"otherdb.public.orders", id="catalog"),
Review Comment:
Could you add `SELECT * FROM sales.orders` against `["orders"]` here? With
no default schema that's refused, which is the one place the operator's rules
differ from `SQLToolset`'s, and nothing pins it.
--
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]
