pavansai1997 opened a new pull request, #76230: URL: https://github.com/apache/airflow/pull/76230
`LLMSQLQueryOperator` had no way to restrict the tables generated SQL may read: `table_names` only selects the schemas shown to the LLM and `validate_sql` only checks the statement type. With `table_names=["orders"]`, a model answer such as `SELECT name, salary FROM employee_salaries` was returned unchanged. This adds an opt-in `allowed_tables` parameter (also available through `@task.llm_sql`). When set, the operator rejects SQL that references a table outside the list, or uses a construct that cannot be checked against it (for example `dblink`), with `SQLSafetyError`. The check covers subqueries and joins, does not mistake a CTE for a table, and applies both to the generated SQL and to SQL a reviewer edits during approval. The rules are the ones `SQLToolset(allowed_tables=...)` already enforces. Instead of a second implementation, the table comparison moves from `SQLToolset` into helpers in `utils/sql_validation.py` (`canonical_allowed_tables`, `find_allowed_tables_violations`) that both use. Each caller keeps its own error wording, so `SQLToolset`'s behaviour and messages are unchanged. `allowed_tables` defaults to `None`, so existing Dags behave as before. An empty list and `validate_sql=False` are rejected at construction, so an allow-list can never be silently skipped. The docs describe the check as defence in depth next to least-privilege database grants, as `agent_security.rst` does for `SQLToolset`. closes: #76229 --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes — Code Puppy Generated-by: Code Puppy following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) -- 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]
