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 5aadb5571ae Add a Databricks toolset for Unity Gateway MCP Services
(#74198)
5aadb5571ae is described below
commit 5aadb5571ae9f24bb2fd1b497f9f013bc8c9c4d1
Author: Firas Bouzazi <[email protected]>
AuthorDate: Sun Nov 8 02:16:14 2026 +0100
Add a Databricks toolset for Unity Gateway MCP Services (#74198)
Dag authors want their agents to use governed Databricks tools such as
Unity Catalog functions, Genie and AI Search through a Unity AI Gateway
MCP Service, without putting bearer tokens or gateway URLs in Dag code.
The generic MCP toolset needs both in a separate connection and sends a
fixed token, which expires during long agent runs with OAuth identities.
closes: #74193
---
docs/spelling_wordlist.txt | 1 +
providers/databricks/docs/index.rst | 5 +-
providers/databricks/docs/toolsets/index.rst | 33 +
providers/databricks/docs/toolsets/unity_mcp.rst | 126 +++
providers/databricks/provider.yaml | 8 +
providers/databricks/pyproject.toml | 10 +-
.../src/airflow/providers/databricks/exceptions.py | 31 +
.../providers/databricks/get_provider_info.py | 5 +
.../providers/databricks/toolsets/__init__.py | 17 +
.../providers/databricks/toolsets/unity_mcp.py | 629 ++++++++++++++
.../airflow/providers/databricks/version_compat.py | 2 +
.../databricks/example_databricks_unity_mcp.py | 69 ++
.../tests/unit/databricks/toolsets/__init__.py | 16 +
.../unit/databricks/toolsets/test_unity_mcp.py | 902 +++++++++++++++++++++
uv.lock | 12 +-
15 files changed, 1855 insertions(+), 11 deletions(-)
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index 3ba09fddae2..1c34f312e1d 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -1111,6 +1111,7 @@ Maxime
MaxRuntimeInSeconds
mb
MCP
+mcp
md
mem
memcached
diff --git a/providers/databricks/docs/index.rst
b/providers/databricks/docs/index.rst
index 0d7d40427a4..d3be73b253e 100644
--- a/providers/databricks/docs/index.rst
+++ b/providers/databricks/docs/index.rst
@@ -36,6 +36,7 @@
Connection types <connections/databricks>
Operators <operators/index>
+ Toolsets <toolsets/index>
Plugins <plugins/index>
.. toctree::
@@ -147,16 +148,16 @@ Install them when installing from PyPI. For example:
.. code-block:: bash
- pip install apache-airflow-providers-databricks[common.ai]
+ pip install apache-airflow-providers-databricks[avro]
==================
================================================================================================================================================================
Extra Dependencies
==================
================================================================================================================================================================
-``common.ai`` ``apache-airflow-providers-common-ai>=1.0.0``
``avro`` ``fastavro>=1.9.0; python_version<"3.14"``,
``fastavro>=1.10.0; python_version>="3.12" and python_version<"3.14"``,
``fastavro>=1.12.1; python_version>="3.14"``
``amazon`` ``apache-airflow-providers-amazon>=9.22.0``
``azure-identity`` ``azure-identity>=1.25.3``
+``common.ai`` ``apache-airflow-providers-common-ai[mcp]>=1.0.0``,
``fastmcp-slim[client]>=4.0.0``
``fab`` ``apache-airflow-providers-fab>=2.2.0``
``google`` ``apache-airflow-providers-google>=10.24.0``
``sdk`` ``databricks-sdk==0.10.0``
diff --git a/providers/databricks/docs/toolsets/index.rst
b/providers/databricks/docs/toolsets/index.rst
new file mode 100644
index 00000000000..cdad1abfd69
--- /dev/null
+++ b/providers/databricks/docs/toolsets/index.rst
@@ -0,0 +1,33 @@
+ .. 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.
+
+
+Databricks Toolsets
+===================
+
+Toolsets give agents from the :doc:`Common AI provider
<apache-airflow-providers-common-ai:index>`
+tools backed by Databricks. They need the ``common.ai`` extra:
+
+.. code-block:: bash
+
+ pip install 'apache-airflow-providers-databricks[common.ai]'
+
+.. toctree::
+ :maxdepth: 1
+ :glob:
+
+ *
diff --git a/providers/databricks/docs/toolsets/unity_mcp.rst
b/providers/databricks/docs/toolsets/unity_mcp.rst
new file mode 100644
index 00000000000..b61c5d64f58
--- /dev/null
+++ b/providers/databricks/docs/toolsets/unity_mcp.rst
@@ -0,0 +1,126 @@
+ .. 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.
+
+.. _howto/toolset:DatabricksUnityMCPToolset:
+
+Unity Gateway MCP Services
+==========================
+
+Use
:class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset`
to give an
+agent the tools of a Unity Gateway MCP Service: either one Databricks provides
for workspace tools
+and SaaS applications, such as ``system.ai.google_calendar``, or an external
MCP server registered as
+an MCP Service in Unity Catalog. The toolset works with
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`,
``@task.agent``, and the LangChain
+bridge of the Common AI provider.
+
+The Dag names the service by its three-level Unity Catalog name,
``catalog.schema.service``, and the
+:ref:`Databricks connection <howto/connection:databricks>` to use. The toolset
builds the service URL,
+``https://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``,
from the connection's
+host, so neither the gateway URL nor a token appears in Dag code, and the
connection's token is only
+sent to that workspace, over HTTPS. A connection whose ``schema`` is ``http``
is rejected unless its
+host is a loopback address. Requests go through the proxy in the connection's
``proxies`` extra, if
+set. Service names may contain only ASCII letters, digits and ``_`` in each
part; any other name is
+rejected before a request is made.
+
+.. exampleinclude::
/../../databricks/tests/system/databricks/example_databricks_unity_mcp.py
+ :language: python
+ :start-after: [START howto_toolset_databricks_unity_mcp]
+ :end-before: [END howto_toolset_databricks_unity_mcp]
+
+The service name and connection ID are templated when the toolset is passed to
``AgentOperator`` or
+``@task.agent``.
+
+Caller identity and authentication
+----------------------------------
+
+MCP Services need a workspace enabled for Unity Catalog, in a region where
Model Serving is
+supported.
+
+The gateway runs every tool call as the identity of the connection's
credentials. That identity needs
+``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its
parent catalog and schema
+(``EXECUTE`` alone is not enough), and an assignment to the workspace. It
needs no privilege on the
+Unity Catalog connection behind the service. On the built-in ``system.ai``
services, account users
+hold these privileges by default (see `MCP Services
+<https://docs.databricks.com/aws/en/agents/mcp-tools/mcp-services>`__). The
gateway exposes only the
+tools selected for the service, and the service's policies apply. Grant the
identity only the services
+the agent should use.
+
+Built-in services that act on a user's own data, such as
``system.ai.google_calendar`` or
+``system.ai.gmail``, need that identity to complete a one-time OAuth login
first, for example by
+opening the service in Catalog Explorer and clicking **Login**.
+
+The toolset sends the connection's token as a bearer token, so the connection
must use one of these
+authentication modes of the Databricks connection:
+
+* a personal access token (the identity is the token's user or service
principal);
+* service principal OAuth (``service_principal_oauth``);
+* Azure AD: a service principal, a managed identity, or
``DefaultAzureCredential``;
+* workload identity federation (Kubernetes, AWS IAM, or a supplied token
provider).
+
+For service principal OAuth, the OAuth secret must allow the ``all-apis``
scope, which the connection
+requests; a secret restricted to narrower scopes fails (see `OAuth for service
principals
+<https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m>`__). Username
and password
+authentication is not supported. A token is fetched from the connection for
+every request, so OAuth and Azure AD tokens are refreshed during a long agent
run and when the toolset
+reconnects. Tokens minted for the toolset are masked in task logs.
+
+Errors and retries
+------------------
+
+When the MCP server answers a tool call with an error, such as invalid
arguments, the error goes to the
+model, which can correct the call and try again.
+
+When the gateway gives no answer of its own, a tool call that was already sent
may or may not have
+run, and repeating a tool that changes data could apply the change twice. The
toolset never retries
+it: the task fails with one of these exceptions from
:mod:`airflow.providers.databricks.exceptions`,
+which are also raised for failures while connecting and listing tools:
+
+.. list-table::
+ :header-rows: 1
+
+ * - Exception
+ - Cause
+ * - ``DatabricksUnityMCPAccessDeniedError``
+ - The gateway will not let the identity invoke the service: no service
has that name, the
+ identity lacks ``EXECUTE`` on the service or ``USE CATALOG`` / ``USE
SCHEMA`` on its parents,
+ or the credentials are invalid. A missing service is reported like a
missing privilege.
+ * - ``DatabricksUnityMCPThrottledError``
+ - HTTP 429. Its ``retry_after`` attribute holds the ``Retry-After``
delay when the gateway gives
+ it as a number of seconds.
+ * - ``DatabricksUnityMCPTransportError``
+ - The gateway could not be reached, the connection dropped, a request
got no answer in time,
+ or the MCP handshake did not finish within ``init_timeout``.
+ * - ``DatabricksUnityMCPError``
+ - Any other gateway error, or a failure to get a token from the
connection.
+
+A connection that cannot be used at all, such as one without a host or without
token-based
+authentication, raises ``ValueError`` when the toolset connects.
+
+Unity Gateway forwards a response only once the MCP server has finished
sending it, so connecting
+to a server that keeps its response stream open takes as long as the server
keeps it open. The toolset
+waits ``init_timeout`` seconds (30 by default) for the handshake; raise it for
slow services.
+
+The exceptions carry the HTTP status in ``http_status_code`` when it is known.
When the agent runs
+several tool calls of the toolset at once, the status of a failed call cannot
be told apart from the
+others', so a call that fails with HTTP 401, 403 or 429 is reported as
``DatabricksUnityMCPError``
+without it. Transport errors, timeouts and a denial the gateway reports in its
JSON-RPC answer keep
+their class. Use task retries for calls that are safe to repeat.
+
+These exception classes and their attributes reach the task in Pydantic AI
runs. Through the
+LangChain bridge, ``airflow_toolset_to_langchain_tools``, a failed tool call
raises
+``ToolCallError`` instead, whose message names the exception class and repeats
its message;
+``http_status_code`` and ``retry_after`` are not carried over.
diff --git a/providers/databricks/provider.yaml
b/providers/databricks/provider.yaml
index c11d7a3e8e2..894c986dd59 100644
--- a/providers/databricks/provider.yaml
+++ b/providers/databricks/provider.yaml
@@ -147,6 +147,9 @@ integrations:
how-to-guide:
- /docs/apache-airflow-providers-databricks/operators/workflow.rst
tags: [service]
+ - integration-name: Databricks Unity Gateway
+ external-doc-url: https://docs.databricks.com/aws/en/unity-gateway/concepts
+ tags: [service]
operators:
- integration-name: Databricks
@@ -201,6 +204,11 @@ sensors:
- airflow.providers.databricks.sensors.databricks_sql
- airflow.providers.databricks.sensors.databricks_partition
+toolsets:
+ - integration-name: Databricks Unity Gateway
+ python-modules:
+ - airflow.providers.databricks.toolsets.unity_mcp
+
connection-types:
- hook-class-name:
airflow.providers.databricks.hooks.databricks.DatabricksHook
hook-name: "Databricks"
diff --git a/providers/databricks/pyproject.toml
b/providers/databricks/pyproject.toml
index 1faa8dde64f..27a48f2b7de 100644
--- a/providers/databricks/pyproject.toml
+++ b/providers/databricks/pyproject.toml
@@ -77,9 +77,6 @@ dependencies = [
# The optional dependencies should be modified in place in the generated file
# Any change in the dependencies is preserved when the file is regenerated
[project.optional-dependencies]
-"common.ai" = [
- "apache-airflow-providers-common-ai>=1.0.0",
-]
# pip install apache-airflow-providers-databricks[sdk]
"avro" = [
'fastavro>=1.9.0; python_version<"3.14"',
@@ -92,6 +89,11 @@ dependencies = [
"azure-identity" = [
"azure-identity>=1.25.3",
]
+"common.ai" = [
+ "apache-airflow-providers-common-ai[mcp]>=1.0.0",
+ # 3.x is built on httpx and MCP SDK 1.x; the toolset needs the httpx2 and
MCP SDK 2.x of 4.0.
+ "fastmcp-slim[client]>=4.0.0",
+]
"fab" = [
"apache-airflow-providers-fab>=2.2.0"
]
@@ -129,7 +131,7 @@ dev = [
"apache-airflow-providers-microsoft-azure",
"apache-airflow-providers-common-sql[pandas,polars]",
"apache-airflow-providers-fab",
- "apache-airflow-providers-databricks[sqlalchemy]",
+ "apache-airflow-providers-databricks[common.ai,sqlalchemy]",
]
# To build docs:
diff --git
a/providers/databricks/src/airflow/providers/databricks/exceptions.py
b/providers/databricks/src/airflow/providers/databricks/exceptions.py
index b48e5dd0059..01d47523929 100644
--- a/providers/databricks/src/airflow/providers/databricks/exceptions.py
+++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py
@@ -54,3 +54,34 @@ class DatabricksApiError(AirflowException):
def __init__(self, message: str, *, http_status_code: int | None = None)
-> None:
super().__init__(message)
self.http_status_code = http_status_code
+
+
+class DatabricksUnityMCPError(DatabricksApiError):
+ """Raised when a call to a Unity Gateway MCP Service fails."""
+
+
+class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError):
+ """
+ Raised when the gateway will not let the caller invoke the service.
+
+ The service may not exist, the caller may lack a privilege on it, or its
credentials may be invalid.
+ """
+
+
+class DatabricksUnityMCPThrottledError(DatabricksUnityMCPError):
+ """Raised when the gateway rate-limits the caller."""
+
+ def __init__(
+ self, message: str, *, http_status_code: int | None = None,
retry_after: float | None = None
+ ) -> None:
+ super().__init__(message, http_status_code=http_status_code)
+ self.retry_after = retry_after
+
+
+class DatabricksUnityMCPTransportError(DatabricksUnityMCPError):
+ """
+ Raised when the gateway cannot be reached or the connection drops.
+
+ When this interrupts a tool call, the tool may or may not have run, so the
call is not
+ retried automatically.
+ """
diff --git
a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
index 397706c99e2..76af2eebfc1 100644
--- a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
+++ b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
@@ -68,6 +68,11 @@ def get_provider_info():
"how-to-guide":
["/docs/apache-airflow-providers-databricks/operators/workflow.rst"],
"tags": ["service"],
},
+ {
+ "integration-name": "Databricks Unity Gateway",
+ "external-doc-url":
"https://docs.databricks.com/aws/en/unity-gateway/concepts";,
+ "tags": ["service"],
+ },
],
"operators": [
{
diff --git
a/providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py
b/providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py
new file mode 100644
index 00000000000..217e5db9607
--- /dev/null
+++ b/providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py
@@ -0,0 +1,17 @@
+#
+# 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.
diff --git
a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py
b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py
new file mode 100644
index 00000000000..b69f9c169f9
--- /dev/null
+++
b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py
@@ -0,0 +1,629 @@
+# 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.
+"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP
Service."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import ipaddress
+import json
+import logging
+import re
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlsplit
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, mask_secret
+from airflow.providers.databricks.exceptions import (
+ DatabricksUnityMCPAccessDeniedError,
+ DatabricksUnityMCPError,
+ DatabricksUnityMCPThrottledError,
+ DatabricksUnityMCPTransportError,
+)
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.providers.databricks.version_compat import AIRFLOW_V_3_2_PLUS
+
+try:
+ # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the
common.ai extra installs.
+ import httpx2
+ from fastmcp.client.transports import StreamableHttpTransport
+ from fastmcp.exceptions import McpError
+ from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset
+
+ from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "DatabricksUnityMCPToolset needs the 'common.ai' extra of the
databricks provider: "
+ "pip install 'apache-airflow-providers-databricks[common.ai]'"
+ ) from e
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncGenerator, AsyncIterator, Callable,
Iterator, Sequence
+
+ from pydantic_ai._run_context import RunContext
+ from pydantic_ai.toolsets.abstract import ToolsetTool
+
+ from airflow.providers.common.ai.tools import AirflowTool
+
+log = logging.getLogger(__name__)
+
+GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services"
+
+# Restricting each part to these characters keeps the name from adding a path,
query or
+# fragment to the URL, so the request can only reach the service path on the
connection's host.
+_SERVICE_NAME_PART = r"[A-Za-z0-9_]+"
+_SERVICE_NAME =
re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}")
+
+# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the
caller may not invoke
+# the service, including when no service has that name.
+_NOT_AUTHORIZED_CODE = -32007
+
+# Errors the MCP client makes up when the gateway gave no JSON-RPC answer,
such as an HTTP error
+# with a plain body or a response that never arrived (see
mcp.client.streamable_http). Every other
+# MCP error is the server's own answer. Matched on code and message, the only
things they carry.
+_CLIENT_ERRORS = {
+ (-32603, "Server returned an error response"),
+ (-32601, "Not Found"),
+ (-32600, "Session terminated"),
+ (-32000, "Connection closed"),
+}
+# What fastmcp raises, in place of the TimeoutError it loses, when the MCP
handshake outlasts
+# init_timeout (see fastmcp.client.client).
+_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session"
+# What fastmcp raises for a call that reaches the session after another
request ended it
+# (see fastmcp.client.client).
+_SESSION_ENDED_MESSAGE = "Client is not connected"
+# What the MCP client raises when a request outlasts its read timeout (see
+# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the
message is matched too.
+_REQUEST_TIMEOUT_CODE = -32001
+_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out")
+_CLIENT_ERROR_PREFIXES = (
+ "Unexpected content type:",
+ "Failed to parse JSON response:",
+ "Failed to parse SSE message:",
+ "SSE stream ",
+ "server answered a request with 202 Accepted",
+ "Redirect to ",
+)
+
+
+def validate_service_name(service_name: str) -> None:
+ """
+ Raise ``ValueError`` unless ``service_name`` is a three-level
``catalog.schema.service`` name.
+
+ Each part may contain only ASCII letters, digits and underscores.
+ """
+ if not isinstance(service_name, str) or not
_SERVICE_NAME.fullmatch(service_name):
+ raise ValueError(
+ f"Invalid Unity Gateway MCP Service name {service_name!r}:
expected "
+ "'catalog.schema.service', each part made of ASCII letters, digits
or '_'."
+ )
+
+
+def _is_loopback(host: str) -> bool:
+ if host == "localhost":
+ return True
+ try:
+ return ipaddress.ip_address(host.strip("[]")).is_loopback
+ except ValueError:
+ return False
+
+
+def _parse_retry_after(value: str | None) -> float | None:
+ if value is None:
+ return None
+ try:
+ seconds = float(value)
+ except ValueError:
+ # The HTTP-date form is not worth parsing: the caller only uses this
as a hint.
+ return None
+ return seconds if seconds >= 0 else None
+
+
+def _walk_exceptions(exc: BaseException) -> Iterator[BaseException]:
+ """Yield ``exc``, its causes and the exceptions inside any exception
groups, each once."""
+ seen: set[int] = set()
+ pending: list[BaseException] = [exc]
+ while pending:
+ current = pending.pop(0)
+ if id(current) in seen:
+ continue
+ seen.add(id(current))
+ yield current
+ # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio
uses the backport on 3.10.
+ if isinstance(grouped := getattr(current, "exceptions", None), (list,
tuple)):
+ pending.extend(e for e in grouped if isinstance(e, BaseException))
+ pending.extend(e for e in (current.__cause__, current.__context__) if
e is not None)
+
+
+def _describe(error: BaseException) -> str:
+ return str(error) or type(error).__name__
+
+
+def _find(exc: BaseException, types: type[BaseException] |
tuple[type[BaseException], ...]) -> Any:
+ """Return the first exception of ``types`` in ``exc``, its causes and any
exception groups."""
+ return next((e for e in _walk_exceptions(exc) if isinstance(e, types)),
None)
+
+
+def _is_handshake_timeout(error: BaseException) -> bool:
+ return any(
+ isinstance(e, TimeoutError)
+ or (isinstance(e, RuntimeError) and
str(e).endswith(_HANDSHAKE_TIMEOUT_MESSAGE))
+ for e in _walk_exceptions(error)
+ )
+
+
+def _is_request_timeout(error: McpError | None) -> bool:
+ return (
+ error is not None
+ and error.error.code == _REQUEST_TIMEOUT_CODE
+ and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None
+ )
+
+
+def _is_session_ended(error: BaseException) -> bool:
+ return any(
+ isinstance(e, RuntimeError) and
str(e).startswith(_SESSION_ENDED_MESSAGE)
+ for e in _walk_exceptions(error)
+ )
+
+
+def _is_client_error(error: McpError) -> bool:
+ code, message = error.error.code, error.error.message
+ return (code, message) in _CLIENT_ERRORS or
message.startswith(_CLIENT_ERROR_PREFIXES)
+
+
+class _TokenUnavailableError(Exception):
+ """Raised, for this module only, when no token could be fetched for a
gateway request."""
+
+
+@dataclass(eq=False)
+class _Operation:
+ """
+ One toolset operation (connecting, listing tools or calling a tool) and
the gateway errors seen while it ran.
+
+ The MCP client does not say which HTTP response made a call fail, so the
error statuses are
+ only attributed to the operation when no other operation of the toolset
ran at the same time.
+ Request failures are recorded on every operation running then and kept
whatever the overlap:
+ they only explain a call that saw no answer, and a request that got no
response usually ends
+ the session for all of them. An overlapping call that fails for its own
reason may still get
+ another call's failure, which changes its class and message but not that
it fails.
+ """
+
+ overlapped: bool = False
+ errors: list[tuple[int, float | None]] = field(default_factory=list)
+ failures: list[Exception] = field(default_factory=list)
+
+ @property
+ def http_error(self) -> tuple[int, float | None] | None:
+ """Return the status and ``Retry-After`` of the last error response it
got, if it is known."""
+ return self.errors[-1] if self.errors and not self.overlapped else None
+
+
+class _DatabricksTokenAuth(httpx2.Auth):
+ """
+ Authenticate each gateway request with a token from the Databricks
connection.
+
+ Asking the hook on every request, rather than once, lets OAuth tokens
refresh during a long
+ agent run and on reconnection; the hook caches tokens until they are about
to expire.
+ """
+
+ def __init__(self, hook: DatabricksHook) -> None:
+ self._hook = hook
+ self._token_lock = threading.Lock()
+ self._masked_token: str | None = None
+
+ def fetch_token(self) -> str:
+ with self._token_lock:
+ token = self._hook._get_token(raise_error=False)
+ if not token:
+ raise ValueError(
+ f"Connection {self._hook.databricks_conn_id!r} has no
token-based authentication "
+ "configured. This toolset sends a bearer token: use a personal
access token, "
+ "service principal OAuth, Azure AD, or workload identity
federation. Username and "
+ "password authentication is not supported."
+ )
+ return token
+
+ def fetch_and_mask_token(self) -> None:
+ """Fetch a token and mask it; only call it through
``AirflowToolset.run_blocking``."""
+ self._mask(self.fetch_token())
+
+ def _mask(self, token: str) -> None:
+ # A personal access token is masked when the connection is fetched;
mask minted
+ # OAuth tokens too, so they never reach task logs.
+ mask_secret(token)
+ self._masked_token = token
+
+ async def async_auth_flow(
+ self, request: httpx2.Request
+ ) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+ try:
+ if AIRFLOW_V_3_2_PLUS:
+ # Not AirflowToolset.run_blocking: its lock is shared by every
toolset, so a long
+ # SQL query in another toolset would hold up each request
here. The hook is this
+ # toolset's own, so a lock of its own is enough.
+ token = await asyncio.to_thread(self.fetch_token)
+ else:
+ # Reading the connection's extras masks them, which messages
the supervisor over
+ # a socket that has no lock of its own before Airflow 3.2.
+ token = await AirflowToolset.run_blocking(self.fetch_token)
+ except Exception as e:
+ raise _TokenUnavailableError(str(e)) from e
+ if token != self._masked_token:
+ # mask_secret messages the supervisor too.
+ await AirflowToolset.run_blocking(self._mask, token)
+ request.headers["Authorization"] = f"Bearer {token}"
+ yield request
+
+
+class _RecordingStream(httpx2.AsyncByteStream):
+ """A response body that records a transport error raised while it is
read."""
+
+ def __init__(self, stream: httpx2.AsyncByteStream, record:
Callable[[Exception], None]) -> None:
+ self._stream = stream
+ self._record = record
+
+ async def __aiter__(self) -> AsyncIterator[bytes]:
+ try:
+ async for chunk in self._stream:
+ yield chunk
+ except httpx2.TransportError as e:
+ self._record(e)
+ raise
+
+ async def aclose(self) -> None:
+ await self._stream.aclose()
+
+
+class _GatewayClient(httpx2.AsyncClient):
+ """
+ Record, on the toolset's running operations, why each JSON-RPC request to
the gateway failed.
+
+ Once the MCP session is open, requests are sent from the MCP client's own
tasks, and a call
+ waiting for an answer only learns that the connection closed. These
records say why.
+ """
+
+ def __init__(self, *args: Any, operations: set[_Operation], **kwargs: Any)
-> None:
+ super().__init__(*args, **kwargs)
+ self._operations = operations
+
+ async def send(self, request: httpx2.Request, **kwargs: Any) ->
httpx2.Response:
+ if not _is_jsonrpc_request(request):
+ return await super().send(request, **kwargs)
+
+ def record(failure: Exception) -> None:
+ for operation in self._operations:
+ operation.failures.append(failure)
+
+ try:
+ response = await super().send(request, **kwargs)
+ except (httpx2.TransportError, _TokenUnavailableError) as e:
+ record(e)
+ raise
+ if response.status_code >= 400:
+ retry_after =
_parse_retry_after(response.headers.get("Retry-After"))
+ for operation in self._operations:
+ operation.errors.append((response.status_code, retry_after))
+ if isinstance(response.stream, httpx2.AsyncByteStream):
+ response.stream = _RecordingStream(response.stream, record)
+ return response
+
+
+def _is_jsonrpc_request(request: httpx2.Request) -> bool:
+ """
+ Whether ``request`` sends a JSON-RPC request, which expects an answer.
+
+ Not a notification or a GET for server events: the MCP client tolerates
errors on those, so
+ they say nothing about why an operation failed.
+ """
+ if request.method != "POST":
+ return False
+ try:
+ return "id" in json.loads(request.content)
+ except (ValueError, TypeError):
+ return False
+
+
+class DatabricksUnityMCPToolset(AirflowToolset):
+ """
+ Give an agent the tools of a Unity Gateway MCP Service, authenticated as
the connection's identity.
+
+ The service is named by its three-level Unity Catalog name,
``catalog.schema.service``, and
+ reached at ``https://<workspace
host>/ai-gateway/mcp-services/<catalog.schema.service>``. The
+ workspace host and the credentials both come from the Databricks
connection, so Dag code holds
+ neither a gateway URL nor a token, and the token is only ever sent to that
workspace, over
+ HTTPS. The connection's ``proxies`` extra applies to gateway requests.
+
+ The gateway runs every tool call as the identity of the connection's
credentials (a user's
+ personal access token, a service principal, or an Azure AD / federated
identity). That identity
+ needs ``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA``
on its catalog and
+ schema, and an assignment to the workspace. It sees only the tools
selected for the service, and
+ the service's policies apply.
+
+ Tokens are fetched from the connection for each request, so OAuth tokens
refresh during a long
+ agent run and on reconnection. An error the server returns for a tool call
reaches the model,
+ which can correct its arguments and try again. When the gateway gives no
answer of its own, a
+ tool call may or may not have run, so it is never retried; it fails the
task with
+
:class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`,
+
:class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`,
+
:class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`,
or
+ :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`,
as do failures
+ while connecting and listing tools.
+
+ .. code-block:: python
+
+ from airflow.providers.common.ai.operators.agent import AgentOperator
+ from airflow.providers.databricks.toolsets.unity_mcp import
DatabricksUnityMCPToolset
+
+ AgentOperator(
+ task_id="ask_mcp_service",
+ prompt="Which tools do you have?",
+ llm_conn_id="pydanticai_default",
+ toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp",
databricks_conn_id="databricks")],
+ )
+
+ :param service_name: Three-level name of the MCP Service,
``catalog.schema.service``. Templated
+ when the toolset is passed to ``AgentOperator`` / ``@task.agent``.
+ :param databricks_conn_id: Databricks connection whose host and
credentials are used. Templated
+ when the toolset is passed to ``AgentOperator`` / ``@task.agent``.
+ :param tool_prefix: Optional prefix prepended to tool names.
+ :param init_timeout: Seconds to wait for the MCP handshake when
connecting. Unity Gateway forwards
+ a response only once the MCP server has finished sending it, so a
server that keeps its
+ response stream open needs more than pydantic-ai's default of 5
seconds.
+ """
+
+ # Rendered, on a copy, by AgentOperator.
+ agent_template_fields: Sequence[str] = ("_databricks_conn_id",
"_service_name")
+
+ def __init__(
+ self,
+ service_name: str,
+ *,
+ databricks_conn_id: str = DatabricksHook.default_conn_name,
+ tool_prefix: str | None = None,
+ init_timeout: float = 30,
+ ) -> None:
+ if not init_timeout > 0:
+ raise ValueError(f"init_timeout must be a positive number of
seconds, got {init_timeout!r}.")
+ self._databricks_conn_id = databricks_conn_id
+ self._service_name = service_name
+ self._tool_prefix = tool_prefix
+ self._init_timeout = init_timeout
+ self._server: Any = None
+ self._operations: set[_Operation] = set()
+ # Whether a gateway failure was reported while connected, so closing
the session, which
+ # re-raises that failure, must not replace the report.
+ self._failure_reported = False
+ # A templated name is checked once it has been rendered, when the
toolset connects.
+ if "{{" not in service_name:
+ validate_service_name(service_name)
+
+ @property
+ def id(self) -> str:
+ return
f"databricks-unity-mcp-{self._databricks_conn_id}-{self._service_name}"
+
+ def get_service_url(self, hook: DatabricksHook) -> str:
+ """Return the gateway URL of the MCP Service on the connection's
workspace."""
+ validate_service_name(self._service_name)
+ if not hook.host:
+ raise ValueError(f"Connection {self._databricks_conn_id!r} has no
workspace host.")
+ url =
hook._endpoint_url(f"{GATEWAY_MCP_SERVICES_PATH}/{self._service_name}")
+ parts = urlsplit(url)
+ if parts.scheme != "https" and not _is_loopback(parts.hostname or ""):
+ raise ValueError(
+ f"Connection {self._databricks_conn_id!r} uses the
{parts.scheme!r} scheme. Unity Gateway "
+ "requests carry a bearer token, so they are only sent over
HTTPS."
+ )
+ return url
+
+ def _get_server(self) -> Any:
+ hook = DatabricksHook(self._databricks_conn_id,
caller=type(self).__name__)
+ url = self.get_service_url(hook)
+ auth = _DatabricksTokenAuth(hook)
+ # Fail here, with a clear message, rather than inside the MCP client,
which reports
+ # any failure as a generic connection error.
+ try:
+ auth.fetch_and_mask_token()
+ except ValueError:
+ raise
+ except Exception as e:
+ raise _TokenUnavailableError(str(e)) from e
+ proxy = (hook.proxies or {}).get(urlsplit(url).scheme)
+ transport = StreamableHttpTransport(
+ url,
+ headers=hook.user_agent_header,
+ auth=auth,
+ httpx_client_factory=_build_client_factory(self._operations,
proxy),
+ )
+ toolset = PydanticAIMCPToolset(transport,
init_timeout=self._init_timeout)
+ return toolset.prefixed(self._tool_prefix) if self._tool_prefix else
toolset
+
+ async def _resolve_server(self) -> Any:
+ if self._server is None:
+ # Resolving the connection talks to the supervisor, so it takes
the blocking-call lock.
+ self._server = await self.run_blocking(self._get_server)
+ return self._server
+
+ def _translate_error(
+ self, error: Exception, operation: _Operation, *, during_tool_call:
bool
+ ) -> Exception | None:
+ """Return the provider exception for a gateway failure, or ``None`` to
re-raise ``error`` as is."""
+ service = self._service_name
+ mcp_error = _find(error, McpError)
+ # When a request fails once the session is open, the call only sees an
error the MCP
+ # client makes up; what the request ran into is in the operation's
failures.
+ session_ended = _is_session_ended(error)
+ saw_no_answer = session_ended or (mcp_error is not None and
_is_client_error(mcp_error))
+ failure = operation.failures[0] if saw_no_answer and
operation.failures else None
+ token_error = _find(error, _TokenUnavailableError)
+ if token_error is None and isinstance(failure, _TokenUnavailableError):
+ token_error = failure
+ if token_error is not None:
+ return DatabricksUnityMCPError(
+ f"Could not get a token from connection
{self._databricks_conn_id!r} to call MCP Service "
+ f"{service!r}: {token_error}"
+ )
+ status, retry_after = operation.http_error or (None, None)
+ if (mcp_error is not None and mcp_error.error.code ==
_NOT_AUTHORIZED_CODE) or status in (401, 403):
+ return DatabricksUnityMCPAccessDeniedError(
+ f"Unity Gateway denied access to MCP Service {service!r}"
+ f"{f' (HTTP {status})' if status else ''}. Either no service
has that name, or the "
+ f"identity of connection {self._databricks_conn_id!r} lacks
EXECUTE on the service or "
+ "USE CATALOG and USE SCHEMA on its catalog and schema, or its
credentials are not valid.",
+ http_status_code=status,
+ )
+ if status == 429:
+ hint = f" Retry after {retry_after:g} seconds." if retry_after is
not None else ""
+ return DatabricksUnityMCPThrottledError(
+ f"Unity Gateway rate-limited calls to MCP Service {service!r}
(HTTP 429).{hint}",
+ http_status_code=status,
+ retry_after=retry_after,
+ )
+ ambiguous = (
+ " The tool call may or may not have run, so it was not retried."
if during_tool_call else ""
+ )
+ # The MCP client reports its own timeouts as MCP errors, so they are
told apart from the
+ # server's answers before those are.
+ if _is_request_timeout(mcp_error):
+ return DatabricksUnityMCPTransportError(
+ f"Timed out waiting for Unity Gateway to answer MCP Service
{service!r}.{ambiguous}"
+ )
+ transport_error = _find(error, httpx2.TransportError)
+ if transport_error is None and isinstance(failure,
httpx2.TransportError):
+ transport_error = failure
+ if transport_error is not None:
+ return DatabricksUnityMCPTransportError(
+ f"Could not reach Unity Gateway for MCP Service {service!r}: "
+ f"{_describe(transport_error)}.{ambiguous}"
+ )
+ if during_tool_call and _is_handshake_timeout(error):
+ return DatabricksUnityMCPTransportError(
+ f"Timed out waiting for Unity Gateway to answer MCP Service
{service!r}.{ambiguous}"
+ )
+ if session_ended:
+ return DatabricksUnityMCPTransportError(
+ f"The connection to Unity Gateway for MCP Service {service!r}
ended.{ambiguous}"
+ )
+ if _is_handshake_timeout(error):
+ return DatabricksUnityMCPTransportError(
+ f"Unity Gateway did not complete the MCP handshake for MCP
Service {service!r} within "
+ f"{self._init_timeout:g} seconds. Raise init_timeout if the
service is slow to answer."
+ )
+ if mcp_error is not None:
+ if not _is_client_error(mcp_error):
+ # The server's own answer: for a tool call, it reaches the
model as a retry.
+ if during_tool_call:
+ return None
+ return DatabricksUnityMCPError(
+ f"MCP Service {service!r} returned an error:
{mcp_error.error.message} "
+ f"(code {mcp_error.error.code}).",
+ http_status_code=status,
+ )
+ what = f"HTTP {status}" if status else f"an error without an
answer ({mcp_error.error.message})"
+ return DatabricksUnityMCPError(
+ f"Unity Gateway returned {what} for MCP Service
{service!r}.{ambiguous}",
+ http_status_code=status,
+ )
+ return None
+
+ @contextlib.asynccontextmanager
+ async def _gateway_operation(self, *, during_tool_call: bool = False) ->
AsyncIterator[None]:
+ operation = _Operation(overlapped=bool(self._operations))
+ for other in self._operations:
+ other.overlapped = True
+ self._operations.add(operation)
+ try:
+ yield
+ except Exception as e:
+ translated = self._translate_error(e, operation,
during_tool_call=during_tool_call)
+ if translated is None:
+ raise
+ self._failure_reported = True
+ raise translated from e
+ finally:
+ self._operations.discard(operation)
+
+ async def __aenter__(self) -> DatabricksUnityMCPToolset:
+ self._failure_reported = False
+ async with self._gateway_operation():
+ await (await self._resolve_server()).__aenter__()
+ return self
+
+ async def __aexit__(self, *args: Any) -> bool | None:
+ if self._server is None:
+ return None
+ # pydantic-ai's CombinedToolset closes its toolsets without passing on
the exception
+ # that ends the run, so a reported failure is remembered rather than
read from args.
+ if args[0] is None and not self._failure_reported:
+ return await self._server.__aexit__(*args)
+ try:
+ return await self._server.__aexit__(*args)
+ except Exception:
+ # Closing the session re-raises the failure of a request that
already failed, such
+ # as the read timeout of a tool call, which would hide the error
explaining it.
+ log.warning(
+ "Error while closing the session of MCP Service %r",
self._service_name, exc_info=True
+ )
+ return None
+
+ async def get_tools(self, ctx: RunContext[Any]) -> dict[str,
ToolsetTool[Any]]:
+ async with self._gateway_operation():
+ return await (await self._resolve_server()).get_tools(ctx)
+
+ async def execute_tool(
+ self,
+ name: str,
+ tool_args: dict[str, Any],
+ *,
+ ctx: RunContext[Any],
+ tool: ToolsetTool[Any],
+ ) -> Any:
+ async with self._gateway_operation(during_tool_call=True):
+ return await (await self._resolve_server()).call_tool(name,
tool_args, ctx, tool)
+
+ def airflow_tools(self) -> list[AirflowTool]:
+ """
+ Not supported: use the agent framework's own MCP client instead.
+
+ An MCP session belongs to the event loop that opened it, and the
framework-neutral
+ tools run outside the Pydantic AI run that manages it, so every call
would reconnect.
+ """
+ raise NotImplementedError(
+ "DatabricksUnityMCPToolset works in Pydantic AI agents and through
the LangChain bridge, "
+ "not through the framework-neutral tools."
+ )
+
+
+def _build_client_factory(operations: set[_Operation], proxy: str | None) ->
Any:
+ """Return an MCP HTTP client factory for gateway requests, sent through
``proxy`` if one is set."""
+
+ def factory(
+ headers: dict[str, str] | None = None,
+ timeout: httpx2.Timeout | None = None,
+ auth: httpx2.Auth | None = None,
+ **kwargs: Any,
+ ) -> httpx2.AsyncClient:
+ # The MCP client's defaults: a response stream may stay open, so reads
wait longer.
+ timeout = timeout or httpx2.Timeout(30.0, read=300.0)
+ if proxy:
+ kwargs["proxy"] = proxy
+ return _GatewayClient(headers=headers, timeout=timeout, auth=auth,
operations=operations, **kwargs)
+
+ return factory
diff --git
a/providers/databricks/src/airflow/providers/databricks/version_compat.py
b/providers/databricks/src/airflow/providers/databricks/version_compat.py
index 0956edd2111..402b4f953ae 100644
--- a/providers/databricks/src/airflow/providers/databricks/version_compat.py
+++ b/providers/databricks/src/airflow/providers/databricks/version_compat.py
@@ -33,7 +33,9 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
+AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
+ "AIRFLOW_V_3_2_PLUS",
]
diff --git
a/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py
b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py
new file mode 100644
index 00000000000..02a5e5235db
--- /dev/null
+++
b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py
@@ -0,0 +1,69 @@
+# 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.
+"""
+Example Dag giving an agent the tools of a Unity Gateway MCP Service.
+
+Needs a Databricks connection whose identity has EXECUTE on the MCP Service
and USE CATALOG and
+USE SCHEMA on its catalog and schema, an LLM connection for the agent, and
these environment
+variables:
+
+- ``UNITY_MCP_SERVICE``: the service's three-level name,
``catalog.schema.service``.
+- ``UNITY_MCP_READ_ONLY_TOOL``: the name of one of its tools that only reads,
which the agent calls
+ once.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.compat.sdk import DAG
+from airflow.providers.databricks.toolsets.unity_mcp import
DatabricksUnityMCPToolset
+
+DAG_ID = "example_databricks_unity_mcp"
+UNITY_MCP_SERVICE = os.environ.get("UNITY_MCP_SERVICE",
"main.default.my_mcp_service")
+UNITY_MCP_READ_ONLY_TOOL = os.environ.get("UNITY_MCP_READ_ONLY_TOOL",
"my_read_only_tool")
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule=None,
+ start_date=datetime(2021, 1, 1),
+ tags=["example"],
+ catchup=False,
+) as dag:
+ # [START howto_toolset_databricks_unity_mcp]
+ ask_unity_mcp_service = AgentOperator(
+ task_id="ask_unity_mcp_service",
+ prompt=(
+ "List the names of the tools you have. Then call the "
+ f"{UNITY_MCP_READ_ONLY_TOOL} tool once and summarize what it
returned."
+ ),
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ DatabricksUnityMCPToolset(UNITY_MCP_SERVICE,
databricks_conn_id="databricks_default"),
+ ],
+ )
+ # [END howto_toolset_databricks_unity_mcp]
+
+ from tests_common.test_utils.watcher import watcher
+
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+test_run = get_test_run(dag)
diff --git a/providers/databricks/tests/unit/databricks/toolsets/__init__.py
b/providers/databricks/tests/unit/databricks/toolsets/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/providers/databricks/tests/unit/databricks/toolsets/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git
a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py
b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py
new file mode 100644
index 00000000000..4d75f195986
--- /dev/null
+++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py
@@ -0,0 +1,902 @@
+# 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.
+from __future__ import annotations
+
+import ast
+import asyncio
+import functools
+import importlib
+import inspect
+import json
+import re
+import sys
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from unittest import mock
+
+import pytest
+
+pytest.importorskip("fastmcp")
+pytest.importorskip("airflow.providers.common.ai")
+
+import httpx2
+from fastmcp.exceptions import McpError
+from mcp_types import jsonrpc
+from pydantic_ai import ModelRetry, RunContext
+from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset
+from pydantic_ai.toolsets import CombinedToolset
+
+from airflow.models import Connection
+from airflow.providers.common.ai.utils import toolset_base
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.exceptions import (
+ DatabricksUnityMCPAccessDeniedError,
+ DatabricksUnityMCPError,
+ DatabricksUnityMCPThrottledError,
+ DatabricksUnityMCPTransportError,
+)
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+from airflow.providers.databricks.toolsets import unity_mcp
+from airflow.providers.databricks.toolsets.unity_mcp import
DatabricksUnityMCPToolset
+
+CONN_ID = "databricks_unity"
+SERVICE = "main.tools.genie"
+
+# What Unity Gateway answers, at HTTP 403, for a service the caller may not
invoke or that does not exist.
+NOT_AUTHORIZED = {"code": -32007, "message": "Not authorized to invoke MCP
service."}
+INVALID_PARAMS = {"code": -32602, "message": "Invalid params: 'text' is
required"}
+
+
+class _StubGateway(BaseHTTPRequestHandler):
+ """
+ A Unity Gateway stand-in that serves one MCP Service with a single
``echo`` tool.
+
+ A call of ``echo`` whose text is ``"reply <status>"`` gets that HTTP
status with an empty body,
+ and ``"reply <status> <error>"`` gets it with the JSON-RPC error named by
``<error>``. With
+ ``"drop"`` the connection closes without a response, and with ``"drop
body"`` it closes
+ partway through the response body.
+ """
+
+ protocol_version = "HTTP/1.1"
+ errors = {"not_authorized": NOT_AUTHORIZED, "invalid_params":
INVALID_PARAMS}
+ # Per-test behaviour, set by the ``gateway`` fixture: a reply for every
request to a service,
+ # a status for notifications only, and a barrier tool calls wait at, so
that they overlap.
+ service_reply: dict[str, tuple[int, dict[str, str], str | None]] = {}
+ notification_status: int | None = None
+ tool_call_barrier: threading.Barrier | None = None
+ # Seconds to hold the answer to each MCP method; ``release`` ends any wait
early on teardown.
+ delays: dict[str, float] = {}
+ release = threading.Event()
+ requests: list[tuple[str, str | None, str | None]] = []
+
+ def do_POST(self):
+ body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))
or b"{}")
+ method = body.get("method")
+ self.requests.append((self.path, self.headers.get("Authorization"),
method))
+ service = self.path.rsplit("/", 1)[-1]
+ if service in self.service_reply:
+ status, headers, error = self.service_reply[service]
+ return self._reply_error(status, body.get("id"), error, headers)
+ if "id" not in body:
+ return self._reply(self.notification_status or 202)
+ self.release.wait(self.delays.get(method, 0))
+ if method == "initialize":
+ result = {
+ "protocolVersion": body["params"]["protocolVersion"],
+ "capabilities": {"tools": {}},
+ "serverInfo": {"name": "stub", "version": "1"},
+ }
+ elif method == "tools/list":
+ result = {
+ "tools": [
+ {
+ "name": "echo",
+ "description": "Echo the text back.",
+ "inputSchema": {"type": "object", "properties":
{"text": {"type": "string"}}},
+ }
+ ]
+ }
+ elif method == "tools/call":
+ if self.tool_call_barrier:
+ self.tool_call_barrier.wait()
+ if "text" not in body["params"]["arguments"]:
+ return self._reply(200, {"jsonrpc": "2.0", "id": body["id"],
"error": INVALID_PARAMS})
+ text = body["params"]["arguments"]["text"]
+ if text.startswith("drop"):
+ self.close_connection = True
+ if text == "drop body":
+ self.send_response(200)
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", "1000")
+ self.end_headers()
+ self.wfile.write(b'{"jsonrpc": "2.0"')
+ return None
+ if text.startswith("reply "):
+ status, _, error = text.removeprefix("reply ").partition(" ")
+ return self._reply_error(int(status), body["id"], error or
None)
+ result = {"content": [{"type": "text", "text": f"echo: {text}"}],
"isError": False}
+ else:
+ result = {}
+ self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "result":
result})
+
+ def do_GET(self):
+ self._reply(405)
+
+ def do_DELETE(self):
+ self._reply(200)
+
+ def _reply_error(self, status, request_id, error, headers=None):
+ payload = {"jsonrpc": "2.0", "id": request_id, "error":
self.errors[error]} if error else None
+ self._reply(status, payload, headers)
+
+ def _reply(self, status, payload=None, headers=None):
+ data = json.dumps(payload).encode() if payload is not None else b""
+ self.send_response(status)
+ for name, value in (headers or {}).items():
+ self.send_header(name, value)
+ if payload is not None:
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", str(len(data)))
+ self.end_headers()
+ self.wfile.write(data)
+
+ def log_message(self, *args):
+ pass
+
+
[email protected]
+def gateway():
+ _StubGateway.service_reply = {}
+ _StubGateway.notification_status = None
+ _StubGateway.tool_call_barrier = None
+ _StubGateway.delays = {}
+ _StubGateway.release = threading.Event()
+ _StubGateway.requests = []
+ server = ThreadingHTTPServer(("127.0.0.1", 0), _StubGateway)
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ yield server
+ _StubGateway.release.set()
+ server.shutdown()
+ server.server_close()
+
+
[email protected]
+def gateway_conn(gateway, create_connection_without_db):
+ create_connection_without_db(
+ Connection(
+ conn_id=CONN_ID,
+ conn_type="databricks",
+ host="127.0.0.1",
+ port=gateway.server_address[1],
+ schema="http",
+ password="dapi-token",
+ )
+ )
+
+
+def _run(coro):
+ return asyncio.run(coro)
+
+
+def _ctx():
+ ctx = mock.MagicMock(spec=RunContext)
+ ctx.max_retries = 1
+ return ctx
+
+
+async def _list_and_call(toolset, text="hi"):
+ async with toolset:
+ tools = await toolset.get_tools(_ctx())
+ result = await toolset.execute_tool("echo", {"text": text},
ctx=_ctx(), tool=tools["echo"])
+ return tools, result
+
+
+async def _list_and_call_in_an_agent_run(toolset, text="hi"):
+ """Call ``echo`` the way an agent run does, through a CombinedToolset,
which closes ``toolset``."""
+ combined = CombinedToolset([toolset])
+ async with combined:
+ tools = await combined.get_tools(_ctx())
+ return await combined.call_tool("echo", {"text": text}, _ctx(),
tools["echo"])
+
+
+async def _call_concurrently(*texts):
+ """Call ``echo`` once per text, all at the same time, and return each
call's result or exception."""
+ toolset = DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID)
+ async with toolset:
+ tools = await toolset.get_tools(_ctx())
+ return await asyncio.gather(
+ *(toolset.execute_tool("echo", {"text": t}, ctx=_ctx(),
tool=tools["echo"]) for t in texts),
+ return_exceptions=True,
+ )
+
+
+class TestServiceName:
+ @pytest.mark.parametrize(
+ "name",
+ [
+ "main.tools",
+ "main.tools.genie.extra",
+ "main.tools.genie/../../api/2.0/jobs",
+ "main.tools.genie?x=1",
+ "main.tools.genie#frag",
+ "main.tools.ge nie",
+ "main.tools.my-genie",
+ "main..genie",
+ "evil.example.com/main.tools.genie",
+ "",
+ ],
+ )
+ def test_rejects_names_that_are_not_three_safe_parts(self, name):
+ with pytest.raises(ValueError, match="catalog.schema.service"):
+ DatabricksUnityMCPToolset(name)
+
+ def test_templated_name_is_validated_once_rendered(self, gateway_conn):
+ toolset = DatabricksUnityMCPToolset("{{ params.service }}",
databricks_conn_id=CONN_ID)
+ toolset._service_name = "main.tools/../../api"
+
+ with pytest.raises(ValueError, match="catalog.schema.service"):
+ toolset._get_server()
+
+ def test_id_names_connection_and_service(self):
+ assert (
+ DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID).id
+ == f"databricks-unity-mcp-{CONN_ID}-{SERVICE}"
+ )
+
+ def test_connection_and_service_are_agent_template_fields(self):
+ assert set(DatabricksUnityMCPToolset.agent_template_fields) == {
+ "_databricks_conn_id",
+ "_service_name",
+ }
+
+
+class TestServiceUrl:
+ @pytest.mark.parametrize(
+ ("host", "extra", "expected"),
+ [
+ ("xx.cloud.databricks.com", None,
"https://xx.cloud.databricks.com";),
+ ("https://xx.cloud.databricks.com/";, None,
"https://xx.cloud.databricks.com";),
+ (None, {"host": "https://yy.azuredatabricks.net"},
"https://yy.azuredatabricks.net";),
+ ],
+ )
+ def test_url_is_built_on_the_connection_host(self,
create_connection_without_db, host, extra, expected):
+ create_connection_without_db(
+ Connection(conn_id=CONN_ID, conn_type="databricks", host=host,
password="t", extra=extra)
+ )
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ assert (
+ toolset.get_service_url(DatabricksHook(CONN_ID))
+ == f"{expected}/ai-gateway/mcp-services/{SERVICE}"
+ )
+
+ def test_connection_without_host_is_rejected(self,
create_connection_without_db):
+ create_connection_without_db(Connection(conn_id=CONN_ID,
conn_type="databricks", password="t"))
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ with pytest.raises(ValueError, match="no workspace host"):
+ toolset.get_service_url(DatabricksHook(CONN_ID))
+
+ @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "[::1]"])
+ def test_plain_http_is_allowed_to_loopback_only(self,
create_connection_without_db, host):
+ create_connection_without_db(
+ Connection(conn_id=CONN_ID, conn_type="databricks", host=host,
schema="http", password="t")
+ )
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ assert
toolset.get_service_url(DatabricksHook(CONN_ID)).startswith(f"http://{host}/";)
+
+ def test_plain_http_to_a_workspace_is_rejected(self,
create_connection_without_db):
+ create_connection_without_db(
+ Connection(
+ conn_id=CONN_ID,
+ conn_type="databricks",
+ host="xx.cloud.databricks.com",
+ schema="http",
+ password="t",
+ )
+ )
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ with pytest.raises(ValueError, match="only sent over HTTPS"):
+ toolset.get_service_url(DatabricksHook(CONN_ID))
+
+ def test_requests_go_through_the_connection_proxy(self, gateway,
create_connection_without_db):
+ # The stub gateway doubles as the proxy: a proxied request names the
full URL in its path.
+ create_connection_without_db(
+ Connection(
+ conn_id=CONN_ID,
+ conn_type="databricks",
+ host="127.0.0.2",
+ port=1,
+ schema="http",
+ password="dapi-token",
+ extra={"proxies": {"http":
f"http://127.0.0.1:{gateway.server_address[1]}"}},
+ )
+ )
+
+ _, result = _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert result == "echo: hi"
+ assert {path for path, _, _ in _StubGateway.requests} == {
+ f"http://127.0.0.2:1/ai-gateway/mcp-services/{SERVICE}";
+ }
+
+
+class TestAuthentication:
+ def test_tools_are_discovered_and_called_with_the_connection_token(self,
gateway, gateway_conn):
+ tools, result = _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert list(tools) == ["echo"]
+ assert result == "echo: hi"
+ assert _StubGateway.requests
+ assert {path for path, _, _ in _StubGateway.requests} ==
{f"/ai-gateway/mcp-services/{SERVICE}"}
+ assert {auth for _, auth, _ in _StubGateway.requests} == {"Bearer
dapi-token"}
+
+ @mock.patch.object(DatabricksHook, "_get_token", autospec=True)
+ def test_token_is_fetched_for_every_request_so_it_can_refresh(self,
get_token, gateway, gateway_conn):
+ tokens = (f"oauth-{i}" for i in range(1000))
+ get_token.side_effect = lambda *_, **__: next(tokens)
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ sent = [auth for _, auth, _ in _StubGateway.requests]
+ # One token is fetched up front to fail fast, then one per request.
+ assert sent == [f"Bearer oauth-{i}" for i in range(1, len(sent) + 1)]
+
+ @mock.patch.object(unity_mcp, "AIRFLOW_V_3_2_PLUS", True)
+ def test_requests_do_not_wait_for_other_toolsets_blocking_calls(self,
gateway, gateway_conn):
+ async def call_while_another_toolset_holds_the_lock():
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+ async with toolset:
+ tools = await toolset.get_tools(_ctx())
+ with toolset_base._blocking_call_lock:
+ return await asyncio.wait_for(
+ toolset.execute_tool("echo", {"text": "hi"},
ctx=_ctx(), tool=tools["echo"]),
+ timeout=10,
+ )
+
+ assert _run(call_while_another_toolset_holds_the_lock()) == "echo: hi"
+
+ @mock.patch.object(unity_mcp, "AIRFLOW_V_3_2_PLUS", False)
+ @mock.patch.object(DatabricksHook, "_get_token", autospec=True)
+ def
test_before_airflow_3_2_tokens_are_fetched_under_the_blocking_call_lock(
+ self, get_token, gateway, gateway_conn
+ ):
+ # Reading the connection's extras messages the supervisor, whose
socket has no lock of its own.
+ held = []
+ get_token.side_effect = lambda *_, **__:
held.append(toolset_base._blocking_call_lock.locked()) or "t"
+
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert len(held) > 1
+ assert all(held)
+
+ @mock.patch.object(DatabricksHook, "_get_token", autospec=True)
+ @mock.patch("airflow.providers.databricks.toolsets.unity_mcp.mask_secret",
autospec=True)
+ def test_each_new_token_is_masked_once_under_the_blocking_call_lock(
+ self, mask_secret, get_token, gateway, gateway_conn
+ ):
+ token = "oauth-1"
+ get_token.side_effect = lambda *_, **__: token
+ # Before Airflow 3.2, the supervisor socket mask_secret writes to has
no lock of its own.
+ mask_secret.side_effect = lambda _: (
+ toolset_base._blocking_call_lock.locked()
+ or pytest.fail("mask_secret was called without the blocking-call
lock")
+ )
+
+ async def refresh_between_calls():
+ nonlocal token
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+ async with toolset:
+ tools = await toolset.get_tools(_ctx())
+ await toolset.execute_tool("echo", {"text": "a"}, ctx=_ctx(),
tool=tools["echo"])
+ token = "oauth-2"
+ await toolset.execute_tool("echo", {"text": "b"}, ctx=_ctx(),
tool=tools["echo"])
+
+ _run(refresh_between_calls())
+
+ assert mask_secret.call_args_list == [mock.call("oauth-1"),
mock.call("oauth-2")]
+
+ def test_username_and_password_connection_is_rejected(self,
create_connection_without_db):
+ create_connection_without_db(
+ Connection(conn_id=CONN_ID, conn_type="databricks", host="h",
login="user", password="pass")
+ )
+
+ with pytest.raises(ValueError, match="sends a bearer token"):
+ DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)._get_server()
+
+ @mock.patch.object(DatabricksHook, "_get_token", autospec=True)
+ def test_token_failure_during_the_run_is_reported(self, get_token,
gateway, gateway_conn):
+ # The up-front check and the first request get a token; the next
request's fetch fails.
+ get_token.side_effect = ["dapi-token", "dapi-token", OSError("token
endpoint unreachable")]
+
+ with pytest.raises(DatabricksUnityMCPError, match="Could not get a
token") as err:
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert "token endpoint unreachable" in str(err.value)
+
+ @mock.patch.object(DatabricksHook, "_get_token", autospec=True)
+ def test_token_failure_before_connecting_is_reported(self, get_token,
gateway, gateway_conn):
+ # A hung token endpoint must not be mistaken for a slow MCP handshake.
+ get_token.side_effect = TimeoutError("token endpoint timed out")
+
+ with pytest.raises(DatabricksUnityMCPError, match="Could not get a
token.*token endpoint timed out"):
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert not _StubGateway.requests
+
+ @mock.patch.object(DatabricksHook, "_get_token", autospec=True)
+ def test_token_failure_during_a_tool_call_is_reported(self, get_token,
gateway, gateway_conn):
+ failing = False
+
+ def get_token_unless_failing(*_, **__):
+ if failing:
+ raise OSError("token endpoint unreachable")
+ return "dapi-token"
+
+ get_token.side_effect = get_token_unless_failing
+
+ async def call_after_the_token_endpoint_fails():
+ nonlocal failing
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+ async with toolset:
+ tools = await toolset.get_tools(_ctx())
+ failing = True
+ await toolset.execute_tool("echo", {"text": "hi"}, ctx=_ctx(),
tool=tools["echo"])
+
+ with pytest.raises(
+ DatabricksUnityMCPError, match="Could not get a token.*token
endpoint unreachable"
+ ):
+ _run(call_after_the_token_endpoint_fails())
+
+
+class TestIsJsonRpcRequest:
+ @pytest.mark.parametrize(
+ ("method", "body", "expected"),
+ [
+ ("POST", {"jsonrpc": "2.0", "id": 1, "method": "tools/call"},
True),
+ ("POST", {"jsonrpc": "2.0", "method":
"notifications/initialized"}, False),
+ ("GET", None, False),
+ ],
+ )
+ def test_only_requests_that_expect_an_answer_count(self, method, body,
expected):
+ request = httpx2.Request(method, "https://h/mcp";, json=body)
+
+ assert unity_mcp._is_jsonrpc_request(request) is expected
+
+
+class TestGatewayClient:
+ @pytest.mark.parametrize(
+ ("method", "body", "recorded"),
+ [
+ ("POST", {"jsonrpc": "2.0", "id": 1, "method": "tools/call"},
True),
+ ("POST", {"jsonrpc": "2.0", "method":
"notifications/initialized"}, False),
+ ("GET", None, False),
+ ],
+ )
+ def test_records_failures_of_json_rpc_requests_only(self, method, body,
recorded):
+ operation = unity_mcp._Operation()
+
+ def fail(request):
+ if request.headers.get("X-Fail") == "connect":
+ raise httpx2.ConnectError("refused", request=request)
+ return httpx2.Response(503)
+
+ async def send_both():
+ async with unity_mcp._GatewayClient(
+ transport=httpx2.MockTransport(fail), operations={operation}
+ ) as client:
+ await client.request(method, "https://h/mcp";, json=body)
+ with pytest.raises(httpx2.ConnectError):
+ await client.request(method, "https://h/mcp";, json=body,
headers={"X-Fail": "connect"})
+
+ _run(send_both())
+
+ assert operation.errors == ([(503, None)] if recorded else [])
+ assert len(operation.failures) == (1 if recorded else 0)
+
+
+class TestErrors:
+ @pytest.mark.parametrize(
+ ("status", "headers", "error", "expected", "retry_after"),
+ [
+ (401, {}, None, DatabricksUnityMCPAccessDeniedError, None),
+ (403, {}, None, DatabricksUnityMCPAccessDeniedError, None),
+ (403, {}, "not_authorized", DatabricksUnityMCPAccessDeniedError,
None),
+ (404, {}, None, DatabricksUnityMCPError, None),
+ (429, {"Retry-After": "12"}, None,
DatabricksUnityMCPThrottledError, 12.0),
+ (429, {"Retry-After": "-5"}, None,
DatabricksUnityMCPThrottledError, None),
+ (
+ 429,
+ {"Retry-After": "Wed, 21 Oct 2026 07:28:00 GMT"},
+ None,
+ DatabricksUnityMCPThrottledError,
+ None,
+ ),
+ ],
+ )
+ def test_gateway_errors_while_connecting_are_translated(
+ self, gateway, gateway_conn, status, headers, error, expected,
retry_after
+ ):
+ _StubGateway.service_reply[SERVICE] = (status, headers, error)
+
+ with pytest.raises(expected) as err:
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert type(err.value) is expected
+ assert err.value.http_status_code == status
+ # No tool call was sent, so nothing may have run.
+ assert "may or may not have run" not in str(err.value)
+ if expected is DatabricksUnityMCPThrottledError:
+ assert err.value.retry_after == retry_after
+
+ def test_missing_service_is_reported_as_access_denied(self, gateway,
gateway_conn):
+ _StubGateway.service_reply[SERVICE] = (403, {}, "not_authorized")
+
+ with pytest.raises(DatabricksUnityMCPAccessDeniedError, match="no
service has that name"):
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ def test_server_error_during_a_tool_call_says_the_call_may_have_run(self,
gateway, gateway_conn):
+ with pytest.raises(DatabricksUnityMCPError, match="may or may not have
run") as err:
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID), "reply 503"))
+
+ assert type(err.value) is DatabricksUnityMCPError
+ assert err.value.http_status_code == 503
+ # The failed call was sent once and not retried.
+ assert [m for _, _, m in _StubGateway.requests].count("tools/call") ==
1
+
+ def test_denied_tool_call_is_not_left_to_the_model(self, gateway,
gateway_conn):
+ with pytest.raises(DatabricksUnityMCPAccessDeniedError):
+ _run(
+ _list_and_call(
+ DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID), "reply 403 not_authorized"
+ )
+ )
+
+ @pytest.mark.parametrize(
+ "text", ["reply 200 invalid_params", "reply 400 invalid_params",
"reply 500 invalid_params"]
+ )
+ def test_json_rpc_error_from_the_server_is_left_to_the_model(self,
gateway, gateway_conn, text):
+ # A server may answer bad arguments with a JSON-RPC error at an HTTP
error status; it is
+ # still the server's answer, so the model can correct the call.
+ with pytest.raises(ModelRetry, match="'text' is required"):
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID), text))
+
+ def test_concurrent_calls_are_each_judged_by_their_own_failure(self,
gateway, gateway_conn):
+ # Calls run in parallel, so another call's response must not change
how a failure is
+ # reported: a server error stays a non-retried failure however the
other call ends.
+ _StubGateway.tool_call_barrier = threading.Barrier(2, timeout=10)
+ for _ in range(10):
+ failed, succeeded = _run(_call_concurrently("reply 503", "hi"))
+ assert type(failed) is DatabricksUnityMCPError
+ assert "may or may not have run" in str(failed)
+ assert succeeded == "echo: hi"
+
+ failed, throttled = _run(_call_concurrently("reply 503", "reply
429"))
+ # Neither status can be attributed to its call, so neither is
reported as throttling.
+ assert type(failed) is DatabricksUnityMCPError
+ assert type(throttled) is DatabricksUnityMCPError
+ assert failed.http_status_code is None
+
+ invalid, failed = _run(_call_concurrently("reply 400
invalid_params", "reply 503"))
+ assert isinstance(invalid, ModelRetry)
+ assert type(failed) is DatabricksUnityMCPError
+
+ # The dropped call ends the session; the other call either got its
answer first or
+ # finds the session gone, which is the same transport failure,
never a raw error.
+ dropped, other = _run(_call_concurrently("drop", "hi"))
+ assert type(dropped) is DatabricksUnityMCPTransportError
+ assert other == "echo: hi" or type(other) is
DatabricksUnityMCPTransportError, repr(other)
+
+ # The status is lost, but the gateway's JSON-RPC code still marks
the denial.
+ denied, succeeded = _run(_call_concurrently("reply 403
not_authorized", "hi"))
+ assert type(denied) is DatabricksUnityMCPAccessDeniedError
+ assert succeeded == "echo: hi"
+
+ def
test_tool_errors_after_a_tolerated_gateway_error_are_left_to_the_agent(self,
gateway, gateway_conn):
+ # The MCP client carries on when the gateway rejects a notification; a
later tool error the
+ # model can fix must still reach pydantic-ai as a retry, not as that
earlier gateway error.
+ _StubGateway.notification_status = 400
+
+ async def call_with_missing_argument():
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+ async with toolset:
+ tools = await toolset.get_tools(_ctx())
+ await toolset.execute_tool("echo", {}, ctx=_ctx(),
tool=tools["echo"])
+
+ with pytest.raises(ModelRetry, match="'text' is required"):
+ _run(call_with_missing_argument())
+
+ @pytest.mark.parametrize("call", [_list_and_call,
_list_and_call_in_an_agent_run])
+ @pytest.mark.parametrize("text", ["drop", "drop body"])
+ def test_connection_lost_during_a_tool_call_is_a_transport_error(self,
gateway, gateway_conn, text, call):
+ with pytest.raises(DatabricksUnityMCPTransportError, match="may or may
not have run"):
+ _run(call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID), text))
+
+ assert [m for _, _, m in _StubGateway.requests].count("tools/call") ==
1
+
+ @pytest.mark.parametrize(
+ ("mcp_error", "expected"),
+ [
+ # The server answered, so the overlapping call's failure says
nothing about this call.
+ (McpError(jsonrpc.INVALID_PARAMS, "Invalid params"), None),
+ (McpError(jsonrpc.CONNECTION_CLOSED, "Connection closed"),
DatabricksUnityMCPTransportError),
+ (RuntimeError("Client is not connected."),
DatabricksUnityMCPTransportError),
+ ],
+ )
+ def
test_recorded_request_failure_only_explains_a_call_without_an_answer(self,
mcp_error, expected):
+ operation = unity_mcp._Operation(overlapped=True,
failures=[httpx2.ConnectError("refused")])
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ translated = toolset._translate_error(mcp_error, operation,
during_tool_call=True)
+
+ if expected is None:
+ assert translated is None
+ else:
+ assert type(translated) is expected
+
+ @staticmethod
+ def _with_context(error, context):
+ error.__context__ = context
+ return error
+
+ @pytest.mark.parametrize(
+ ("error", "during_tool_call", "expected", "match"),
+ [
+ pytest.param(
+ ExceptionGroup("task group", [httpx2.ConnectError("refused")]),
+ False,
+ DatabricksUnityMCPTransportError,
+ "refused",
+ id="transport error inside an exception group",
+ ),
+ pytest.param(
+ _with_context(RuntimeError("wrapped"),
httpx2.ConnectError("refused")),
+ False,
+ DatabricksUnityMCPTransportError,
+ "refused",
+ id="transport error in the context",
+ ),
+ pytest.param(
+ httpx2.ReadTimeout(""),
+ False,
+ DatabricksUnityMCPTransportError,
+ "ReadTimeout",
+ id="transport error without a message is named",
+ ),
+ pytest.param(
+ TimeoutError(),
+ False,
+ DatabricksUnityMCPTransportError,
+ "Raise init_timeout",
+ id="bare timeout while connecting",
+ ),
+ pytest.param(
+ RuntimeError("Failed to initialize server session"),
+ True,
+ DatabricksUnityMCPTransportError,
+ "Timed out waiting.*may or may not have run",
+ id="handshake timeout during a tool call",
+ ),
+ pytest.param(
+ McpError(-32001, "The warehouse is still starting"),
+ True,
+ None,
+ None,
+ id="server's own -32001 answer reaches the model",
+ ),
+ ],
+ )
+ def test_error_is_translated_from_its_chain(self, error, during_tool_call,
expected, match):
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ translated = toolset._translate_error(
+ error, unity_mcp._Operation(), during_tool_call=during_tool_call
+ )
+
+ if expected is None:
+ assert translated is None
+ else:
+ assert type(translated) is expected
+ assert re.search(match, str(translated))
+
+ def test_status_of_the_last_failed_request_is_used(self):
+ operation = unity_mcp._Operation(errors=[(503, None), (429, 5.0)])
+
+ assert operation.http_error == (429, 5.0)
+
+ def test_call_that_finds_the_session_ended_is_a_transport_error(self):
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+
+ translated = toolset._translate_error(
+ RuntimeError("Client is not connected."), unity_mcp._Operation(),
during_tool_call=True
+ )
+
+ assert type(translated) is DatabricksUnityMCPTransportError
+ assert "may or may not have run" in str(translated)
+
+ def test_unreachable_gateway_is_a_transport_error(self,
create_connection_without_db):
+ create_connection_without_db(
+ Connection(
+ conn_id=CONN_ID, conn_type="databricks", host="127.0.0.1",
port=1, schema="http", password="t"
+ )
+ )
+
+ with pytest.raises(DatabricksUnityMCPTransportError, match="Could not
reach Unity Gateway"):
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+
+# The MCP client gives up on a request after pydantic-ai's read_timeout, 5
minutes by default.
[email protected](
+ unity_mcp, "PydanticAIMCPToolset", functools.partial(PydanticAIMCPToolset,
read_timeout=0.5)
+)
+class TestRequestTimeout:
+ @pytest.mark.parametrize("call", [_list_and_call,
_list_and_call_in_an_agent_run])
+ def test_tool_call_without_an_answer_is_not_left_to_the_model(self,
gateway, gateway_conn, call):
+ _StubGateway.delays = {"tools/call": 2}
+
+ with pytest.raises(DatabricksUnityMCPTransportError, match="may or may
not have run"):
+ _run(call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert [m for _, _, m in _StubGateway.requests].count("tools/call") ==
1
+
+ def test_listing_tools_without_an_answer_is_a_transport_error(self,
gateway, gateway_conn):
+ _StubGateway.delays = {"tools/list": 2}
+
+ with pytest.raises(DatabricksUnityMCPTransportError, match="Timed out
waiting") as err:
+ _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)))
+
+ assert "init_timeout" not in str(err.value)
+
+
+class TestClientErrorMessages:
+ """
+ The errors the MCP client makes up are told from the server's own by their
text.
+
+ A failure here means an MCP SDK or fastmcp update changed that text:
update the constants in
+ ``toolsets/unity_mcp.py`` to match, or such failures are mistaken for the
server's answers.
+ """
+
+ MADE_UP_ERRORS = [
+ ("mcp.client.streamable_http", "INTERNAL_ERROR", "Server returned an
error response"),
+ ("mcp.client.streamable_http", "METHOD_NOT_FOUND", "Not Found"),
+ ("mcp.client.streamable_http", "INVALID_REQUEST", "Session
terminated"),
+ ("mcp.shared.jsonrpc_dispatcher", "CONNECTION_CLOSED", "Connection
closed"),
+ ]
+
+ @pytest.mark.parametrize(("module", "code", "message"), MADE_UP_ERRORS)
+ def test_error_is_still_made_up_by_the_client(self, module, code, message):
+ assert f'code={code}, message="{message}"' in
inspect.getsource(importlib.import_module(module))
+
+ def test_every_matched_error_is_checked(self):
+ expected = {(getattr(jsonrpc, code), message) for _, code, message in
self.MADE_UP_ERRORS}
+ assert expected == unity_mcp._CLIENT_ERRORS
+
+ @staticmethod
+ def _find_made_up_errors() -> list[tuple[int, str]]:
+ """Return the code and message, or message prefix, of each error the
HTTP client makes up."""
+ module = importlib.import_module("mcp.client.streamable_http")
+ found = []
+ for node in ast.walk(ast.parse(inspect.getsource(module))):
+ if not isinstance(node, ast.Call):
+ continue
+ keywords = {k.arg: k.value for k in node.keywords}
+ func = node.func
+ if isinstance(func, ast.Name) and func.id == "ErrorData":
+ code, message = keywords.get("code"), keywords.get("message")
+ elif isinstance(func, ast.Attribute) and func.attr ==
"_resolve_abandoned_request":
+ code, message = keywords.get("code"), node.args[2]
+ else:
+ continue
+ if isinstance(message, ast.JoinedStr):
+ message = message.values[0]
+ # Messages built elsewhere, such as the redirect one, are checked
by the prefix test.
+ if not isinstance(message, ast.Constant) or not
isinstance(message.value, str):
+ continue
+ code_name = code.id if isinstance(code, ast.Name) else
"CONNECTION_CLOSED"
+ found.append((getattr(jsonrpc, code_name), message.value))
+ return found
+
+ # mcp 2.3.0 adds this one, for an SSE error while reading the answer to a
request.
+ def test_failed_sse_stream_is_recognized(self):
+ assert unity_mcp._is_client_error(McpError(jsonrpc.CONNECTION_CLOSED,
"SSE stream failed: boom"))
+
+ def test_every_error_the_client_makes_up_is_recognized(self):
+ made_up = self._find_made_up_errors()
+
+ assert len(made_up) >= 8
+ for code, message in made_up:
+ assert unity_mcp._is_client_error(McpError(code, message)), message
+
+ @pytest.mark.parametrize("prefix", unity_mcp._CLIENT_ERROR_PREFIXES)
+ def test_message_prefix_is_still_made_up_by_the_client(self, prefix):
+ assert prefix in
inspect.getsource(importlib.import_module("mcp.client.streamable_http"))
+
+ def test_session_ended_message_is_still_raised_by_fastmcp(self):
+ source =
inspect.getsource(importlib.import_module("fastmcp.client.client"))
+ assert unity_mcp._SESSION_ENDED_MESSAGE in source
+
+ def test_handshake_timeout_message_is_still_raised_by_fastmcp(self):
+ source =
inspect.getsource(importlib.import_module("fastmcp.client.client"))
+ assert unity_mcp._HANDSHAKE_TIMEOUT_MESSAGE in source
+
+
+class TestHandshakeTimeout:
+ def test_slow_handshake_within_init_timeout_connects(self, gateway,
gateway_conn):
+ _StubGateway.delays = {"initialize": 1}
+
+ _, result = _run(
+ _list_and_call(DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID, init_timeout=10))
+ )
+
+ assert result == "echo: hi"
+
+ def test_handshake_beyond_init_timeout_is_a_transport_error(self, gateway,
gateway_conn):
+ # Shorter than pydantic-ai's own 5 second default, so this only fails
if init_timeout is used.
+ _StubGateway.delays = {"initialize": 2}
+
+ with pytest.raises(DatabricksUnityMCPTransportError, match="within 0.5
seconds. Raise init_timeout"):
+ _run(
+ _list_and_call(
+ DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID, init_timeout=0.5)
+ )
+ )
+
+ @pytest.mark.parametrize("init_timeout", [0, -1])
+ def test_init_timeout_must_be_positive(self, init_timeout):
+ with pytest.raises(ValueError, match="init_timeout must be a positive
number"):
+ DatabricksUnityMCPToolset(SERVICE, init_timeout=init_timeout)
+
+
+class TestClosing:
+ @pytest.mark.parametrize("failed_before", [False, True])
+ def test_error_while_closing_after_a_clean_run_is_raised(self, gateway,
gateway_conn, failed_before):
+ toolset = DatabricksUnityMCPToolset(SERVICE,
databricks_conn_id=CONN_ID)
+ if failed_before:
+ # A failure in an earlier run must not hide close errors in the
next one.
+ with pytest.raises(DatabricksUnityMCPError):
+ _run(_list_and_call(toolset, "reply 503"))
+
+ async def close_with_an_error():
+ async with toolset:
+ await toolset.get_tools(_ctx())
+ real_exit = toolset._server.__aexit__
+
+ async def failing_exit(*args):
+ await real_exit(*args)
+ raise RuntimeError("close failed")
+
+ toolset._server.__aexit__ = failing_exit
+
+ with pytest.raises(RuntimeError, match="close failed"):
+ _run(close_with_an_error())
+
+
+class TestFrameworkNeutralTools:
+ def test_are_not_supported(self):
+ with pytest.raises(NotImplementedError, match="framework-neutral"):
+ DatabricksUnityMCPToolset(SERVICE).airflow_tools()
+
+
+class TestOptionalDependency:
+ def test_import_without_common_ai_names_the_extra(self):
+ module = "airflow.providers.databricks.toolsets.unity_mcp"
+ with mock.patch.dict(sys.modules,
{"airflow.providers.common.ai.utils.toolset_base": None}):
+ sys.modules.pop(module)
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match=r"databricks\[common.ai\]"):
+ importlib.import_module(module)
diff --git a/uv.lock b/uv.lock
index bb6963fd7f3..76132824ca5 100644
--- a/uv.lock
+++ b/uv.lock
@@ -5129,7 +5129,8 @@ azure-identity = [
{ name = "azure-identity" },
]
common-ai = [
- { name = "apache-airflow-providers-common-ai" },
+ { name = "apache-airflow-providers-common-ai", extra = ["mcp"] },
+ { name = "fastmcp-slim", extra = ["client"] },
]
fab = [
{ name = "apache-airflow-providers-fab" },
@@ -5158,7 +5159,7 @@ dev = [
{ name = "apache-airflow-providers-common-ai" },
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-sql", extra = ["pandas",
"polars"] },
- { name = "apache-airflow-providers-databricks", extra = ["sqlalchemy"] },
+ { name = "apache-airflow-providers-databricks", extra = ["common-ai",
"sqlalchemy"] },
{ name = "apache-airflow-providers-fab" },
{ name = "apache-airflow-providers-google" },
{ name = "apache-airflow-providers-microsoft-azure" },
@@ -5175,7 +5176,7 @@ requires-dist = [
{ name = "aiohttp", specifier = ">=3.14.0,<4" },
{ name = "apache-airflow", editable = "." },
{ name = "apache-airflow-providers-amazon", marker = "extra == 'amazon'",
editable = "providers/amazon" },
- { name = "apache-airflow-providers-common-ai", marker = "extra ==
'common-ai'", editable = "providers/common/ai" },
+ { name = "apache-airflow-providers-common-ai", extras = ["mcp"], marker =
"extra == 'common-ai'", editable = "providers/common/ai" },
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
{ name = "apache-airflow-providers-fab", marker = "extra == 'fab'",
editable = "providers/fab" },
@@ -5189,6 +5190,7 @@ requires-dist = [
{ name = "fastavro", marker = "python_full_version >= '3.12' and
python_full_version < '3.14' and extra == 'avro'", specifier = ">=1.10.0" },
{ name = "fastavro", marker = "python_full_version >= '3.14' and extra ==
'avro'", specifier = ">=1.12.1" },
{ name = "fastavro", marker = "python_full_version < '3.14' and extra ==
'avro'", specifier = ">=1.9.0" },
+ { name = "fastmcp-slim", extras = ["client"], marker = "extra ==
'common-ai'", specifier = ">=4.0.0" },
{ name = "mergedeep", specifier = ">=1.3.4" },
{ name = "pandas", marker = "python_full_version < '3.13'", specifier =
">=2.1.2" },
{ name = "pandas", marker = "python_full_version == '3.13.*'", specifier =
">=2.2.3" },
@@ -5199,7 +5201,7 @@ requires-dist = [
{ name = "requests", specifier = ">=2.32.0,<3" },
{ name = "tenacity", specifier = ">=8.3.0" },
]
-provides-extras = ["common-ai", "avro", "amazon", "azure-identity", "fab",
"google", "sdk", "standard", "openlineage", "sqlalchemy"]
+provides-extras = ["avro", "amazon", "azure-identity", "common-ai", "fab",
"google", "sdk", "standard", "openlineage", "sqlalchemy"]
[package.metadata.requires-dev]
dev = [
@@ -5210,7 +5212,7 @@ dev = [
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
{ name = "apache-airflow-providers-common-sql", extras = ["pandas",
"polars"], editable = "providers/common/sql" },
- { name = "apache-airflow-providers-databricks", extras = ["sqlalchemy"],
editable = "providers/databricks" },
+ { name = "apache-airflow-providers-databricks", extras = ["common-ai",
"sqlalchemy"], editable = "providers/databricks" },
{ name = "apache-airflow-providers-fab", editable = "providers/fab" },
{ name = "apache-airflow-providers-google", editable = "providers/google"
},
{ name = "apache-airflow-providers-microsoft-azure", editable =
"providers/microsoft/azure" },
