Lee-W commented on code in PR #70134:
URL: https://github.com/apache/airflow/pull/70134#discussion_r3615315006
##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py:
##########
@@ -65,6 +65,13 @@ def example_agent_operator_sql():
"the schema and answer the question with data."
),
toolsets=[
+ # ``allowed_tables`` scopes the agent's intent, but it is an
+ # application-level guardrail, not a security boundary. Point
+ # ``postgres_default`` at a least-privilege role whose SELECT
grants
+ # are limited to these tables -- that is the boundary that
holds even
+ # if the agent (which may be under prompt injection) reaches
for data
+ # through a function the parser cannot see. See the "Security"
section
+ # of the toolsets docs.
SQLToolset(
db_conn_id="postgres_default",
Review Comment:
do we want to also add `allow_functions` as part of the example
##########
providers/common/ai/src/airflow/providers/common/ai/utils/sql_validation.py:
##########
@@ -289,6 +311,38 @@ def collect_table_references(statements: list[exp.Expr])
-> TableScan:
if isinstance(stmt, (exp.Command, exp.Execute)):
unverifiable.append(f"a {type(stmt).__name__.lower()} statement")
continue
+ # COPY moves data between a table and the server filesystem or a
spawned program
+ # (``COPY t FROM/TO PROGRAM '...'``, ``COPY t FROM/TO '<file>'``). The
table it
+ # names is real and allow-listed, but the data channel -- a file or a
program --
+ # is not a table the allow-list can describe, and ``FROM PROGRAM`` is
arbitrary
+ # command execution. Top-level COPY is already blocked in read-only
mode by the
+ # statement-type allow-list; this also refuses it under
``allow_writes`` while an
+ # allow-list is active, where only this scan runs.
+ if isinstance(stmt, exp.Copy):
+ unverifiable.append("a COPY statement")
+ continue
+ # A function whose string argument reaches a file, another table, or a
program
+ # (``pg_read_file``, ``query_to_xml``, scalar ``dblink``, ...) carries
no
+ # ``exp.Table`` node, so the table scan below cannot see it. sqlglot
parses any
+ # function it cannot type as ``exp.Anonymous`` (typed builtins like
``count`` are
+ # ``exp.Func`` subclasses), so reject every ``exp.Anonymous`` not
explicitly
+ # allow-listed rather than chase an unbounded denylist of dangerous
names. A
+ # schema-qualified call (``pg_catalog.pg_read_file(...)``) parses as
+ # ``Dot(this=..., expression=Anonymous)`` with the bare name on the
nested
+ # ``Anonymous``, so ``find_all`` still reaches it.
+ unknown = sorted(
+ {
+ name
+ for fn in stmt.find_all(exp.Anonymous)
+ if (name := fn.name.casefold()) not in allowed_functions
+ }
+ )
+ if unknown:
+ unverifiable.append(
+ f"function(s) the parser cannot verify against allowed_tables "
+ f"({', '.join(unknown)}); if safe, permit them via
allowed_functions"
Review Comment:
```suggestion
f"({', '.join(unknown)}); if these functions are trusted,
permit them via allowed_functions"
```
--
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]
