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.
--
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]
