Lee-W commented on code in PR #74448:
URL: https://github.com/apache/airflow/pull/74448#discussion_r4227296980
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -546,15 +535,36 @@ def _query(self, sql: str) -> str:
total_rows=fetch.total_rows,
)
+ def _validate_for_execution(self, sql: str, *, dialect: str | None,
require_parse: bool) -> None:
+ """
+ Apply the checks ``query`` runs before executing ``sql``, raising when
it would be refused.
+
+ ``check_query`` shares this so it never reports valid for a statement
``query`` refuses.
+ ``require_parse`` makes ``check_query`` syntax-check writes too;
``query`` leaves a write
+ unparsed unless ``allowed_tables`` needs the AST, so a statement
sqlglot cannot parse
+ still reaches the database.
+ """
+ statements: list[Any] | None = None
+ if not self._allow_writes:
+ # allow_read_only_metadata lets agents inspect schemas with
DESCRIBE/SHOW
+ # (a common first move) instead of hard-failing; the deep scan
still
+ # rejects any data-modifying statement, including EXPLAIN <write>.
+ statements = _validate_sql(sql, dialect=dialect,
allow_read_only_metadata=True)
+ elif require_parse or self._allowed_canonical is not None:
+ # Writes are allowed, so only parse: the allow-list, when set,
still
+ # governs which tables a write may touch.
+ statements = _parse_sql(sql, dialect=dialect)
Review Comment:
`_parse_sql` now gets `allow_multiple_statements=self._allowed_canonical is
None`, so
* without an allow-list: `check_query` accepts `INSERT ...; DELETE ...` like
`query`
* with an allow-list: both still reject multiple statements
--
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]
