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 6c2e7b36460 Fail the attempt when `SQLToolset` cannot reach its
database (#76783)
6c2e7b36460 is described below
commit 6c2e7b3646062af0c72528869a211274a775e8c3
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Nov 12 14:55:50 2026 +0000
Fail the attempt when `SQLToolset` cannot reach its database (#76783)
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.
The error is SQLToolsetConnectionError in
airflow.providers.common.ai.exceptions, so a retry policy can match it.
---
providers/common/ai/docs/code_mode.rst | 3 +
providers/common/ai/docs/durable_execution.rst | 5 +-
providers/common/ai/docs/toolsets/index.rst | 11 +-
providers/common/ai/docs/toolsets/sql.rst | 39 ++++-
providers/common/ai/docs/troubleshooting.rst | 11 ++
.../src/airflow/providers/common/ai/exceptions.py | 14 ++
.../airflow/providers/common/ai/toolsets/sql.py | 99 ++++++++---
.../unit/common/ai/durable/test_capability.py | 66 ++++++-
.../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 | 190 +++++++++++++++++----
11 files changed, 472 insertions(+), 69 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 01e6dc07408..bfb1f9b223c 100644
--- a/providers/common/ai/docs/durable_execution.rst
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -259,7 +259,10 @@ A tool that raises ``ModelRetry``, ``ToolFailed``,
``ApprovalRequired`` or
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; once it has failed the run on two
attempts in a
-row, the next retry asks the model again, as described above.
+row, the next retry asks the model again, as described above. ``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 fef8edee861..5c66466b514 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -271,8 +271,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 399262d5f37..914451fb2d2 100644
--- a/providers/common/ai/docs/toolsets/sql.rst
+++ b/providers/common/ai/docs/toolsets/sql.rst
@@ -458,15 +458,36 @@ 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 fails with
+ :class:`~airflow.providers.common.ai.exceptions.SQLToolsetConnectionError`,
+ whose message 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. Match that class in a retry policy
+ (:doc:`../retry_policies`) to retry an unreachable database on a different
+ schedule from other failures. When the driver's message carries a secret the
+ masker cannot remove from the exception, the task fails with a
``RuntimeError``
+ holding the masked message instead, and a rule on the class does not match
it.
+ 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 52ff66d09ec..40dd6e7c025 100644
--- a/providers/common/ai/docs/troubleshooting.rst
+++ b/providers/common/ai/docs/troubleshooting.rst
@@ -49,6 +49,17 @@ 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 with ``SQLToolsetConnectionError`` 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/exceptions.py
b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
index 201445725db..829084d806e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
@@ -76,6 +76,20 @@ class UnsupportedToolDeferralError(AirflowFailException):
"""
+class SQLToolsetConnectionError(AirflowException):
+ """
+ Raised when a ``SQLToolset`` tool cannot reach its database.
+
+ The connection is not defined or does not supply a ``DbApiHook``,
connecting fails, or
+ SQLAlchemy cannot be set up or loses the connection. No change to the SQL
fixes it, so
+ the attempt fails for Airflow to retry instead of the model getting the
error. Match it
+ in a retry policy to treat an unreachable database differently from other
failures.
+
+ When the driver's message carries a secret the masker cannot remove from
the exception,
+ the toolset raises a ``RuntimeError`` with the masked message instead.
+ """
+
+
class DurableToolMetadataError(AirflowFailException):
"""
Raised when a tool's ``airflow`` metadata is not ``True``, ``False`` or
``None``.
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 5da4eaaf48e..4c647276d24 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,
@@ -40,6 +42,7 @@ from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets.abstract import ToolsetTool
+from airflow.providers.common.ai.exceptions import SQLToolsetConnectionError
from airflow.providers.common.ai.utils.masking import dumps_masked
from airflow.providers.common.ai.utils.query_results import (
DEFAULT_MAX_COLUMNS,
@@ -55,9 +58,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
@@ -105,6 +109,33 @@ _CHECK_QUERY_SCHEMA: dict[str, Any] = {
}
+def _build_connection_error(conn_id: str, error: Exception) ->
SQLToolsetConnectionError:
+ return SQLToolsetConnectionError(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 _build_connection_error(conn_id, e) from e
+
+ return wrapper
+
+
+def _is_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.
@@ -184,12 +215,14 @@ class SQLToolset(ToolPrefixMixin, 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
@@ -384,16 +417,30 @@ class SQLToolset(ToolPrefixMixin, 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 _build_connection_error(self._db_conn_id, e) from e
if not isinstance(hook, DbApiHook):
- raise ValueError(
+ raise SQLToolsetConnectionError(
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.
+ try:
+ return hook.inspector
+ except Exception as e:
+ raise _build_connection_error(self._db_conn_id, e) from e
+
# ------------------------------------------------------------------
# AbstractToolset interface
# ------------------------------------------------------------------
@@ -446,12 +493,15 @@ class SQLToolset(ToolPrefixMixin, AirflowToolset):
raise ValueError(f"Unknown tool: {name!r}")
try:
return await self.run_blocking(self._run_tool, base, tool_args)
+ except SQLToolsetConnectionError:
+ # 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 _is_lost_connection(e):
+ raise _build_connection_error(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"
f"Use the {self._build_tool_name('list_tables')} and
{self._build_tool_name('get_schema')} tools "
@@ -480,7 +530,7 @@ class SQLToolset(ToolPrefixMixin, 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
@@ -496,13 +546,13 @@ class SQLToolset(ToolPrefixMixin, 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)
@@ -511,8 +561,10 @@ class SQLToolset(ToolPrefixMixin, 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)
+ columns = [
+ {"name": col["name"], "type": str(col["type"])}
+ for col in self._get_inspector().get_columns(table, schema=schema)
+ ]
return build_schema_result(
columns,
max_columns=self._max_columns,
@@ -524,7 +576,12 @@ class SQLToolset(ToolPrefixMixin, 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 _build_connection_error(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 7d643479abf..f23827c94c9 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
@@ -40,6 +40,7 @@ from pydantic_ai.mcp import MCPToolset
from pydantic_ai.messages import (
ModelMessage,
ModelResponse,
+ RetryPromptPart,
TextPart,
ToolCallPart,
ToolReturn,
@@ -59,11 +60,13 @@ from airflow.providers.common.ai.durable.capability import (
_JournalCodec,
)
from airflow.providers.common.ai.durable.journal import DurableJournal,
Unrecordable, journal_scope
-from airflow.providers.common.ai.exceptions import DurableJournalError
+from airflow.providers.common.ai.exceptions import DurableJournalError,
SQLToolsetConnectionError
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
ensure_masked
from airflow.providers.common.compat.sdk import AirflowFailException
from unit.common.ai.durable.memory_storage import MemoryStorage
+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:
@@ -942,6 +945,21 @@ def clock_then_lookup(calls: Calls, *, fail_final:
list[bool]):
return model_fn
+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")
+ answered = messages[-1].parts[0]
+ assert isinstance(answered, (RetryPromptPart, ToolReturnPart))
+ return ModelResponse(parts=[TextPart(f"answer: {answered.content}")])
+
+ return model_fn
+
+
class TestToolOptOut:
"""A tool whose metadata says ``{"airflow": False}`` runs live on every
attempt."""
@@ -1266,6 +1284,52 @@ class TestGetCallableToolsetId:
assert _get_callable_toolset_id(function) is None
+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(
+ SQLToolsetConnectionError, 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 68f6d7233a7..f3b65545d43 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
@@ -42,6 +42,7 @@ from airflow.providers.common.ai.tools._from_toolset import (
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
@@ -114,11 +115,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"}))
@@ -132,7 +142,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 57256030a57..cdbe3bf86a8 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,15 @@ 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.exceptions import SQLToolsetConnectionError
from airflow.providers.common.ai.toolsets.sql import SQLToolset
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:
"""
@@ -63,7 +69,7 @@ def _make_mock_db_hook(
mock = MagicMock(spec=DbApiHook)
mock.inspector = MagicMock(spec=Inspector)
mock.inspector.get_table_names.return_value = table_names or ["users",
"orders"]
- mock.get_table_schema.return_value = table_schema or [
+ mock.inspector.get_columns.return_value = table_schema or [
{"name": "id", "type": "INTEGER"},
{"name": "name", "type": "VARCHAR"},
]
@@ -83,12 +89,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:
@@ -281,17 +290,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):
@@ -307,7 +305,7 @@ class TestSQLToolsetGetSchema:
"columns": [{"name": "id", "type": "INTEGER"}, {"name": "name",
"type": "VARCHAR"}],
"column_count": 2,
}
- mock_hook.get_table_schema.assert_called_once_with("users",
schema=None)
+ mock_hook.inspector.get_columns.assert_called_once_with("users",
schema=None)
def test_name_contains_filters_the_columns(self):
"""``name_contains`` threads from the tool call through to the bounded
result."""
@@ -368,13 +366,13 @@ class TestSQLToolsetGetSchema:
assert "error" in data
assert "secrets" in data["error"]
# The guard must short-circuit before touching the database.
- ts._hook.get_table_schema.assert_not_called()
+ ts._hook.inspector.get_columns.assert_not_called()
def test_introspection_error_raises_model_retry(self):
"""A failure while reading a table's schema is returned to the agent
as a retry."""
ts = SQLToolset("pg_default")
ts._hook = _make_mock_db_hook()
- ts._hook.get_table_schema.side_effect = Exception('relation "users"
does not exist')
+ ts._hook.inspector.get_columns.side_effect = Exception('relation
"users" does not exist')
with pytest.raises(ModelRetry) as exc_info:
asyncio.run(
@@ -600,13 +598,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
@@ -626,8 +622,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:
@@ -782,7 +778,7 @@ class TestSQLToolsetHookResolution:
ts = SQLToolset("bad_conn")
- with pytest.raises(ValueError, match="does not provide a DbApiHook"):
+ with pytest.raises(SQLToolsetConnectionError, match="does not provide
a DbApiHook"):
ts._get_db_hook()
@patch("airflow.providers.common.ai.toolsets.sql.BaseHook", autospec=True)
@@ -800,6 +796,134 @@ 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(
+ SQLToolsetConnectionError, 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, SQLToolsetConnectionError), (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(
+ SQLToolsetConnectionError,
+ 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(SQLToolsetConnectionError, 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(
+ SQLToolsetConnectionError, 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(SQLToolsetConnectionError) 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."""
@@ -808,7 +932,7 @@ class TestSQLToolsetMultiSchema:
hook = MagicMock(spec=DbApiHook)
hook.inspector = MagicMock()
hook.inspector.get_table_names.side_effect = lambda schema=None:
tables_by_schema.get(schema, [])
- hook.get_table_schema.return_value = [{"name": "id", "type":
"INTEGER"}]
+ hook.inspector.get_columns.return_value = [{"name": "id", "type":
"INTEGER"}]
return hook
def test_list_tables_spans_multiple_schemas(self):
@@ -860,7 +984,9 @@ class TestSQLToolsetMultiSchema:
)
)
assert result == {"columns": [{"name": "id", "type": "INTEGER"}],
"column_count": 1}
-
ts._hook.get_table_schema.assert_called_once_with("DEPLOYMENT_IMAGE_DETAILS",
schema="MODEL_ASTRO")
+ ts._hook.inspector.get_columns.assert_called_once_with(
+ "DEPLOYMENT_IMAGE_DETAILS", schema="MODEL_ASTRO"
+ )
def test_get_schema_blocks_table_outside_allowed_schema(self):
ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.X"])
@@ -874,14 +1000,14 @@ class TestSQLToolsetMultiSchema:
)
)
assert "error" in result
- ts._hook.get_table_schema.assert_not_called()
+ ts._hook.inspector.get_columns.assert_not_called()
def test_get_schema_unqualified_uses_default_schema(self):
ts = SQLToolset("pg", schema="public")
ts._hook = self._schema_aware_hook({"public": ["users"]})
asyncio.run(ts.call_tool("get_schema", {"table_name": "users"},
ctx=MagicMock(), tool=MagicMock()))
- ts._hook.get_table_schema.assert_called_once_with("users",
schema="public")
+ ts._hook.inspector.get_columns.assert_called_once_with("users",
schema="public")
def test_list_tables_matches_case_insensitively(self):
"""Snowflake reflects unquoted names lowercased; uppercase
allowed_tables still match."""
@@ -914,7 +1040,9 @@ class TestSQLToolsetMultiSchema:
)
)
assert "error" not in result
-
ts._hook.get_table_schema.assert_called_once_with("deployment_image_details",
schema="MODEL_ASTRO")
+ ts._hook.inspector.get_columns.assert_called_once_with(
+ "deployment_image_details", schema="MODEL_ASTRO"
+ )
def test_list_tables_deduplicates_same_table(self):
"""A table listed both qualified and unqualified appears once."""
