pavansai1997 opened a new issue, #76229:
URL: https://github.com/apache/airflow/issues/76229
### Description
Add an opt-in `allowed_tables` parameter to `LLMSQLQueryOperator` (and
`@task.llm_sql`) that rejects generated SQL referencing any table outside the
list, reusing the enforcement `SQLToolset(allowed_tables=...)` already has.
Today `table_names` only selects which schemas are put into the model's
context, and `validate_sql` only checks the statement type. Nothing checks
which tables the generated query reads. With
```python
LLMSQLQueryOperator(
task_id="report",
prompt="Total order value per customer last month",
llm_conn_id="my_llm",
table_names=["orders"],
schema_context="orders(customer_id int, amount numeric, created_at
date)",
)
```
a model that returns `SELECT name, salary FROM employee_salaries` passes
validation and the query is handed to the downstream task unchanged (verified
on `main` with a stubbed model returning that query).
The same provider already solves this for agents:
`SQLToolset(allowed_tables=[...])` resolves every table a query reaches through
`collect_table_references` (CTE scoping, subqueries, catalog qualification) and
refuses constructs it cannot verify, such as table-valued functions and unknown
functions. That check is not available to `LLMSQLQueryOperator`.
### Use case/motivation
`LLMSQLQueryOperator` turns free text into SQL that a downstream task
executes, often with a database user that can read more than the tables the
pipeline is meant for. Restricting the generated query to known tables limits
what a confused model or a prompt-injected input can reach, as defence in depth
next to least-privilege grants (the same caveat `agent_security.rst` gives for
`SQLToolset`'s allow-list). Users of `SQLToolset` already get this check; users
of the single-shot operator cannot express it.
Proposed behaviour, aligned with `SQLToolset` so both paths enforce the same
rules:
- `allowed_tables: list[str] | None = None`. `None` keeps today's behaviour,
so the change is opt-in and backward compatible.
- Matching follows `SQLToolset`: case-insensitive, an unqualified entry
matches unqualified references, a `schema.table` entry matches that qualified
reference, and catalog-qualified references are refused. An empty list raises
instead of meaning "no restriction".
- The check runs on the generated SQL and again on SQL edited by a reviewer
in the approval flow, next to the existing statement-type validation, and
raises `SQLSafetyError`.
- `allowed_tables` together with `validate_sql=False` raises at construction
rather than silently skipping the check.
Implementation-wise, the table-reference comparison would move out of
`SQLToolset` into a helper in `utils/sql_validation.py` used by both, so there
is one implementation of the rules. `SQLToolset` behaviour would stay unchanged.
### Related issues
None found.
### Are you willing to submit a PR?
- [X] Yes I am willing to submit a PR!
### Code of Conduct
- [X] I agree to follow this project's Code of Conduct
--
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]
