This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9100ca8a4e5 Reject allowed_tables=None in SQLToolset so only omitting
it allows every table (#73452)
9100ca8a4e5 is described below
commit 9100ca8a4e514bd21a2ffe63b7a3e22b0c0ecca3
Author: Kaxil Naik <[email protected]>
AuthorDate: Mon Sep 21 13:58:02 2026 +0100
Reject allowed_tables=None in SQLToolset so only omitting it allows every
table (#73452)
---
providers/common/ai/docs/changelog.rst | 14 ++++++----
providers/common/ai/docs/choosing_a_toolset.rst | 11 ++++----
providers/common/ai/docs/toolsets.rst | 8 ++++--
.../airflow/providers/common/ai/toolsets/sql.py | 32 ++++++++++++++++------
.../ai/tests/unit/common/ai/toolsets/test_sql.py | 12 ++++++--
5 files changed, 52 insertions(+), 25 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index dd5daeb3047..5fc707ae7c1 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -37,12 +37,14 @@ Changelog
:doc:`retry_policies`, "When the connection also carries a fallback chain".
.. note::
- ``SQLToolset(allowed_tables=[])`` now raises ``ValueError``. Up to 0.9.0 an
empty list
- was accepted and exposed every table in the schema -- the same as
``allowed_tables=None``
- -- so a Dag that builds the list dynamically (a ``Variable.get``, a config
file, a
- filtered comprehension) silently handed the agent the whole schema whenever
the list
- came back empty. Such a Dag now fails at import instead. Pass ``None``
explicitly if
- exposing every table is what you meant.
+ ``SQLToolset`` now rejects ``allowed_tables=None`` and ``allowed_tables=[]``
with
+ ``ValueError``. Up to 0.9.0 both were accepted and exposed every table in
the schema, so
+ a Dag that builds the list dynamically (a ``Variable.get`` with a ``None``
default, a
+ config file, a filtered comprehension) silently handed the agent the whole
schema
+ whenever the lookup came back empty. Such a Dag now fails at import instead.
Exposing
+ every table is still the default, but only by omitting the argument: no
value you can
+ pass requests it, so a runtime lookup can never widen the allow-list by
accident. Dags
+ that passed ``allowed_tables=None`` explicitly should drop the argument.
0.9.0
.....
diff --git a/providers/common/ai/docs/choosing_a_toolset.rst
b/providers/common/ai/docs/choosing_a_toolset.rst
index 3ac32937e33..cc7f8a54a2a 100644
--- a/providers/common/ai/docs/choosing_a_toolset.rst
+++ b/providers/common/ai/docs/choosing_a_toolset.rst
@@ -144,11 +144,12 @@ subqueries and joins.
something other than a DBAPI cursor — ``ExasolHook`` and its pyexasol
statement, for instance — fall back to a full fetch. The payload handed to
the
model is still bounded; the transfer is not. See
:ref:`bounded-query-results`.
-- Its parser-level closure is opt-in, not the default. ``allowed_tables``
- defaults to ``None`` and the table walk returns immediately while it is
- unset, so out of the box the agent reaches every table the connection can
- see. ``DESCRIBE`` and ``SHOW`` both pass, on dialects that parse them,
- while ``allowed_tables`` stays unset. Set ``allowed_tables`` and the walk
+- Its parser-level closure is opt-in, not the default. The table walk returns
+ immediately while ``allowed_tables`` is unset, so out of the box the agent
+ reaches every table the connection can see. Only omitting the argument grants
+ that: an explicit ``None`` or empty list is rejected at construction.
+ ``DESCRIBE`` and ``SHOW`` both pass, on dialects that parse them, while
+ ``allowed_tables`` stays unset. Set ``allowed_tables`` and the walk
turns fail-closed for ``SHOW`` — but not for ``DESCRIBE``: it instead
becomes an ordinary table reference, allowed only when the table it names
is on the list. See :ref:`allowed-tables-enforcement` for what else the
diff --git a/providers/common/ai/docs/toolsets.rst
b/providers/common/ai/docs/toolsets.rst
index 3461935fe62..63cbfa01399 100644
--- a/providers/common/ai/docs/toolsets.rst
+++ b/providers/common/ai/docs/toolsets.rst
@@ -206,9 +206,11 @@ Parameters
^^^^^^^^^^
- ``db_conn_id``: Airflow connection ID for the database.
-- ``allowed_tables``: Restrict the agent to a fixed set of tables. ``None``
- (default) exposes all tables in ``schema``; an empty list raises
``ValueError``
- rather than silently exposing them all. Entries may be schema-qualified
+- ``allowed_tables``: Restrict the agent to a fixed set of tables. Omit the
+ argument (the default) to expose all tables in ``schema``. No value means
+ allow-all: ``None`` and an empty list both raise ``ValueError``, so an
allow-list
+ built at runtime that resolves to nothing fails at import instead of silently
+ exposing every table. Entries may be schema-qualified
(``"SCHEMA.TABLE"``) to span multiple schemas; see above. Matching is
case-insensitive. When set, the list is enforced on ``query`` and
``check_query`` as well as discovery -- every table a query references must
be
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
index 5011ba99c20..bbfd32108b6 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py
@@ -52,6 +52,13 @@ from airflow.providers.common.compat.sdk import BaseHook
if TYPE_CHECKING:
from pydantic_ai._run_context import RunContext
+# Sentinel distinguishing "caller did not pass ``allowed_tables``" (expose
every
+# table) from an explicit value. Every explicit falsy value -- ``None`` as
much as
+# ``[]`` -- is rejected, so an allow-list a Dag builds at runtime can never
widen to
+# allow-all by resolving to nothing. Allow-all is reachable only by omitting
the
+# argument, which is a static, visible decision in the Dag file.
+_UNSET: Any = object()
+
# JSON Schemas for the four SQL tools.
_LIST_TABLES_SCHEMA: dict[str, Any] = {
"type": "object",
@@ -170,9 +177,12 @@ class SQLToolset(AbstractToolset[Any]):
toolset does not inspect the error type or message.
:param db_conn_id: Airflow connection ID for the database.
- :param allowed_tables: Restrict the agent to a fixed set of tables.
``None``
- (default) exposes every table in ``schema``; an empty list raises
- ``ValueError`` rather than silently exposing every table. Entries may
be
+ :param allowed_tables: Restrict the agent to a fixed set of tables. Omit
the
+ argument (the default) to expose every table in ``schema``. No *value*
means
+ allow-all: ``None`` and an empty list both raise ``ValueError``, so an
allow-list
+ built at runtime (a ``Variable.get`` with a ``None`` default, a config
lookup, a
+ filtered comprehension) that resolves to nothing fails at import
instead of
+ silently handing the agent the whole schema. Entries may be
schema-qualified (``"SCHEMA.TABLE"``) to span multiple schemas in one
database
-- common on warehouses such as Snowflake. ``list_tables`` introspects
each referenced
schema and returns the matching tables fully qualified, and
``get_schema``
@@ -240,20 +250,26 @@ class SQLToolset(AbstractToolset[Any]):
self,
db_conn_id: str,
*,
- allowed_tables: list[str] | None = None,
+ allowed_tables: list[str] = _UNSET,
allowed_functions: list[str] | None = None,
schema: str | None = None,
allow_writes: bool = False,
max_rows: int = 50,
max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES,
) -> None:
- if allowed_tables is not None and not allowed_tables:
+ self._allowed_tables: frozenset[str] | None
+ if allowed_tables is _UNSET:
+ self._allowed_tables = None
+ elif not allowed_tables:
raise ValueError(
- "allowed_tables must not be empty. Pass None to allow every
table in the schema, "
- "or list the tables the agent may access."
+ f"allowed_tables must name at least one table, got
{allowed_tables!r}. Omit the "
+ "argument to expose every table in the schema. An empty or
missing value is "
+ "rejected rather than read as 'no restriction' so that an
allow-list built at "
+ "runtime cannot silently widen to every table."
)
+ else:
+ self._allowed_tables = frozenset(allowed_tables)
self._db_conn_id = db_conn_id
- self._allowed_tables: frozenset[str] | None =
frozenset(allowed_tables) if allowed_tables else None
# Case-folded so matching a query's function names (also case-folded)
is
# case-insensitive, mirroring how allowed_tables is compared.
self._allowed_functions: frozenset[str] = (
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
index 455805d1043..57e0db38b0d 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py
@@ -96,9 +96,15 @@ class TestSQLToolsetInit:
ts = SQLToolset("my_pg")
assert ts.id == "sql-my_pg"
- def test_empty_allowed_tables_raises(self):
- with pytest.raises(ValueError, match="allowed_tables must not be
empty"):
- SQLToolset("my_pg", allowed_tables=[])
+ @pytest.mark.parametrize("value", [None, []], ids=["none", "empty-list"])
+ def test_falsy_allowed_tables_raises(self, value):
+ """Only omitting the argument grants allow-all; no explicit value
does."""
+ with pytest.raises(ValueError, match="allowed_tables must name at
least one table"):
+ SQLToolset("my_pg", allowed_tables=value)
+
+ def test_omitted_allowed_tables_means_no_restriction(self):
+ ts = SQLToolset("my_pg")
+ assert ts._allowed_tables is None
class TestSQLToolsetGetTools: