This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch sql-toolset-infra-errors in repository https://gitbox.apache.org/repos/asf/airflow.git
commit fcafefb5bc92cb99336a652e8ee654c76e12284d Author: Kaxil Naik <[email protected]> AuthorDate: Mon Nov 9 18:49:12 2026 +0000 Fail the attempt when SQLToolset cannot reach its database SQLToolset handed every exception back to the model as a ModelRetry, including an unreachable host, rejected credentials and a missing connection. No SQL change fixes those: the model either answered without data, so the task succeeded with an apology as its output, or spent its retries. Under durable execution the recorded ModelRetry also replayed the outage on retry instead of querying again. Failing to reach the database now raises: resolving the connection's hook, building the hook's URI (which can fetch credentials), the hook's get_conn (which DbApiHook.run uses to connect), building the SQLAlchemy inspector for list_tables and get_schema, and a SQLAlchemy DBAPIError without a statement or with connection_invalidated. A raised call is not recorded, so a durable retry runs it again. An error from the statement itself stays a ModelRetry, so the model can still correct its SQL and the correction replays. Inside CodeMode the harness still hands the error back to the model; the docs say so. --- providers/common/ai/docs/code_mode.rst | 3 + providers/common/ai/docs/durable_execution.rst | 4 +- providers/common/ai/docs/toolsets/index.rst | 11 +- providers/common/ai/docs/toolsets/sql.rst | 33 +++-- providers/common/ai/docs/troubleshooting.rst | 10 ++ .../airflow/providers/common/ai/toolsets/sql.py | 101 ++++++++++--- .../unit/common/ai/durable/test_capability.py | 61 ++++++++ .../ai/tests/unit/common/ai/sqlite_connection.py | 85 +++++++++++ .../unit/common/ai/tools/test__from_toolset.py | 18 ++- .../ai/tests/unit/common/ai/toolsets/test_sql.py | 165 ++++++++++++++++++--- 10 files changed, 431 insertions(+), 60 deletions(-) diff --git a/providers/common/ai/docs/code_mode.rst b/providers/common/ai/docs/code_mode.rst index b4a5bb7b3b8..242c6917103 100644 --- a/providers/common/ai/docs/code_mode.rst +++ b/providers/common/ai/docs/code_mode.rst @@ -98,6 +98,9 @@ Code mode turns off :doc:`tool_approval`. A tool that needs approval and is called from ``run_code`` does not run: the model gets an error back and may carry on without it, so the task can still succeed. To keep such a tool out of ``run_code``, list the other tools in ``CodeMode(tools=[...])``; called directly, it fails the task. +The same goes for any exception a tool raises from inside ``run_code``: the model gets it +back as an error to work around, so a ``SQLToolset`` call that cannot reach its database +does not fail the task as it does when called directly. Importing ``CodeMode`` in the Dag file loads pydantic-ai-harness and Monty each time the file is parsed. Keep code mode agents in their own Dag file if the rest of the file does not need them. diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index 87bffbbed7a..f46bd62b96c 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/providers/common/ai/docs/durable_execution.rst @@ -187,7 +187,9 @@ A tool that raises ``ModelRetry``, ``ToolFailed``, ``ApprovalRequired`` or ``CallDeferred`` replays as the same signal. ``AgentOperator`` does not support tool approval or deferred tools with ``durable=True``, so there such a call fails the task with ``UnsupportedToolDeferralError``, as it would without ``durable``. A tool that raises -any other exception runs again on retry. +any other exception runs again on retry. ``SQLToolset`` uses that split: a statement the +database rejects is a ``ModelRetry`` and replays, while failing to reach the database +raises, so the retry queries it again rather than handing the model the old failure. A task can run more than one agent, and an agent's tool can run another agent; each agent run is replayed on its own, and one a tool started is replayed under that tool diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst index 230b6e9c0fd..8752188077a 100644 --- a/providers/common/ai/docs/toolsets/index.rst +++ b/providers/common/ai/docs/toolsets/index.rst @@ -267,8 +267,15 @@ When the model calls a tool with arguments that fail its schema, or the tool ask model to try again (``ModelRetry``), the error goes back to the model so it can correct the call. What counts differs by toolset: -- ``SQLToolset`` and ``DataFusionToolset`` turn every query error into ``ModelRetry``, - so a misspelled column and a dropped connection both count. +- ``SQLToolset`` counts an error the database returns for a statement, such as a + misspelled column, and a statement it refuses itself (a write while + ``allow_writes`` is off, a table outside ``allowed_tables``). Failing to reach the + database fails the run straight away, except through a hook that only connects when + the statement runs; see :doc:`sql`. +- ``DataFusionToolset`` counts a query naming a column or table that does not exist and, + while ``allow_writes`` is off, a statement it refuses itself (a write, or SQL it cannot + parse). Other errors from running a query go back to the model as a failed result + without using the budget. - ``HookToolset`` counts invalid arguments and a call that supplies a pinned argument. An exception from the hook itself fails the run straight away. - ``ObjectStorageToolset`` counts invalid arguments only. A path that does not exist or diff --git a/providers/common/ai/docs/toolsets/sql.rst b/providers/common/ai/docs/toolsets/sql.rst index 1e44b4b87ad..453ccc62cbe 100644 --- a/providers/common/ai/docs/toolsets/sql.rst +++ b/providers/common/ai/docs/toolsets/sql.rst @@ -450,15 +450,30 @@ subqueries and joins. walk rejects once ``allowed_tables`` is active. Statements that modify data, ``COPY`` among them, are rejected either way while ``allow_writes`` is ``False``. -- It does not classify failures. A connection error or a typo in a column - name reaching ``list_tables``, ``get_schema`` or ``query`` becomes one - ``ModelRetry``, so the two are treated the same way until the retry budget - (:ref:`toolset-retry-budget`) runs out and the task fails for Airflow to retry. Two paths do not raise: - ``check_query`` catches its own errors and reports them back as a normal - ``{"valid": false, ...}`` result, and ``get_schema`` returns a normal - ``{"error": ...}`` result instead of raising when the requested table is - outside ``allowed_tables``; other ``get_schema`` failures still raise and - still become a ``ModelRetry``. +- It tells an unreachable database from a bad statement by when the error + happens, not by what it says. Failing to reach the database fails the + attempt straight away for Airflow to retry: the connection is not defined, + connecting fails (an unreachable host, rejected credentials), SQLAlchemy + cannot be set up for ``list_tables`` or ``get_schema`` (a dialect it cannot + load, a token it cannot fetch), or SQLAlchemy reports the connection lost + while they run. With :doc:`../durable_execution` the retry queries the + database again instead of replaying the failure. The task log reads + ``Could not reach the database of connection '<id>': ...``, or + ``Connection '<id>' does not provide a DbApiHook`` when the connection is + of the wrong type. Inside :doc:`../code_mode` the split does not hold: a + tool the model's code calls hands any exception back to the model as a + retry, this one included. + Once connected, an error from the statement becomes a ``ModelRetry`` so the + model can correct it, bounded by the retry budget + (:ref:`toolset-retry-budget`). A connection that drops while a ``query`` + runs falls in this second group, since the driver reports it like any + other statement error; the model usually runs the query again, on a new + connection. So does any failure of a hook whose ``get_conn`` does not open + a network connection, since nothing is reached until the statement runs. + Two paths do not raise: ``check_query`` catches its own errors and reports + them back as a normal ``{"valid": false, ...}`` result, and ``get_schema`` + returns a normal ``{"error": ...}`` result instead of raising when the + requested table is outside ``allowed_tables``. **A real example.** ``example_pydantic_ai_hook.py`` builds an agent around ``SQLToolset`` inside a plain ``@task`` function, with no operator involved: diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst index 545318a9392..b6b70b7394d 100644 --- a/providers/common/ai/docs/troubleshooting.rst +++ b/providers/common/ai/docs/troubleshooting.rst @@ -49,6 +49,16 @@ Model and connection errors version suffixes, for example), use the matching vendor connection type instead of the generic one. +``Could not reach the database of connection '...'`` + A ``SQLToolset`` tool could not reach the database the connection points at: the + connection is not defined, the host is unreachable, the credentials were rejected, + the driver or SQLAlchemy dialect failed to load or connect, or the connection was + lost while listing tables or reading a schema. + The attempt fails instead of handing the error to the model, and Airflow's retry + runs the call again. The driver's own message follows the colon. A connection that + is not defined, or a SQLAlchemy dialect that is not installed, fails every retry the + same way: define the connection or install the package. See :doc:`toolsets/sql`. + An ``ImportError`` for ``pydantic_ai.models.<vendor>``, the vendor SDK, or ``pydantic_ai_harness`` The provider is installed without the extra for that vendor, or without ``code-mode`` for ``pydantic_ai_harness``. Install it, quoting the 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 1311c9ee157..55c04629fd1 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 @@ -22,6 +22,8 @@ from contextlib import suppress from typing import TYPE_CHECKING, Any try: + from sqlalchemy.exc import DBAPIError + from airflow.providers.common.ai.utils.sql_validation import ( SQLSafetyError, collect_table_references, @@ -54,9 +56,10 @@ from airflow.providers.common.ai.utils.toolset_base import AirflowToolset, valid from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence from pydantic_ai._run_context import RunContext + from sqlalchemy.engine.reflection import Inspector # Sentinel distinguishing "caller did not pass ``allowed_tables``" (expose every @@ -104,6 +107,37 @@ _CHECK_QUERY_SCHEMA: dict[str, Any] = { } +class _ConnectionFailed(Exception): + """The toolset could not get a connection to its database, which no change to the SQL can fix.""" + + +def _unreachable(conn_id: str, error: Exception) -> _ConnectionFailed: + return _ConnectionFailed(f"Could not reach the database of connection {conn_id!r}: {error}") + + +def _raise_connection_failures(get_conn: Callable[[], Any], conn_id: str) -> Callable[[], Any]: + """Wrap a hook's ``get_conn`` so a failure to connect is told apart from a failing statement.""" + + def wrapper() -> Any: + try: + return get_conn() + except Exception as e: + raise _unreachable(conn_id, e) from e + + return wrapper + + +def _lost_connection(error: Exception) -> bool: + """ + Whether SQLAlchemy reports ``error`` as a connection failure rather than a failing statement. + + ``list_tables`` and ``get_schema`` go through the SQLAlchemy inspector, which raises a + failure to connect without a statement and marks a connection lost mid-statement as + invalidated. + """ + return isinstance(error, DBAPIError) and (error.statement is None or error.connection_invalidated) + + def _trusted_row_count(cursor: Any, *, fetched: int) -> int | None: """ Return the driver's row count for the query, or ``None`` when it cannot mean that. @@ -183,12 +217,14 @@ class SQLToolset(AirflowToolset): Uses a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook` resolved lazily from the given ``db_conn_id``. - When a tool fails, the database's error message is returned to the agent as a - retry (:class:`pydantic_ai.ModelRetry`) so the model can correct its SQL within - the run instead of failing the task. ``pydantic-ai`` bounds this by the tool's - ``max_retries``, so an unrecoverable error -- a bad connection or an auth - failure -- exhausts the retries and fails the task for Airflow to retry. The - toolset does not inspect the error type or message. + When the statement fails, the database's error message is returned to the agent as + a retry (:class:`pydantic_ai.ModelRetry`) so the model can correct its SQL within + the run, bounded by the tool's ``max_retries``. When the toolset cannot reach the + database -- the connection is missing or misconfigured, the host is unreachable, + the credentials are rejected, or SQLAlchemy reports the connection lost -- the + error is raised instead, failing the attempt for Airflow to retry: no change to + the SQL can fix it. Under durable execution a raised call is not recorded, so the + retry runs it against the database again rather than replaying the failure. :param db_conn_id: Airflow connection ID for the database. Templated when the toolset is passed to ``AgentOperator`` / ``@task.agent``, so each task @@ -377,16 +413,31 @@ class SQLToolset(AirflowToolset): def _get_db_hook(self) -> DbApiHook: if self._hook is None: - connection = BaseHook.get_connection(self._db_conn_id) - hook = connection.get_hook() + try: + hook = BaseHook.get_connection(self._db_conn_id).get_hook() + except Exception as e: + raise _unreachable(self._db_conn_id, e) from e if not isinstance(hook, DbApiHook): - raise ValueError( + raise _ConnectionFailed( f"Connection {self._db_conn_id!r} does not provide a DbApiHook. " f"Got {type(hook).__name__}." ) + # DbApiHook.run and every in-tree override connect through get_conn, so this is + # where the query tool's connection failures surface. + hook.get_conn = _raise_connection_failures(hook.get_conn, self._db_conn_id) # type: ignore[method-assign] self._hook = hook return self._hook + def _get_inspector(self) -> Inspector: + hook = self._get_db_hook() + # Building the inspector fetches credentials (Databricks), loads the SQLAlchemy dialect + # and connects, none of which goes through get_conn. The hook caches it, so the + # get_table_schema call after this reuses it. + try: + return hook.inspector + except Exception as e: + raise _unreachable(self._db_conn_id, e) from e + # ------------------------------------------------------------------ # AbstractToolset interface # ------------------------------------------------------------------ @@ -433,12 +484,15 @@ class SQLToolset(AirflowToolset): raise ValueError(f"Unknown tool: {name!r}") try: return await self.run_blocking(self._run_tool, name, tool_args) + except _ConnectionFailed: + # Handed to the model as a retry, a connection failure makes it give up and + # answer without data, or spend its retries on something no SQL can fix. + raise except Exception as e: - # Hand the database's own error back to the agent as a retry so it can - # read the message and fix its SQL within the run. pydantic-ai bounds - # this by the tool's max_retries, so an unrecoverable error (a bad - # connection, an auth failure) exhausts the budget and fails the task - # for Airflow to retry, rather than being silently worked around. + if _lost_connection(e): + raise _unreachable(self._db_conn_id, e) from e + # The statement failed: hand the database's own error back to the agent so it + # can fix its SQL within the run, bounded by the tool's max_retries. raise ModelRetry( f"The {name} tool failed: {e}\n" "Use the list_tables and get_schema tools to inspect the database, " @@ -466,7 +520,7 @@ class SQLToolset(AirflowToolset): return schema, table def _list_tables(self) -> str: - hook = self._get_db_hook() + inspector = self._get_inspector() tables: list[str] = [] # Dedupe by (schema, table) so a table reachable both qualified and via the # default schema (e.g. "public.users" and "users" with schema="public") is @@ -482,13 +536,13 @@ class SQLToolset(AirflowToolset): # Schemas referenced by qualified allowed_tables entries: introspect each # and return matching tables fully qualified so they round-trip to get_schema. for schema in sorted(self._qualified_schemas): - for name in hook.inspector.get_table_names(schema=schema): + for name in inspector.get_table_names(schema=schema): add(schema, name, f"{schema}.{name}") # Default schema: used for allow-all and unqualified allowed_tables entries. # Names stay bare to preserve the single-schema behaviour. if self._include_default_schema: - for name in hook.inspector.get_table_names(schema=self._schema): + for name in inspector.get_table_names(schema=self._schema): add(self._schema, name, name) return dumps_masked(tables) @@ -497,8 +551,8 @@ class SQLToolset(AirflowToolset): schema, table = self._split_table_identifier(table_name) if not self._is_ref_allowed("", schema, table): return dumps_masked({"error": f"Table {table_name!r} is not in the allowed tables list."}) - hook = self._get_db_hook() - columns = hook.get_table_schema(table, schema=schema) + self._get_inspector() # get_table_schema reads through it + columns = self._get_db_hook().get_table_schema(table, schema=schema) return build_schema_result( columns, max_columns=self._max_columns, @@ -509,7 +563,12 @@ class SQLToolset(AirflowToolset): def _dialect_for_validation(self) -> str | None: """Resolve the hook's sqlglot dialect so DESCRIBE/SHOW validate correctly.""" hook = self._get_db_hook() - return resolve_sqlglot_dialect(getattr(hook, "dialect_name", None)) + # dialect_name builds the hook's URI, which can fetch credentials (Databricks). + try: + dialect_name = getattr(hook, "dialect_name", None) + except Exception as e: + raise _unreachable(self._db_conn_id, e) from e + return resolve_sqlglot_dialect(dialect_name) def _query(self, sql: str) -> str: hook = self._get_db_hook() diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py index 9aba084a323..59b32994be2 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py @@ -47,8 +47,10 @@ from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits from airflow.providers.common.ai.durable import AirflowDurability from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope from airflow.providers.common.ai.exceptions import DurableJournalError +from airflow.providers.common.ai.toolsets.sql import SQLToolset, _ConnectionFailed from airflow.providers.common.ai.utils.toolset_base import AirflowToolset, ensure_masked +from unit.common.ai.sqlite_connection import create_orders_table, run_sql, sqlite_connection from unit.common.ai.tool_parts import is_tool_result, is_tool_retry if TYPE_CHECKING: @@ -554,6 +556,65 @@ class TestAirflowToolsets: assert calls["model"] == 3 +def _query_then_answer(sql: str, *, fail_final: list[bool]): + """A model that runs ``sql`` once, then answers with whatever the tool sent back, error or rows.""" + + def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + if responses_so_far(messages) == 0: + return ModelResponse(parts=[ToolCallPart("query", {"sql": sql})]) + if fail_final[0]: + raise RuntimeError("worker died") + return ModelResponse(parts=[TextPart(f"answer: {messages[-1].parts[0].content}")]) + + return model_fn + + +class TestSQLToolsetErrors: + """A database that could not be reached is queried again on retry; a rejected statement replays.""" + + @staticmethod + def _agent(model_fn) -> Agent[None, str]: + return Agent( + FunctionModel(model_fn), + name="analyst", + toolsets=[SQLToolset("warehouse")], + capabilities=[AirflowDurability()], + ) + + @pytest.mark.asyncio + async def test_a_query_that_could_not_connect_runs_again_on_retry(self, memory_storage, tmp_path): + path = tmp_path / "warehouse" / "warehouse.db" # its directory does not exist yet + model_fn = _query_then_answer("SELECT count(*) FROM orders", fail_final=[False]) + + with sqlite_connection("warehouse", path): + # A model handed the failure would answer with it; the attempt fails instead. + with pytest.raises( + _ConnectionFailed, match="connection 'warehouse': .*unable to open database file" + ): + await attempt(memory_storage, self._agent(model_fn)) + create_orders_table(path) + result = await attempt(memory_storage, self._agent(model_fn)) + + assert '"rows":[[2]]' in result.output + + @pytest.mark.asyncio + async def test_a_query_naming_a_missing_column_replays_on_retry(self, memory_storage, tmp_path): + path = tmp_path / "warehouse.db" + create_orders_table(path) + fail_final = [True] + model_fn = _query_then_answer("SELECT amount FROM orders", fail_final=fail_final) + + with sqlite_connection("warehouse", path): + with pytest.raises(RuntimeError): + await attempt(memory_storage, self._agent(model_fn)) + # Run live, the query would now succeed; replayed, it fails as it did. + run_sql(path, "ALTER TABLE orders ADD COLUMN amount REAL") + fail_final[0] = False + result = await attempt(memory_storage, self._agent(model_fn)) + + assert "no such column: amount" in result.output + + class Answer(BaseModel): rows: int diff --git a/providers/common/ai/tests/unit/common/ai/sqlite_connection.py b/providers/common/ai/tests/unit/common/ai/sqlite_connection.py new file mode 100644 index 00000000000..e5ef7f801df --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/sqlite_connection.py @@ -0,0 +1,85 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +A real ``DbApiHook`` over a SQLite file, for tests that need a database to fail like one. + +The sqlite provider is not a dependency of this one, so the hook is defined here over the +standard library's ``sqlite3``. Statements run through ``DbApiHook.run`` and introspection +through the SQLAlchemy inspector, so the errors the toolset sees are the driver's own. +""" + +from __future__ import annotations + +import sqlite3 +from collections.abc import Iterator +from contextlib import closing, contextmanager +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +from airflow.providers.common.compat.sdk import AirflowNotFoundException +from airflow.providers.common.sql.hooks.sql import DbApiHook + + +class SQLiteFileHook(DbApiHook): + conn_name_attr = "sqlite_conn_id" + default_conn_name = "sqlite_default" + + def __init__(self, path: Path) -> None: + super().__init__() + self.path = path + + def get_conn(self) -> sqlite3.Connection: + return sqlite3.connect(self.path) + + def get_uri(self) -> str: + return f"sqlite:///{self.path}" + + +def run_sql(path: Path, *statements: str) -> None: + """Run ``statements`` against the database at ``path``, creating it if needed.""" + path.parent.mkdir(parents=True, exist_ok=True) + with closing(sqlite3.connect(path)) as connection: + for statement in statements: + connection.execute(statement) + connection.commit() + + +def create_orders_table(path: Path) -> None: + """Create the database at ``path`` with an ``orders`` table of two rows.""" + run_sql( + path, + "CREATE TABLE orders (id INTEGER PRIMARY KEY, status TEXT)", + "INSERT INTO orders (status) VALUES ('shipped'), ('pending')", + ) + + +@contextmanager +def sqlite_connection(conn_id: str, path: Path) -> Iterator[None]: + """Resolve the Airflow connection ``conn_id`` to a fresh :class:`SQLiteFileHook` over ``path``.""" + + def get_connection(requested: str) -> SimpleNamespace: + if requested != conn_id: + raise AirflowNotFoundException(f"The conn_id `{requested}` isn't defined") + return SimpleNamespace(get_hook=lambda: SQLiteFileHook(path)) + + with patch( + "airflow.providers.common.ai.toolsets.sql.BaseHook.get_connection", + autospec=True, + side_effect=get_connection, + ): + yield diff --git a/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py b/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py index 76f5030eaec..a2bef09c39a 100644 --- a/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py +++ b/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py @@ -37,6 +37,7 @@ from airflow.providers.common.ai.tools._from_toolset import airflow_tools_from_t from airflow.providers.common.ai.toolsets.hook import HookToolset from airflow.providers.common.ai.toolsets.sql import SQLToolset +from unit.common.ai.sqlite_connection import sqlite_connection from unit.common.ai.toolsets.test_sql import _make_mock_db_hook @@ -102,11 +103,20 @@ class TestSQLToolsetAirflowTools: assert "sql" in result.content ts._hook.run.assert_not_called() - def test_a_database_that_keeps_failing_fails_the_call(self): - """Bad credentials are not the model's to fix; the tool's retry budget ends the loop.""" + def test_a_database_that_cannot_be_reached_fails_the_first_call(self, tmp_path): + ts = SQLToolset("warehouse") + + with ( + sqlite_connection("warehouse", tmp_path / "missing" / "warehouse.db"), + pytest.raises(ToolCallError, match="Could not reach the database of connection 'warehouse'"), + ): + asyncio.run(_by_name(ts.airflow_tools())["query"].call({"sql": "SELECT status FROM orders"})) + + def test_a_statement_that_keeps_failing_fails_the_call(self): + """The tool's retry budget ends a loop the model cannot correct its way out of.""" ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() - ts._hook.run.side_effect = ConnectionError("authentication failed") + ts._hook.run.side_effect = RuntimeError("canceling statement due to statement timeout") query = _by_name(ts.airflow_tools())["query"] first = asyncio.run(query.call({"sql": "SELECT 1"})) @@ -120,7 +130,7 @@ class TestSQLToolsetAirflowTools: ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() ts._hook.run.side_effect = RuntimeError( - f"could not connect to postgresql://svc:{registered_secret}@db" + f'invalid input syntax for type uuid: "postgresql://svc:{registered_secret}@db"' ) result = asyncio.run(_by_name(ts.airflow_tools())["query"].call({"sql": "SELECT 1"})) 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 dd16a1c1b17..4e5a5cca296 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 @@ -18,6 +18,8 @@ from __future__ import annotations import asyncio import json +import sqlite3 +import traceback from unittest.mock import MagicMock, PropertyMock, patch import pytest @@ -26,11 +28,14 @@ from pydantic_ai.exceptions import ModelRetry from pydantic_ai.toolsets.abstract import ToolsetTool from pydantic_core import ValidationError from sqlalchemy.engine.reflection import Inspector +from sqlalchemy.exc import OperationalError -from airflow.providers.common.ai.toolsets.sql import SQLToolset +from airflow.providers.common.ai.toolsets.sql import SQLToolset, _ConnectionFailed from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA from airflow.providers.common.sql.hooks.sql import DbApiHook +from unit.common.ai.sqlite_connection import SQLiteFileHook, create_orders_table, sqlite_connection + class _FakeCursor: """ @@ -83,12 +88,15 @@ def _assert_executed(hook, sql: str) -> None: assert hook.run.call_args.args[0] == sql +def _call(ts: SQLToolset, name: str, tool_args: dict) -> str: + return asyncio.run( + ts.call_tool(name, tool_args, ctx=MagicMock(spec=RunContext), tool=MagicMock(spec=ToolsetTool)) + ) + + def _query(ts: SQLToolset, sql: str) -> dict: """Run the ``query`` tool and decode its JSON result.""" - result = asyncio.run( - ts.call_tool("query", {"sql": sql}, ctx=MagicMock(spec=RunContext), tool=MagicMock(spec=ToolsetTool)) - ) - return json.loads(result) + return json.loads(_call(ts, "query", {"sql": sql})) class TestSQLToolsetInit: @@ -167,17 +175,6 @@ class TestSQLToolsetListTables: tables = json.loads(result) assert tables == ["orders"] - def test_introspection_error_raises_model_retry(self): - """A failure while listing tables is returned to the agent as a retry.""" - ts = SQLToolset("pg_default") - mock_hook = _make_mock_db_hook() - mock_hook.inspector.get_table_names.side_effect = Exception("could not connect to server") - ts._hook = mock_hook - - with pytest.raises(ModelRetry) as exc_info: - asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())) - assert "could not connect to server" in exc_info.value.message - class TestSQLToolsetGetSchema: def test_returns_column_info(self): @@ -486,13 +483,11 @@ class TestSQLToolsetQuery: [ Exception("001003 (42000): SQL compilation error: unexpected 'rows'"), RuntimeError("type mismatch"), - ConnectionError("could not connect to server"), ], ) def test_query_error_is_returned_to_agent_as_model_retry(self, error): - """Any error from the query, whatever its type, is handed back to the agent as a retry with - the database's own message. The toolset never inspects the error type or text; pydantic-ai's - max_retries bounds the loop, so an unrecoverable error still fails the task.""" + """Any error once connected, whatever its type, is handed back to the agent as a retry with + the database's own message; pydantic-ai's max_retries bounds the loop.""" ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() ts._hook.run.side_effect = error @@ -512,8 +507,8 @@ class TestSQLToolsetQuery: ): ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() - ts._hook.run.side_effect = ConnectionError( - f'connection to "db:5432" failed: password "{registered_secret}" rejected' + ts._hook.run.side_effect = RuntimeError( + f'syntax error at or near "{registered_secret}": password "{registered_secret}" rejected' ) with pytest.raises(ModelRetry) as exc_info: @@ -668,7 +663,7 @@ class TestSQLToolsetHookResolution: ts = SQLToolset("bad_conn") - with pytest.raises(ValueError, match="does not provide a DbApiHook"): + with pytest.raises(_ConnectionFailed, match="does not provide a DbApiHook"): ts._get_db_hook() @patch("airflow.providers.common.ai.toolsets.sql.BaseHook", autospec=True) @@ -686,6 +681,130 @@ class TestSQLToolsetHookResolution: mock_base_hook.get_connection.assert_called_once() +class TestSQLToolsetConnectionFailures: + """Failing to reach the database raises; a statement the database rejects goes back to the model.""" + + @pytest.mark.parametrize( + ("name", "tool_args"), + [ + ("query", {"sql": "SELECT status FROM orders"}), + ("list_tables", {}), + ("get_schema", {"table_name": "orders"}), + ], + ) + def test_raises_when_the_database_cannot_be_opened(self, tmp_path, name, tool_args): + ts = SQLToolset("warehouse") + + with ( + sqlite_connection("warehouse", tmp_path / "missing" / "warehouse.db"), + pytest.raises(_ConnectionFailed, match="connection 'warehouse': .*unable to open database file"), + ): + _call(ts, name, tool_args) + + def test_query_naming_a_missing_column_goes_back_to_the_model(self, tmp_path): + """SQLite raises OperationalError for this too, so the exception class cannot decide it.""" + path = tmp_path / "warehouse.db" + create_orders_table(path) + ts = SQLToolset("warehouse") + + with sqlite_connection("warehouse", path), pytest.raises(ModelRetry, match="no such column: amount"): + _call(ts, "query", {"sql": "SELECT amount FROM orders"}) + + @pytest.mark.parametrize( + ("connection_invalidated", "expected"), + [(True, _ConnectionFailed), (False, ModelRetry)], + ids=["connection-lost", "statement-rejected"], + ) + def test_introspection_error_raises_only_when_the_connection_was_lost( + self, connection_invalidated, expected + ): + ts = SQLToolset("warehouse") + ts._hook = _make_mock_db_hook() + ts._hook.inspector.get_table_names.side_effect = OperationalError( + "SELECT name FROM sqlite_master", + {}, + sqlite3.OperationalError("disk I/O error"), + connection_invalidated=connection_invalidated, + ) + + with pytest.raises(expected, match="disk I/O error"): + _call(ts, "list_tables", {}) + + def test_a_connection_that_is_not_defined_raises(self, tmp_path): + ts = SQLToolset("missing") + + with ( + sqlite_connection("warehouse", tmp_path / "warehouse.db"), + pytest.raises( + _ConnectionFailed, + match="Could not reach the database of connection 'missing': .*isn't defined", + ), + ): + _call(ts, "query", {"sql": "SELECT 1"}) + + @pytest.mark.parametrize( + ("name", "tool_args"), [("list_tables", {}), ("get_schema", {"table_name": "orders"})] + ) + def test_introspection_raises_when_sqlalchemy_cannot_be_set_up(self, tmp_path, name, tool_args): + """A failure before any connect, such as a dialect SQLAlchemy cannot load, is no SQL error either.""" + create_orders_table(tmp_path / "warehouse.db") + ts = SQLToolset("warehouse") + + with ( + sqlite_connection("warehouse", tmp_path / "warehouse.db"), + patch.object(SQLiteFileHook, "get_uri", autospec=True, return_value="nosuchdialect://"), + pytest.raises(_ConnectionFailed, match="connection 'warehouse': .*nosuchdialect"), + ): + _call(ts, name, tool_args) + + @pytest.mark.parametrize( + ("name", "tool_args"), + [ + ("query", {"sql": "SELECT status FROM orders"}), + ("list_tables", {}), + ("get_schema", {"table_name": "orders"}), + ], + ) + def test_raises_when_building_the_uri_fails(self, tmp_path, name, tool_args): + """Some hooks fetch credentials to build their URI, before any statement is sent.""" + create_orders_table(tmp_path / "warehouse.db") + ts = SQLToolset("warehouse") + + with ( + sqlite_connection("warehouse", tmp_path / "warehouse.db"), + patch.object( + SQLiteFileHook, + "get_uri", + autospec=True, + side_effect=RuntimeError("token endpoint unreachable"), + ), + pytest.raises(_ConnectionFailed, match="connection 'warehouse': token endpoint unreachable"), + ): + _call(ts, name, tool_args) + + @pytest.mark.enable_redact + def test_a_connection_failure_carrying_a_secret_is_raised_masked(self, tmp_path, registered_secret): + """Raised rather than handed to the model, the error still passes through the toolset's masking.""" + ts = SQLToolset("warehouse") + + with ( + sqlite_connection("warehouse", tmp_path / "warehouse.db"), + patch.object( + SQLiteFileHook, + "get_conn", + autospec=True, + side_effect=sqlite3.OperationalError( + f"could not open postgresql://svc:{registered_secret}@db" + ), + ), + pytest.raises(_ConnectionFailed) as exc_info, + ): + _call(ts, "query", {"sql": "SELECT 1"}) + + assert registered_secret not in "".join(traceback.format_exception(exc_info.value)) + assert "postgresql://svc:***@db" in str(exc_info.value) + + class TestSQLToolsetMultiSchema: """Schema-qualified allowed_tables span multiple schemas in one database."""
