diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index dd5daeb30472e..5fc707ae7c1d7 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 3ac32937e331b..cc7f8a54a2af7 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 3461935fe62f7..63cbfa0139972 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 5011ba99c2095..bbfd32108b661 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 @@ 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 @@ def __init__( 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 455805d104364..57e0db38b0d93 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 @@ def test_id_includes_conn_id(self): 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: