From bf153331abeb9e5e3d1e2811d36057605e048dd4 Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:45:42 +0000 Subject: [PATCH 1/8] Add a Databricks toolset for Unity AI Gateway MCP Services 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 --- providers/databricks/docs/index.rst | 3 + providers/databricks/docs/toolsets/index.rst | 33 ++ .../databricks/docs/toolsets/unity_mcp.rst | 87 +++++ providers/databricks/provider.yaml | 10 + providers/databricks/pyproject.toml | 4 + .../providers/databricks/exceptions.py | 31 ++ .../providers/databricks/get_provider_info.py | 6 + .../providers/databricks/toolsets/__init__.py | 17 + .../databricks/toolsets/unity_mcp.py | 337 ++++++++++++++++++ .../example_databricks_unity_mcp.py | 68 ++++ .../unit/databricks/toolsets/__init__.py | 16 + .../databricks/toolsets/test_unity_mcp.py | 323 +++++++++++++++++ 12 files changed, 935 insertions(+) create mode 100644 providers/databricks/docs/toolsets/index.rst create mode 100644 providers/databricks/docs/toolsets/unity_mcp.rst create mode 100644 providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py create mode 100644 providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py create mode 100644 providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py create mode 100644 providers/databricks/tests/unit/databricks/toolsets/__init__.py create mode 100644 providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py diff --git a/providers/databricks/docs/index.rst b/providers/databricks/docs/index.rst index 7aed991add904..37823f8803daf 100644 --- a/providers/databricks/docs/index.rst +++ b/providers/databricks/docs/index.rst @@ -36,6 +36,7 @@ Connection types Operators + Toolsets Plugins .. toctree:: @@ -133,6 +134,7 @@ You can install such cross-provider dependencies when installing from PyPI. For Dependent package Extra ============================================================================================================== =============== `apache-airflow-providers-amazon `_ ``amazon`` +`apache-airflow-providers-common-ai `_ ``common.ai`` `apache-airflow-providers-google `_ ``google`` `apache-airflow-providers-openlineage `_ ``openlineage`` ============================================================================================================== =============== @@ -154,6 +156,7 @@ Extra Dependencies ``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]>=0.10.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 0000000000000..cdad1abfd697c --- /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 ` +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 0000000000000..5909305376872 --- /dev/null +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -0,0 +1,87 @@ + .. 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 AI Gateway MCP Services +============================= + +Use :class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset` to give an +agent the tools of a Unity AI Gateway MCP Service. One service can front Unity Catalog functions, +Genie spaces, AI Search indexes, or an external MCP server registered 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 ` to use. The toolset builds the service URL, +``https:///ai-gateway/mcp-services/``, 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. 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 +---------------------------------- + +The gateway runs every tool call as the identity of the connection's credentials, which needs +``EXECUTE`` on the MCP Service in Unity Catalog. It exposes only the tools selected for the service, +and the service's policies apply. Grant that identity only the services the agent should use. + +The gateway accepts bearer tokens only, 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). + +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 +------------------ + +Gateway errors are raised as dedicated exceptions from +:mod:`airflow.providers.databricks.exceptions`: + +.. list-table:: + :header-rows: 1 + + * - Exception + - Cause + * - ``DatabricksUnityMCPAccessDeniedError`` + - HTTP 401 or 403: the credentials are invalid, or the identity lacks ``EXECUTE`` on the service. + * - ``DatabricksUnityMCPServiceNotFoundError`` + - HTTP 404: the service does not exist on the workspace, or is not visible to the identity. + * - ``DatabricksUnityMCPThrottledError`` + - HTTP 429. Its ``retry_after`` attribute holds the ``Retry-After`` delay in seconds, when given. + * - ``DatabricksUnityMCPTransportError`` + - The gateway could not be reached, or the connection dropped. + * - ``DatabricksUnityMCPError`` + - Any other gateway error, or a failure to get a token from the connection. + +The toolset never retries a tool call. A call that fails after it was sent, with a server or transport +error, may or may not have run, and repeating a tool that changes data could apply the change twice; +the error message says so. Use task retries for calls that are safe to repeat. diff --git a/providers/databricks/provider.yaml b/providers/databricks/provider.yaml index 2b66f546bf2a6..33e26bf8bcd73 100644 --- a/providers/databricks/provider.yaml +++ b/providers/databricks/provider.yaml @@ -144,6 +144,11 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-databricks/operators/workflow.rst tags: [service] + - integration-name: Databricks Unity AI Gateway + external-doc-url: https://docs.databricks.com/aws/en/unity-gateway/concepts + how-to-guide: + - /docs/apache-airflow-providers-databricks/toolsets/unity_mcp.rst + tags: [service] operators: - integration-name: Databricks @@ -195,6 +200,11 @@ sensors: - airflow.providers.databricks.sensors.databricks_sql - airflow.providers.databricks.sensors.databricks_partition +toolsets: + - integration-name: Databricks Unity AI 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 28910f667f5a8..42cb536a25805 100644 --- a/providers/databricks/pyproject.toml +++ b/providers/databricks/pyproject.toml @@ -88,6 +88,9 @@ dependencies = [ "azure-identity" = [ "azure-identity>=1.25.3", ] +"common.ai" = [ + "apache-airflow-providers-common-ai[mcp]>=0.10.0", # use next version +] "fab" = [ "apache-airflow-providers-fab>=2.2.0" ] @@ -113,6 +116,7 @@ dev = [ "apache-airflow-task-sdk", "apache-airflow-devel-common", "apache-airflow-providers-amazon", + "apache-airflow-providers-common-ai", "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", "apache-airflow-providers-google", diff --git a/providers/databricks/src/airflow/providers/databricks/exceptions.py b/providers/databricks/src/airflow/providers/databricks/exceptions.py index a0f52a6077c06..5b3fa490f33fb 100644 --- a/providers/databricks/src/airflow/providers/databricks/exceptions.py +++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py @@ -46,3 +46,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 AI Gateway MCP Service fails.""" + + +class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError): + """Raised when the gateway rejects the caller's credentials or the caller lacks EXECUTE on the service.""" + + +class DatabricksUnityMCPServiceNotFoundError(DatabricksUnityMCPError): + """Raised when the MCP Service does not exist, or is not visible to the caller.""" + + +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 a218ce3546417..9301589d3d1c2 100644 --- a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py +++ b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py @@ -67,6 +67,12 @@ def get_provider_info(): "how-to-guide": ["/docs/apache-airflow-providers-databricks/operators/workflow.rst"], "tags": ["service"], }, + { + "integration-name": "Databricks Unity AI Gateway", + "external-doc-url": "https://docs.databricks.com/aws/en/unity-gateway/concepts", + "how-to-guide": ["/docs/apache-airflow-providers-databricks/toolsets/unity_mcp.rst"], + "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 0000000000000..217e5db960782 --- /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 0000000000000..0acc0ad5740e2 --- /dev/null +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -0,0 +1,337 @@ +# 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 AI Gateway MCP Service.""" + +from __future__ import annotations + +import contextlib +import re +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPServiceNotFoundError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook + +try: + import httpx2 + from fastmcp.client.transports import StreamableHttpTransport + from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + + from airflow.providers.common.ai.toolsets.mcp import MCPToolset + 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, Generator, Sequence + + from pydantic_ai._run_context import RunContext + from pydantic_ai.toolsets.abstract import ToolsetTool + + from airflow.sdk.execution_time.secrets_masker import mask_secret +else: + try: + from airflow.sdk.log import mask_secret + except ImportError: + try: + from airflow.sdk.execution_time.secrets_masker import mask_secret + except ImportError: + from airflow.utils.log.secrets_masker import mask_secret + +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}") + + +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, underscores and hyphens. + """ + if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): + raise ValueError( + f"Invalid Unity AI Gateway MCP Service name {service_name!r}: expected " + "'catalog.schema.service', each part made of ASCII letters, digits, '_' or '-'." + ) + + +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 _find_transport_error(exc: BaseException) -> httpx2.TransportError | None: + """Return the network error behind ``exc``, looking through causes and exception groups.""" + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, httpx2.TransportError): + return 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) + return 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. + + The MCP client reports every HTTP error from the gateway as the same generic error, so the + status and ``Retry-After`` of the last error response, and any failure to get a token, are + kept here for :class:`DatabricksUnityMCPToolset` to report what went wrong. With concurrent + calls on one toolset, the error reported for a failed call can be another call's. + """ + + def __init__(self, hook: DatabricksHook) -> None: + self._hook = hook + self.error_status: int | None = None + self.retry_after: float | None = None + self.token_error: Exception | None = None + + def get_token(self) -> str: + 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. Unity AI Gateway accepts bearer tokens only: use a personal access " + "token, service principal OAuth, Azure AD, or workload identity federation. " + "Username and password authentication is not supported." + ) + # 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) + return token + + def take_error(self) -> tuple[int | None, float | None, Exception | None]: + """Return the last recorded error and forget it.""" + error = (self.error_status, self.retry_after, self.token_error) + self.error_status = self.retry_after = self.token_error = None + return error + + def _record(self, response: httpx2.Response) -> None: + if response.status_code >= 400: + self.error_status = response.status_code + self.retry_after = _parse_retry_after(response.headers.get("Retry-After")) + + def sync_auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + try: + token = self.get_token() + except Exception as e: + self.token_error = e + raise + request.headers["Authorization"] = f"Bearer {token}" + response = yield request + self._record(response) + + async def async_auth_flow( + self, request: httpx2.Request + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + try: + # Minting an OAuth token makes a blocking HTTP call; keep it off the event loop. + token = await AirflowToolset.run_blocking(self.get_token) + except Exception as e: + self.token_error = e + raise + request.headers["Authorization"] = f"Bearer {token}" + response = yield request + self._record(response) + + +class DatabricksUnityMCPToolset(MCPToolset): + """ + Give an agent the tools of a Unity AI 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:///ai-gateway/mcp-services/``. 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. + + 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), which needs + ``EXECUTE`` on the MCP Service in Unity Catalog. 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. Gateway errors are raised as + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPServiceNotFoundError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError` or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`. Tool calls + are never retried by this toolset, because a call interrupted after it was sent may already + have run. + + .. 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_genie", + prompt="How many orders shipped yesterday?", + llm_conn_id="pydanticai_default", + toolsets=[DatabricksUnityMCPToolset("system.ai.genie_one_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. + """ + + 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, + ) -> None: + super().__init__(databricks_conn_id, tool_prefix=tool_prefix) + self._databricks_conn_id = databricks_conn_id + self._service_name = service_name + self._auth: _DatabricksTokenAuth | None = None + # A templated name is checked once it has been rendered, in _get_server. + 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.") + return hook._endpoint_url(f"{GATEWAY_MCP_SERVICES_PATH}/{self._service_name}") + + def _get_server(self) -> Any: + if self._server is None: + 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. + auth.get_token() + transport = StreamableHttpTransport(url, headers=hook.user_agent_header, auth=auth) + toolset = PydanticAIMCPToolset(transport) + self._auth = auth + self._server = toolset.prefixed(self._tool_prefix) if self._tool_prefix else toolset + return self._server + + def _translate_error(self, error: Exception, *, during_tool_call: bool) -> Exception | None: + """Return the provider exception for a gateway failure, or ``None`` to re-raise ``error`` as is.""" + status, retry_after, token_error = self._auth.take_error() if self._auth else (None, None, None) + service = self._service_name + 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}" + ) + if status in (401, 403): + return DatabricksUnityMCPAccessDeniedError( + f"Unity AI Gateway denied access to MCP Service {service!r} (HTTP {status}). The " + f"identity of connection {self._databricks_conn_id!r} needs EXECUTE on the service, " + "and its credentials must be valid.", + http_status_code=status, + ) + if status == 404: + return DatabricksUnityMCPServiceNotFoundError( + f"MCP Service {service!r} was not found on the workspace of connection " + f"{self._databricks_conn_id!r}, or is not visible to its identity (HTTP 404).", + 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 AI 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 "" + ) + if status is not None: + return DatabricksUnityMCPError( + f"Unity AI Gateway returned HTTP {status} for MCP Service {service!r}.{ambiguous}", + http_status_code=status, + ) + if (transport_error := _find_transport_error(error)) is not None: + return DatabricksUnityMCPTransportError( + f"Could not reach Unity AI Gateway for MCP Service {service!r}: {transport_error}.{ambiguous}" + ) + return None + + @contextlib.asynccontextmanager + async def _gateway_errors(self, *, during_tool_call: bool = False) -> AsyncIterator[None]: + try: + yield + except Exception as e: + translated = self._translate_error(e, during_tool_call=during_tool_call) + if translated is None: + raise + raise translated from e + + async def __aenter__(self) -> DatabricksUnityMCPToolset: + async with self._gateway_errors(): + await super().__aenter__() + return self + + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: + async with self._gateway_errors(): + return await super().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_errors(during_tool_call=True): + return await super().execute_tool(name, tool_args, ctx=ctx, tool=tool) 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 0000000000000..e6f90ca84f35b --- /dev/null +++ b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py @@ -0,0 +1,68 @@ +# 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 AI Gateway MCP Service. + +Needs a Databricks connection whose identity has EXECUTE on the MCP Service, 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 0000000000000..13a83393a9124 --- /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 0000000000000..689819e90a16f --- /dev/null +++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py @@ -0,0 +1,323 @@ +# 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 asyncio +import importlib +import json +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") + +from pydantic_ai import RunContext + +from airflow.models import Connection +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPServiceNotFoundError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook +from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + +CONN_ID = "databricks_unity" +SERVICE = "main.tools.genie" + + +class _StubGateway(BaseHTTPRequestHandler): + """A Unity AI Gateway stand-in that serves one MCP Service with a single ``echo`` tool.""" + + protocol_version = "HTTP/1.1" + # Per-test behaviour, set by the ``gateway`` fixture: service name -> status for every request, + # and a status for ``tools/call`` requests only. + service_status: dict[str, tuple[int, dict[str, str]]] = {} + tool_call_status: int | None = None + 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_status: + status, headers = self.service_status[service] + return self._reply(status, headers=headers) + if method == "tools/call" and self.tool_call_status: + return self._reply(self.tool_call_status) + if "id" not in body: + return self._reply(202) + 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": + text = body["params"]["arguments"]["text"] + 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(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 + + +@pytest.fixture +def gateway(): + _StubGateway.service_status = {} + _StubGateway.tool_call_status = None + _StubGateway.requests = [] + server = ThreadingHTTPServer(("127.0.0.1", 0), _StubGateway) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server + server.shutdown() + server.server_close() + + +@pytest.fixture +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 + + +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..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)) + + +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("airflow.providers.databricks.toolsets.unity_mcp.mask_secret", autospec=True) + def test_minted_tokens_are_masked(self, mask_secret, gateway, gateway_conn): + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + mask_secret.assert_called_with("dapi-token") + + 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="bearer tokens only"): + 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) + + +class TestErrors: + @pytest.mark.parametrize( + ("status", "headers", "expected", "retry_after"), + [ + (401, {}, DatabricksUnityMCPAccessDeniedError, None), + (403, {}, DatabricksUnityMCPAccessDeniedError, None), + (404, {}, DatabricksUnityMCPServiceNotFoundError, None), + (429, {"Retry-After": "12"}, DatabricksUnityMCPThrottledError, 12.0), + (429, {"Retry-After": "Wed, 21 Oct 2026 07:28:00 GMT"}, DatabricksUnityMCPThrottledError, None), + ], + ) + def test_gateway_errors_are_translated( + self, gateway, gateway_conn, status, headers, expected, retry_after + ): + _StubGateway.service_status[SERVICE] = (status, headers) + + with pytest.raises(expected) as err: + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + assert err.value.http_status_code == status + if expected is DatabricksUnityMCPThrottledError: + assert err.value.retry_after == retry_after + + def test_server_error_during_a_tool_call_says_the_call_may_have_run(self, gateway, gateway_conn): + _StubGateway.tool_call_status = 503 + + with pytest.raises(DatabricksUnityMCPError, match="may or may not have run") as err: + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + 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_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 AI Gateway"): + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + +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.toolsets.mcp": None}): + sys.modules.pop(module) + with pytest.raises(AirflowOptionalProviderFeatureException, match=r"databricks\[common.ai\]"): + importlib.import_module(module) From b5eea8d5d7f66e47141ee6d956cce25548484061 Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 4 Oct 2026 23:43:05 +0000 Subject: [PATCH 2/8] Keep Unity MCP requests from waiting on other toolsets' blocking calls Each gateway request fetched its token through AirflowToolset.run_blocking, whose lock is shared by every toolset in the process. While another toolset ran a long blocking call, such as a SQL query, every request to the MCP Service waited for it to finish. --- .../providers/databricks/toolsets/unity_mcp.py | 14 ++++++++++---- .../unit/databricks/toolsets/test_unity_mcp.py | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py index 0acc0ad5740e2..806901fa5af91 100644 --- a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -18,8 +18,10 @@ from __future__ import annotations +import asyncio import contextlib import re +import threading from typing import TYPE_CHECKING, Any from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException @@ -38,7 +40,6 @@ from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset from airflow.providers.common.ai.toolsets.mcp import MCPToolset - 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: " @@ -129,9 +130,11 @@ def __init__(self, hook: DatabricksHook) -> None: self.error_status: int | None = None self.retry_after: float | None = None self.token_error: Exception | None = None + self._token_lock = threading.Lock() def get_token(self) -> str: - token = self._hook._get_token(raise_error=False) + 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 " @@ -169,8 +172,11 @@ async def async_auth_flow( self, request: httpx2.Request ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: try: - # Minting an OAuth token makes a blocking HTTP call; keep it off the event loop. - token = await AirflowToolset.run_blocking(self.get_token) + # 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 connection is already + # resolved by _get_server, so fetching a token only calls the token endpoint, and the + # hook is this toolset's own, so a lock of its own is enough. + token = await asyncio.to_thread(self.get_token) except Exception as e: self.token_error = e raise diff --git a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py index 689819e90a16f..a5329ab379be1 100644 --- a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py +++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py @@ -32,6 +32,7 @@ from pydantic_ai import RunContext 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, @@ -244,6 +245,19 @@ def test_token_is_fetched_for_every_request_so_it_can_refresh(self, get_token, g # 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)] + 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("airflow.providers.databricks.toolsets.unity_mcp.mask_secret", autospec=True) def test_minted_tokens_are_masked(self, mask_secret, gateway, gateway_conn): _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) From 6b3e4abe83bbd0248873920ac2f4d3a28a96531d Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 4 Oct 2026 23:48:05 +0000 Subject: [PATCH 3/8] Document the Unity Catalog privileges MCP Services need Databricks requires USE CATALOG and USE SCHEMA on the parent catalog and schema as well as EXECUTE on the service, so an identity granted only EXECUTE is denied. Users following the guide or the access-denied error would grant too little. --- .../databricks/docs/toolsets/unity_mcp.rst | 22 ++++++++++++++----- .../providers/databricks/exceptions.py | 2 +- .../databricks/toolsets/unity_mcp.py | 14 ++++++------ .../example_databricks_unity_mcp.py | 5 +++-- 4 files changed, 27 insertions(+), 16 deletions(-) diff --git a/providers/databricks/docs/toolsets/unity_mcp.rst b/providers/databricks/docs/toolsets/unity_mcp.rst index 5909305376872..3c7f06248c033 100644 --- a/providers/databricks/docs/toolsets/unity_mcp.rst +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -44,12 +44,21 @@ The service name and connection ID are templated when the toolset is passed to ` Caller identity and authentication ---------------------------------- -The gateway runs every tool call as the identity of the connection's credentials, which needs -``EXECUTE`` on the MCP Service in Unity Catalog. It exposes only the tools selected for the service, -and the service's policies apply. Grant that identity only the services the agent should use. +MCP Services need a workspace enabled for Unity Catalog, in a region where Model Serving is +supported. -The gateway accepts bearer tokens only, so the connection must use one of these authentication modes -of the Databricks connection: +The gateway runs every tool call as the identity of the connection's credentials. That identity needs +``EXECUTE`` on the MCP Service, and ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and +schema: ``EXECUTE`` alone is not enough. Account users hold these on the built-in ``system.ai`` +services by default. 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``); @@ -72,7 +81,8 @@ Gateway errors are raised as dedicated exceptions from * - Exception - Cause * - ``DatabricksUnityMCPAccessDeniedError`` - - HTTP 401 or 403: the credentials are invalid, or the identity lacks ``EXECUTE`` on the service. + - HTTP 401 or 403: the credentials are invalid, or the identity lacks ``EXECUTE`` on the service + or ``USE CATALOG`` / ``USE SCHEMA`` on its parents. * - ``DatabricksUnityMCPServiceNotFoundError`` - HTTP 404: the service does not exist on the workspace, or is not visible to the identity. * - ``DatabricksUnityMCPThrottledError`` diff --git a/providers/databricks/src/airflow/providers/databricks/exceptions.py b/providers/databricks/src/airflow/providers/databricks/exceptions.py index 5b3fa490f33fb..3ee06c4f77225 100644 --- a/providers/databricks/src/airflow/providers/databricks/exceptions.py +++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py @@ -53,7 +53,7 @@ class DatabricksUnityMCPError(DatabricksApiError): class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError): - """Raised when the gateway rejects the caller's credentials or the caller lacks EXECUTE on the service.""" + """Raised when the gateway rejects the caller's credentials or the caller lacks a privilege on the service.""" class DatabricksUnityMCPServiceNotFoundError(DatabricksUnityMCPError): diff --git a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py index 806901fa5af91..6f536b88e238d 100644 --- a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -196,8 +196,8 @@ class DatabricksUnityMCPToolset(MCPToolset): 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), which needs - ``EXECUTE`` on the MCP Service in Unity Catalog. It sees only the tools selected for the - service, and the service's policies apply. + ``EXECUTE`` on the MCP Service and ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and schema. + 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. Gateway errors are raised as @@ -214,10 +214,10 @@ class DatabricksUnityMCPToolset(MCPToolset): from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset AgentOperator( - task_id="ask_genie", - prompt="How many orders shipped yesterday?", + task_id="ask_mcp_service", + prompt="Which tools do you have?", llm_conn_id="pydanticai_default", - toolsets=[DatabricksUnityMCPToolset("system.ai.genie_one_mcp", databricks_conn_id="databricks")], + toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp", databricks_conn_id="databricks")], ) :param service_name: Three-level name of the MCP Service, ``catalog.schema.service``. Templated @@ -281,8 +281,8 @@ def _translate_error(self, error: Exception, *, during_tool_call: bool) -> Excep if status in (401, 403): return DatabricksUnityMCPAccessDeniedError( f"Unity AI Gateway denied access to MCP Service {service!r} (HTTP {status}). The " - f"identity of connection {self._databricks_conn_id!r} needs EXECUTE on the service, " - "and its credentials must be valid.", + f"identity of connection {self._databricks_conn_id!r} needs EXECUTE on the service and " + "USE CATALOG and USE SCHEMA on its catalog and schema, and its credentials must be valid.", http_status_code=status, ) if status == 404: diff --git a/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py index e6f90ca84f35b..49e10f0e072e2 100644 --- a/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py +++ b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py @@ -17,8 +17,9 @@ """ Example Dag giving an agent the tools of a Unity AI Gateway MCP Service. -Needs a Databricks connection whose identity has EXECUTE on the MCP Service, an LLM connection for -the agent, and these environment variables: +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 From e60246e3dea28474d4022b07b409f28a5d8e00a0 Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 4 Oct 2026 23:49:30 +0000 Subject: [PATCH 4/8] Note that MCP Service callers must be assigned to the workspace Databricks also requires the calling identity to be assigned to the workspace the request goes to, which the guide did not say. --- providers/databricks/docs/toolsets/unity_mcp.rst | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/providers/databricks/docs/toolsets/unity_mcp.rst b/providers/databricks/docs/toolsets/unity_mcp.rst index 3c7f06248c033..45989ad5be6c9 100644 --- a/providers/databricks/docs/toolsets/unity_mcp.rst +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -49,7 +49,9 @@ supported. The gateway runs every tool call as the identity of the connection's credentials. That identity needs ``EXECUTE`` on the MCP Service, and ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and -schema: ``EXECUTE`` alone is not enough. Account users hold these on the built-in ``system.ai`` +schema: ``EXECUTE`` alone is not enough. It must also be assigned to the workspace. It needs no +privilege on the Unity Catalog connection behind the service. Account users hold these on the +built-in ``system.ai`` services by default. 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. From fd1dbdf95de71aa8882469627640d715163a0010 Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 4 Oct 2026 23:55:51 +0000 Subject: [PATCH 5/8] Report Unity MCP tool errors the model can fix as retries The MCP client carries on after some gateway error responses, such as one to a notification, but the toolset kept that status and later reported an unrelated failure as it. A tool call with bad arguments then ended the agent run as a gateway error instead of letting the model retry. Also use Databricks' name for the product, Unity Gateway, and correct the guide: Databricks provides MCP Services for workspace tools and SaaS applications, and only external MCP servers can be registered as one. --- .../databricks/docs/toolsets/unity_mcp.rst | 24 ++++----- providers/databricks/provider.yaml | 4 +- .../providers/databricks/exceptions.py | 2 +- .../providers/databricks/get_provider_info.py | 2 +- .../databricks/toolsets/unity_mcp.py | 52 +++++++++---------- .../example_databricks_unity_mcp.py | 2 +- .../databricks/toolsets/test_unity_mcp.py | 31 ++++++++--- 7 files changed, 66 insertions(+), 51 deletions(-) diff --git a/providers/databricks/docs/toolsets/unity_mcp.rst b/providers/databricks/docs/toolsets/unity_mcp.rst index 45989ad5be6c9..f2d2e690313f9 100644 --- a/providers/databricks/docs/toolsets/unity_mcp.rst +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -17,14 +17,15 @@ .. _howto/toolset:DatabricksUnityMCPToolset: -Unity AI Gateway MCP Services -============================= +Unity Gateway MCP Services +========================== Use :class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset` to give an -agent the tools of a Unity AI Gateway MCP Service. One service can front Unity Catalog functions, -Genie spaces, AI Search indexes, or an external MCP server registered 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. +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 ` to use. The toolset builds the service URL, @@ -48,12 +49,11 @@ MCP Services need a workspace enabled for Unity Catalog, in a region where Model supported. The gateway runs every tool call as the identity of the connection's credentials. That identity needs -``EXECUTE`` on the MCP Service, and ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and -schema: ``EXECUTE`` alone is not enough. It must also be assigned to the workspace. It needs no -privilege on the Unity Catalog connection behind the service. Account users hold these on the -built-in ``system.ai`` -services by default. 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. +``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. Account users already hold these privileges on the +built-in ``system.ai`` 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 diff --git a/providers/databricks/provider.yaml b/providers/databricks/provider.yaml index 33e26bf8bcd73..36c8e3202c1f4 100644 --- a/providers/databricks/provider.yaml +++ b/providers/databricks/provider.yaml @@ -144,7 +144,7 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-databricks/operators/workflow.rst tags: [service] - - integration-name: Databricks Unity AI Gateway + - integration-name: Databricks Unity Gateway external-doc-url: https://docs.databricks.com/aws/en/unity-gateway/concepts how-to-guide: - /docs/apache-airflow-providers-databricks/toolsets/unity_mcp.rst @@ -201,7 +201,7 @@ sensors: - airflow.providers.databricks.sensors.databricks_partition toolsets: - - integration-name: Databricks Unity AI Gateway + - integration-name: Databricks Unity Gateway python-modules: - airflow.providers.databricks.toolsets.unity_mcp diff --git a/providers/databricks/src/airflow/providers/databricks/exceptions.py b/providers/databricks/src/airflow/providers/databricks/exceptions.py index 3ee06c4f77225..5962b82dc21d7 100644 --- a/providers/databricks/src/airflow/providers/databricks/exceptions.py +++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py @@ -49,7 +49,7 @@ def __init__(self, message: str, *, http_status_code: int | None = None) -> None class DatabricksUnityMCPError(DatabricksApiError): - """Raised when a call to a Unity AI Gateway MCP Service fails.""" + """Raised when a call to a Unity Gateway MCP Service fails.""" class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError): 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 9301589d3d1c2..59229f0fab65c 100644 --- a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py +++ b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py @@ -68,7 +68,7 @@ def get_provider_info(): "tags": ["service"], }, { - "integration-name": "Databricks Unity AI Gateway", + "integration-name": "Databricks Unity Gateway", "external-doc-url": "https://docs.databricks.com/aws/en/unity-gateway/concepts", "how-to-guide": ["/docs/apache-airflow-providers-databricks/toolsets/unity_mcp.rst"], "tags": ["service"], diff --git a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py index 6f536b88e238d..2de62016800ac 100644 --- a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -14,7 +14,7 @@ # 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 AI Gateway MCP Service.""" +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" from __future__ import annotations @@ -47,7 +47,7 @@ ) from e if TYPE_CHECKING: - from collections.abc import AsyncGenerator, AsyncIterator, Generator, Sequence + from collections.abc import AsyncGenerator, AsyncIterator, Sequence from pydantic_ai._run_context import RunContext from pydantic_ai.toolsets.abstract import ToolsetTool @@ -78,7 +78,7 @@ def validate_service_name(service_name: str) -> None: """ if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): raise ValueError( - f"Invalid Unity AI Gateway MCP Service name {service_name!r}: expected " + f"Invalid Unity Gateway MCP Service name {service_name!r}: expected " "'catalog.schema.service', each part made of ASCII letters, digits, '_' or '-'." ) @@ -138,9 +138,9 @@ def get_token(self) -> str: if not token: raise ValueError( f"Connection {self._hook.databricks_conn_id!r} has no token-based authentication " - "configured. Unity AI Gateway accepts bearer tokens only: use a personal access " - "token, service principal OAuth, Azure AD, or workload identity federation. " - "Username and password authentication is not supported." + "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." ) # A personal access token is masked when the connection is fetched; mask minted # OAuth tokens too, so they never reach task logs. @@ -154,19 +154,13 @@ def take_error(self) -> tuple[int | None, float | None, Exception | None]: return error def _record(self, response: httpx2.Response) -> None: + # The MCP client tolerates some error responses, such as one to a notification, so an + # error must not outlive the next success, or a later failure would be reported as it. if response.status_code >= 400: self.error_status = response.status_code self.retry_after = _parse_retry_after(response.headers.get("Retry-After")) - - def sync_auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: - try: - token = self.get_token() - except Exception as e: - self.token_error = e - raise - request.headers["Authorization"] = f"Bearer {token}" - response = yield request - self._record(response) + else: + self.error_status = self.retry_after = None async def async_auth_flow( self, request: httpx2.Request @@ -187,7 +181,7 @@ async def async_auth_flow( class DatabricksUnityMCPToolset(MCPToolset): """ - Give an agent the tools of a Unity AI Gateway MCP Service, authenticated as the connection's identity. + 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:///ai-gateway/mcp-services/``. The @@ -195,18 +189,20 @@ class DatabricksUnityMCPToolset(MCPToolset): neither a gateway URL nor a token, and the token is only ever sent to that workspace. 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), which needs - ``EXECUTE`` on the MCP Service and ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and schema. - It sees only the tools selected for the service, and the service's policies apply. + 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. Gateway errors are raised as :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPServiceNotFoundError`, - :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError` or - :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`. Tool calls - are never retried by this toolset, because a call interrupted after it was sent may already - have run. + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or, for any + other gateway error, :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`. + Tool calls are never retried by this toolset, because a call interrupted after it was sent may + already have run. .. code-block:: python @@ -280,7 +276,7 @@ def _translate_error(self, error: Exception, *, during_tool_call: bool) -> Excep ) if status in (401, 403): return DatabricksUnityMCPAccessDeniedError( - f"Unity AI Gateway denied access to MCP Service {service!r} (HTTP {status}). The " + f"Unity Gateway denied access to MCP Service {service!r} (HTTP {status}). The " f"identity of connection {self._databricks_conn_id!r} needs EXECUTE on the service and " "USE CATALOG and USE SCHEMA on its catalog and schema, and its credentials must be valid.", http_status_code=status, @@ -294,7 +290,7 @@ def _translate_error(self, error: Exception, *, during_tool_call: bool) -> Excep if status == 429: hint = f" Retry after {retry_after:g} seconds." if retry_after is not None else "" return DatabricksUnityMCPThrottledError( - f"Unity AI Gateway rate-limited calls to MCP Service {service!r} (HTTP 429).{hint}", + f"Unity Gateway rate-limited calls to MCP Service {service!r} (HTTP 429).{hint}", http_status_code=status, retry_after=retry_after, ) @@ -303,12 +299,12 @@ def _translate_error(self, error: Exception, *, during_tool_call: bool) -> Excep ) if status is not None: return DatabricksUnityMCPError( - f"Unity AI Gateway returned HTTP {status} for MCP Service {service!r}.{ambiguous}", + f"Unity Gateway returned HTTP {status} for MCP Service {service!r}.{ambiguous}", http_status_code=status, ) if (transport_error := _find_transport_error(error)) is not None: return DatabricksUnityMCPTransportError( - f"Could not reach Unity AI Gateway for MCP Service {service!r}: {transport_error}.{ambiguous}" + f"Could not reach Unity Gateway for MCP Service {service!r}: {transport_error}.{ambiguous}" ) return None diff --git a/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py index 49e10f0e072e2..02a5e5235db12 100644 --- a/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py +++ b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. """ -Example Dag giving an agent the tools of a Unity AI Gateway MCP Service. +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 diff --git a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py index a5329ab379be1..01a922c24b9d2 100644 --- a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py +++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py @@ -29,7 +29,7 @@ pytest.importorskip("fastmcp") pytest.importorskip("airflow.providers.common.ai") -from pydantic_ai import RunContext +from pydantic_ai import ModelRetry, RunContext from airflow.models import Connection from airflow.providers.common.ai.utils import toolset_base @@ -49,13 +49,14 @@ class _StubGateway(BaseHTTPRequestHandler): - """A Unity AI Gateway stand-in that serves one MCP Service with a single ``echo`` tool.""" + """A Unity Gateway stand-in that serves one MCP Service with a single ``echo`` tool.""" protocol_version = "HTTP/1.1" # Per-test behaviour, set by the ``gateway`` fixture: service name -> status for every request, - # and a status for ``tools/call`` requests only. + # a status for ``tools/call`` requests only, and a status for notifications only. service_status: dict[str, tuple[int, dict[str, str]]] = {} tool_call_status: int | None = None + notification_status: int | None = None requests: list[tuple[str, str | None, str | None]] = [] def do_POST(self): @@ -69,7 +70,7 @@ def do_POST(self): if method == "tools/call" and self.tool_call_status: return self._reply(self.tool_call_status) if "id" not in body: - return self._reply(202) + return self._reply(self.notification_status or 202) if method == "initialize": result = { "protocolVersion": body["params"]["protocolVersion"], @@ -87,6 +88,9 @@ def do_POST(self): ] } elif method == "tools/call": + if "text" not in body["params"]["arguments"]: + error = {"code": -32602, "message": "Invalid params: 'text' is required"} + return self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "error": error}) text = body["params"]["arguments"]["text"] result = {"content": [{"type": "text", "text": f"echo: {text}"}], "isError": False} else: @@ -118,6 +122,7 @@ def log_message(self, *args): def gateway(): _StubGateway.service_status = {} _StubGateway.tool_call_status = None + _StubGateway.notification_status = None _StubGateway.requests = [] server = ThreadingHTTPServer(("127.0.0.1", 0), _StubGateway) thread = threading.Thread(target=server.serve_forever, daemon=True) @@ -269,7 +274,7 @@ def test_username_and_password_connection_is_rejected(self, create_connection_wi Connection(conn_id=CONN_ID, conn_type="databricks", host="h", login="user", password="pass") ) - with pytest.raises(ValueError, match="bearer tokens only"): + 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) @@ -317,6 +322,20 @@ def test_server_error_during_a_tool_call_says_the_call_may_have_run(self, gatewa # The failed call was sent once and not retried. assert [m for _, _, m in _StubGateway.requests].count("tools/call") == 1 + 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()) + def test_unreachable_gateway_is_a_transport_error(self, create_connection_without_db): create_connection_without_db( Connection( @@ -324,7 +343,7 @@ def test_unreachable_gateway_is_a_transport_error(self, create_connection_withou ) ) - with pytest.raises(DatabricksUnityMCPTransportError, match="Could not reach Unity AI Gateway"): + with pytest.raises(DatabricksUnityMCPTransportError, match="Could not reach Unity Gateway"): _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) From 6402af7f7edb0009995a800b2e23f8af3dea3025 Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sat, 10 Oct 2026 23:27:03 +0000 Subject: [PATCH 6/8] Judge each Unity MCP failure by its own error, not shared state Tool calls run in parallel, so attributing the last HTTP error seen by the toolset to a failed call could report a server error as a retry the model should make, re-running a call that may already have changed data. A JSON-RPC error the server returns at an HTTP error status, by contrast, is its answer to the call and something the model can correct. Building on common.ai's MCPToolset meant overriding its private internals, which carry no stability guarantee; owning the MCP client keeps the toolset working as common.ai evolves. A missing service is answered with 403, never 404, so it is reported as access denied. Gateway requests now honour the connection's proxies and are only sent over HTTPS, so the bearer token cannot travel in clear text. --- docs/spelling_wordlist.txt | 1 + providers/databricks/docs/index.rst | 2 +- .../databricks/docs/toolsets/unity_mcp.rst | 42 ++- providers/databricks/provider.yaml | 2 - providers/databricks/pyproject.toml | 2 +- .../providers/databricks/exceptions.py | 8 +- .../providers/databricks/get_provider_info.py | 1 - .../databricks/toolsets/unity_mcp.py | 293 +++++++++++++----- .../databricks/toolsets/test_unity_mcp.py | 183 +++++++++-- 9 files changed, 398 insertions(+), 136 deletions(-) diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index c9ec654ce7168..5ffede0502ddc 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1078,6 +1078,7 @@ Maxime MaxRuntimeInSeconds mb MCP +mcp md mem memcached diff --git a/providers/databricks/docs/index.rst b/providers/databricks/docs/index.rst index 37823f8803daf..450c25cbf94b6 100644 --- a/providers/databricks/docs/index.rst +++ b/providers/databricks/docs/index.rst @@ -156,7 +156,7 @@ Extra Dependencies ``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]>=0.10.0`` +``common.ai`` ``apache-airflow-providers-common-ai[mcp]>=0.11.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/unity_mcp.rst b/providers/databricks/docs/toolsets/unity_mcp.rst index f2d2e690313f9..b8b84cfd12ace 100644 --- a/providers/databricks/docs/toolsets/unity_mcp.rst +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -31,8 +31,10 @@ The Dag names the service by its three-level Unity Catalog name, ``catalog.schem :ref:`Databricks connection ` to use. The toolset builds the service URL, ``https:///ai-gateway/mcp-services/``, 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. Service names may contain only ASCII letters, digits, ``_`` and ``-`` in each -part; any other name is rejected before a request is made. +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 @@ -51,9 +53,11 @@ 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. Account users already hold these privileges on the -built-in ``system.ai`` 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. +Unity Catalog connection behind the service. On the built-in ``system.ai`` services, account users +hold these privileges by default (see `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 @@ -67,15 +71,23 @@ authentication modes of the Databricks connection: * Azure AD: a service principal, a managed identity, or ``DefaultAzureCredential``; * workload identity federation (Kubernetes, AWS IAM, or a supplied token provider). -Username and password authentication is not supported. A token is fetched from the connection for +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 +`__). 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 ------------------ -Gateway errors are raised as dedicated exceptions from -:mod:`airflow.providers.databricks.exceptions`: +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 @@ -83,10 +95,9 @@ Gateway errors are raised as dedicated exceptions from * - Exception - Cause * - ``DatabricksUnityMCPAccessDeniedError`` - - HTTP 401 or 403: the credentials are invalid, or the identity lacks ``EXECUTE`` on the service - or ``USE CATALOG`` / ``USE SCHEMA`` on its parents. - * - ``DatabricksUnityMCPServiceNotFoundError`` - - HTTP 404: the service does not exist on the workspace, or is not visible to the identity. + - 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 in seconds, when given. * - ``DatabricksUnityMCPTransportError`` @@ -94,6 +105,7 @@ Gateway errors are raised as dedicated exceptions from * - ``DatabricksUnityMCPError`` - Any other gateway error, or a failure to get a token from the connection. -The toolset never retries a tool call. A call that fails after it was sent, with a server or transport -error, may or may not have run, and repeating a tool that changes data could apply the change twice; -the error message says so. Use task retries for calls that are safe to repeat. +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 the error is reported as ``DatabricksUnityMCPError`` without it. Use task retries for +calls that are safe to repeat. diff --git a/providers/databricks/provider.yaml b/providers/databricks/provider.yaml index 36c8e3202c1f4..becb146b7f1f1 100644 --- a/providers/databricks/provider.yaml +++ b/providers/databricks/provider.yaml @@ -146,8 +146,6 @@ integrations: tags: [service] - integration-name: Databricks Unity Gateway external-doc-url: https://docs.databricks.com/aws/en/unity-gateway/concepts - how-to-guide: - - /docs/apache-airflow-providers-databricks/toolsets/unity_mcp.rst tags: [service] operators: diff --git a/providers/databricks/pyproject.toml b/providers/databricks/pyproject.toml index 42cb536a25805..f89476a76b235 100644 --- a/providers/databricks/pyproject.toml +++ b/providers/databricks/pyproject.toml @@ -89,7 +89,7 @@ dependencies = [ "azure-identity>=1.25.3", ] "common.ai" = [ - "apache-airflow-providers-common-ai[mcp]>=0.10.0", # use next version + "apache-airflow-providers-common-ai[mcp]>=0.11.0", # use next version ] "fab" = [ "apache-airflow-providers-fab>=2.2.0" diff --git a/providers/databricks/src/airflow/providers/databricks/exceptions.py b/providers/databricks/src/airflow/providers/databricks/exceptions.py index 5962b82dc21d7..75487794e9151 100644 --- a/providers/databricks/src/airflow/providers/databricks/exceptions.py +++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py @@ -53,11 +53,11 @@ class DatabricksUnityMCPError(DatabricksApiError): class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError): - """Raised when the gateway rejects the caller's credentials or the caller lacks a privilege on the service.""" - + """ + Raised when the gateway will not let the caller invoke the service. -class DatabricksUnityMCPServiceNotFoundError(DatabricksUnityMCPError): - """Raised when the MCP Service does not exist, or is not visible to the caller.""" + The service may not exist, the caller may lack a privilege on it, or its credentials may be invalid. + """ class DatabricksUnityMCPThrottledError(DatabricksUnityMCPError): 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 59229f0fab65c..7b384389452a2 100644 --- a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py +++ b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py @@ -70,7 +70,6 @@ def get_provider_info(): { "integration-name": "Databricks Unity Gateway", "external-doc-url": "https://docs.databricks.com/aws/en/unity-gateway/concepts", - "how-to-guide": ["/docs/apache-airflow-providers-databricks/toolsets/unity_mcp.rst"], "tags": ["service"], }, ], diff --git a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py index 2de62016800ac..13039b46943f4 100644 --- a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -20,26 +20,31 @@ import asyncio import contextlib +import ipaddress +import json 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 from airflow.providers.databricks.exceptions import ( DatabricksUnityMCPAccessDeniedError, DatabricksUnityMCPError, - DatabricksUnityMCPServiceNotFoundError, DatabricksUnityMCPThrottledError, DatabricksUnityMCPTransportError, ) from airflow.providers.databricks.hooks.databricks import DatabricksHook 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.toolsets.mcp import MCPToolset + 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: " @@ -52,6 +57,7 @@ from pydantic_ai._run_context import RunContext from pydantic_ai.toolsets.abstract import ToolsetTool + from airflow.providers.common.ai.tools import AirflowTool from airflow.sdk.execution_time.secrets_masker import mask_secret else: try: @@ -66,23 +72,54 @@ # 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_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"), +} +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream ended without a response", + "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, underscores and hyphens. + 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 '-'." + "'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 @@ -94,16 +131,16 @@ def _parse_retry_after(value: str | None) -> float | None: return seconds if seconds >= 0 else None -def _find_transport_error(exc: BaseException) -> httpx2.TransportError | None: - """Return the network error behind ``exc``, looking through causes and exception groups.""" +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.""" seen: set[int] = set() pending: list[BaseException] = [exc] while pending: - current = pending.pop() + current = pending.pop(0) if id(current) in seen: continue seen.add(id(current)) - if isinstance(current, httpx2.TransportError): + if isinstance(current, types): return 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)): @@ -112,24 +149,44 @@ def _find_transport_error(exc: BaseException) -> httpx2.TransportError | None: return None +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 errors are only + attributed to the operation when no other operation of the toolset ran at the same time. + """ + + overlapped: bool = False + errors: list[tuple[int, float | None]] = 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. - - The MCP client reports every HTTP error from the gateway as the same generic error, so the - status and ``Retry-After`` of the last error response, and any failure to get a token, are - kept here for :class:`DatabricksUnityMCPToolset` to report what went wrong. With concurrent - calls on one toolset, the error reported for a failed call can be another call's. """ - def __init__(self, hook: DatabricksHook) -> None: + def __init__(self, hook: DatabricksHook, operations: set[_Operation]) -> None: self._hook = hook - self.error_status: int | None = None - self.retry_after: float | None = None - self.token_error: Exception | None = None + self._operations = operations self._token_lock = threading.Lock() def get_token(self) -> str: @@ -147,46 +204,49 @@ def get_token(self) -> str: mask_secret(token) return token - def take_error(self) -> tuple[int | None, float | None, Exception | None]: - """Return the last recorded error and forget it.""" - error = (self.error_status, self.retry_after, self.token_error) - self.error_status = self.retry_after = self.token_error = None - return error - - def _record(self, response: httpx2.Response) -> None: - # The MCP client tolerates some error responses, such as one to a notification, so an - # error must not outlive the next success, or a later failure would be reported as it. - if response.status_code >= 400: - self.error_status = response.status_code - self.retry_after = _parse_retry_after(response.headers.get("Retry-After")) - else: - self.error_status = self.retry_after = None - async def async_auth_flow( self, request: httpx2.Request ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: try: # 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 connection is already - # resolved by _get_server, so fetching a token only calls the token endpoint, and the - # hook is this toolset's own, so a lock of its own is enough. + # resolved when the toolset connects, so fetching a token only calls the token + # endpoint, and the hook is this toolset's own, so a lock of its own is enough. token = await asyncio.to_thread(self.get_token) except Exception as e: - self.token_error = e - raise + raise _TokenUnavailableError(str(e)) from e request.headers["Authorization"] = f"Bearer {token}" response = yield request - self._record(response) + if response.status_code >= 400 and _is_jsonrpc_request(request): + retry_after = _parse_retry_after(response.headers.get("Retry-After")) + for operation in self._operations: + operation.errors.append((response.status_code, retry_after)) -class DatabricksUnityMCPToolset(MCPToolset): +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:///ai-gateway/mcp-services/``. 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. + 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 @@ -195,14 +255,14 @@ class DatabricksUnityMCPToolset(MCPToolset): 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. Gateway errors are raised as + 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.DatabricksUnityMCPServiceNotFoundError`, :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, - :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or, for any - other gateway error, :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`. - Tool calls are never retried by this toolset, because a call interrupted after it was sent may - already have run. + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`, as do failures + while connecting and listing tools. .. code-block:: python @@ -223,6 +283,7 @@ class DatabricksUnityMCPToolset(MCPToolset): :param tool_prefix: Optional prefix prepended to tool names. """ + # Rendered, on a copy, by AgentOperator. agent_template_fields: Sequence[str] = ("_databricks_conn_id", "_service_name") def __init__( @@ -232,11 +293,12 @@ def __init__( databricks_conn_id: str = DatabricksHook.default_conn_name, tool_prefix: str | None = None, ) -> None: - super().__init__(databricks_conn_id, tool_prefix=tool_prefix) self._databricks_conn_id = databricks_conn_id self._service_name = service_name - self._auth: _DatabricksTokenAuth | None = None - # A templated name is checked once it has been rendered, in _get_server. + self._tool_prefix = tool_prefix + self._server: Any = None + self._operations: set[_Operation] = set() + # A templated name is checked once it has been rendered, when the toolset connects. if "{{" not in service_name: validate_service_name(service_name) @@ -249,42 +311,56 @@ def get_service_url(self, hook: DatabricksHook) -> str: validate_service_name(self._service_name) if not hook.host: raise ValueError(f"Connection {self._databricks_conn_id!r} has no workspace host.") - return hook._endpoint_url(f"{GATEWAY_MCP_SERVICES_PATH}/{self._service_name}") + 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, self._operations) + # Fail here, with a clear message, rather than inside the MCP client, which reports + # any failure as a generic connection error. + auth.get_token() + proxy = (hook.proxies or {}).get(urlsplit(url).scheme) + transport = StreamableHttpTransport( + url, + headers=hook.user_agent_header, + auth=auth, + httpx_client_factory=_build_proxied_client_factory(proxy) if proxy else None, + ) + toolset = PydanticAIMCPToolset(transport) + return toolset.prefixed(self._tool_prefix) if self._tool_prefix else toolset + + async def _resolve_server(self) -> Any: if self._server is None: - 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. - auth.get_token() - transport = StreamableHttpTransport(url, headers=hook.user_agent_header, auth=auth) - toolset = PydanticAIMCPToolset(transport) - self._auth = auth - self._server = toolset.prefixed(self._tool_prefix) if self._tool_prefix else toolset + # 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, *, during_tool_call: bool) -> Exception | None: + 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.""" - status, retry_after, token_error = self._auth.take_error() if self._auth else (None, None, None) service = self._service_name - if token_error is not None: + if (token_error := _find(error, _TokenUnavailableError)) 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}" ) - if status in (401, 403): + mcp_error = _find(error, McpError) + 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} (HTTP {status}). The " - f"identity of connection {self._databricks_conn_id!r} needs EXECUTE on the service and " - "USE CATALOG and USE SCHEMA on its catalog and schema, and its credentials must be valid.", - http_status_code=status, - ) - if status == 404: - return DatabricksUnityMCPServiceNotFoundError( - f"MCP Service {service!r} was not found on the workspace of connection " - f"{self._databricks_conn_id!r}, or is not visible to its identity (HTTP 404).", + 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: @@ -297,35 +373,56 @@ def _translate_error(self, error: Exception, *, during_tool_call: bool) -> Excep ambiguous = ( " The tool call may or may not have run, so it was not retried." if during_tool_call else "" ) - if status is not None: + 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 HTTP {status} for MCP Service {service!r}.{ambiguous}", + f"Unity Gateway returned {what} for MCP Service {service!r}.{ambiguous}", http_status_code=status, ) - if (transport_error := _find_transport_error(error)) is not None: + if (transport_error := _find(error, httpx2.TransportError)) is not None: return DatabricksUnityMCPTransportError( f"Could not reach Unity Gateway for MCP Service {service!r}: {transport_error}.{ambiguous}" ) return None @contextlib.asynccontextmanager - async def _gateway_errors(self, *, during_tool_call: bool = False) -> AsyncIterator[None]: + 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, during_tool_call=during_tool_call) + translated = self._translate_error(e, operation, during_tool_call=during_tool_call) if translated is None: raise raise translated from e + finally: + self._operations.discard(operation) async def __aenter__(self) -> DatabricksUnityMCPToolset: - async with self._gateway_errors(): - await super().__aenter__() + 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 not None: + return await self._server.__aexit__(*args) + return None + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: - async with self._gateway_errors(): - return await super().get_tools(ctx) + async with self._gateway_operation(): + return await (await self._resolve_server()).get_tools(ctx) async def execute_tool( self, @@ -335,5 +432,33 @@ async def execute_tool( ctx: RunContext[Any], tool: ToolsetTool[Any], ) -> Any: - async with self._gateway_errors(during_tool_call=True): - return await super().execute_tool(name, tool_args, ctx=ctx, tool=tool) + 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_proxied_client_factory(proxy: str) -> Any: + """Return an MCP HTTP client factory whose clients send requests through ``proxy``.""" + + 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) + return httpx2.AsyncClient(headers=headers, timeout=timeout, auth=auth, proxy=proxy, **kwargs) + + return factory diff --git a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py index 01a922c24b9d2..aaa73a81bc413 100644 --- a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py +++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py @@ -37,7 +37,6 @@ from airflow.providers.databricks.exceptions import ( DatabricksUnityMCPAccessDeniedError, DatabricksUnityMCPError, - DatabricksUnityMCPServiceNotFoundError, DatabricksUnityMCPThrottledError, DatabricksUnityMCPTransportError, ) @@ -47,16 +46,26 @@ 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 Unity Gateway stand-in that serves one MCP Service with a single ``echo`` tool. + + A call of ``echo`` whose text is ``"reply "`` gets that HTTP status with an empty body, + and ``"reply "`` gets it with the JSON-RPC error named by ````. + """ protocol_version = "HTTP/1.1" - # Per-test behaviour, set by the ``gateway`` fixture: service name -> status for every request, - # a status for ``tools/call`` requests only, and a status for notifications only. - service_status: dict[str, tuple[int, dict[str, str]]] = {} - tool_call_status: int | None = None + 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 requests: list[tuple[str, str | None, str | None]] = [] def do_POST(self): @@ -64,11 +73,9 @@ def do_POST(self): 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_status: - status, headers = self.service_status[service] - return self._reply(status, headers=headers) - if method == "tools/call" and self.tool_call_status: - return self._reply(self.tool_call_status) + 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) if method == "initialize": @@ -88,10 +95,14 @@ def do_POST(self): ] } elif method == "tools/call": + if self.tool_call_barrier: + self.tool_call_barrier.wait() if "text" not in body["params"]["arguments"]: - error = {"code": -32602, "message": "Invalid params: 'text' is required"} - return self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "error": error}) + return self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "error": INVALID_PARAMS}) text = body["params"]["arguments"]["text"] + 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 = {} @@ -103,6 +114,10 @@ def do_GET(self): 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) @@ -120,9 +135,9 @@ def log_message(self, *args): @pytest.fixture def gateway(): - _StubGateway.service_status = {} - _StubGateway.tool_call_status = None + _StubGateway.service_reply = {} _StubGateway.notification_status = None + _StubGateway.tool_call_barrier = None _StubGateway.requests = [] server = ThreadingHTTPServer(("127.0.0.1", 0), _StubGateway) thread = threading.Thread(target=server.serve_forever, daemon=True) @@ -163,6 +178,17 @@ async def _list_and_call(toolset, text="hi"): return tools, result +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", @@ -173,6 +199,7 @@ class TestServiceName: "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", "", @@ -229,6 +256,51 @@ def test_connection_without_host_is_rejected(self, create_connection_without_db) 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): @@ -290,38 +362,87 @@ def test_token_failure_during_the_run_is_reported(self, get_token, gateway, gate class TestErrors: @pytest.mark.parametrize( - ("status", "headers", "expected", "retry_after"), + ("status", "headers", "error", "expected", "retry_after"), [ - (401, {}, DatabricksUnityMCPAccessDeniedError, None), - (403, {}, DatabricksUnityMCPAccessDeniedError, None), - (404, {}, DatabricksUnityMCPServiceNotFoundError, None), - (429, {"Retry-After": "12"}, DatabricksUnityMCPThrottledError, 12.0), - (429, {"Retry-After": "Wed, 21 Oct 2026 07:28:00 GMT"}, DatabricksUnityMCPThrottledError, None), + (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": "Wed, 21 Oct 2026 07:28:00 GMT"}, + None, + DatabricksUnityMCPThrottledError, + None, + ), ], ) - def test_gateway_errors_are_translated( - self, gateway, gateway_conn, status, headers, expected, retry_after + def test_gateway_errors_while_connecting_are_translated( + self, gateway, gateway_conn, status, headers, error, expected, retry_after ): - _StubGateway.service_status[SERVICE] = (status, headers) + _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 if expected is DatabricksUnityMCPThrottledError: assert err.value.retry_after == retry_after - def test_server_error_during_a_tool_call_says_the_call_may_have_run(self, gateway, gateway_conn): - _StubGateway.tool_call_status = 503 + def test_missing_service_is_reported_as_access_denied(self, gateway, gateway_conn): + _StubGateway.service_reply[SERVICE] = (403, {}, "not_authorized") - with pytest.raises(DatabricksUnityMCPError, match="may or may not have run") as err: + 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 + 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. @@ -347,10 +468,16 @@ def test_unreachable_gateway_is_a_transport_error(self, create_connection_withou _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) +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.toolsets.mcp": None}): + 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) From 51e27b39b784cc6cfc9778bc99a843202db0ac9c Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:06:04 +0000 Subject: [PATCH 7/8] Install the MCP client for Databricks provider development Without the mcp extra of common.ai, uv sync left out fastmcp and the Unity MCP toolset tests skipped silently. The extra sits in the hand-maintained part of the dev group because the generated part is rewritten without extras. --- providers/databricks/pyproject.toml | 1 + uv.lock | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/providers/databricks/pyproject.toml b/providers/databricks/pyproject.toml index f89476a76b235..6df8b3152b96e 100644 --- a/providers/databricks/pyproject.toml +++ b/providers/databricks/pyproject.toml @@ -126,6 +126,7 @@ dev = [ "deltalake>=1.1.3,!=1.3.0", "apache-airflow-providers-fab>=2.2.0", "apache-airflow-providers-microsoft-azure", + "apache-airflow-providers-common-ai[mcp]", "apache-airflow-providers-common-sql[pandas,polars]", "apache-airflow-providers-fab", "apache-airflow-providers-databricks[sqlalchemy]", diff --git a/uv.lock b/uv.lock index 586f624f0d7ac..907ab80f6a456 100644 --- a/uv.lock +++ b/uv.lock @@ -5017,6 +5017,9 @@ avro = [ azure-identity = [ { name = "azure-identity" }, ] +common-ai = [ + { name = "apache-airflow-providers-common-ai", extra = ["mcp"] }, +] fab = [ { name = "apache-airflow-providers-fab" }, ] @@ -5041,6 +5044,7 @@ dev = [ { name = "apache-airflow" }, { name = "apache-airflow-devel-common" }, { name = "apache-airflow-providers-amazon" }, + { name = "apache-airflow-providers-common-ai", extra = ["mcp"] }, { name = "apache-airflow-providers-common-compat" }, { name = "apache-airflow-providers-common-sql", extra = ["pandas", "polars"] }, { name = "apache-airflow-providers-databricks", extra = ["sqlalchemy"] }, @@ -5060,6 +5064,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", 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" }, @@ -5082,13 +5087,15 @@ requires-dist = [ { name = "pyarrow", marker = "python_full_version >= '3.14'", specifier = ">=22.0.0" }, { name = "requests", specifier = ">=2.32.0,<3" }, ] -provides-extras = ["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 = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-devel-common", editable = "devel-common" }, { name = "apache-airflow-providers-amazon", editable = "providers/amazon" }, + { name = "apache-airflow-providers-common-ai", editable = "providers/common/ai" }, + { name = "apache-airflow-providers-common-ai", extras = ["mcp"], 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-common-sql", extras = ["pandas", "polars"], editable = "providers/common/sql" }, From 47c8d439af8b88db713706f616b929cb4907fdbd Mon Sep 17 00:00:00 2001 From: firasbouzazi <115628206+firasbouzazi@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:12:44 +0000 Subject: [PATCH 8/8] Let Unity MCP connections wait for slow MCP Services Unity Gateway forwards a response only once the MCP server has finished sending it, so a server that keeps its response stream open delays the MCP handshake by as long as it does. Behind the gateway that exceeded pydantic-ai's 5 second default, and the connection failed with a bare RuntimeError that gave no hint of the cause. --- .../databricks/docs/toolsets/unity_mcp.rst | 7 ++- .../databricks/toolsets/unity_mcp.py | 44 ++++++++++++++++--- .../databricks/toolsets/test_unity_mcp.py | 34 ++++++++++++++ 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/providers/databricks/docs/toolsets/unity_mcp.rst b/providers/databricks/docs/toolsets/unity_mcp.rst index b8b84cfd12ace..43f243b1089da 100644 --- a/providers/databricks/docs/toolsets/unity_mcp.rst +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -101,10 +101,15 @@ which are also raised for failures while connecting and listing tools: * - ``DatabricksUnityMCPThrottledError`` - HTTP 429. Its ``retry_after`` attribute holds the ``Retry-After`` delay in seconds, when given. * - ``DatabricksUnityMCPTransportError`` - - The gateway could not be reached, or the connection dropped. + - The gateway could not be reached, the connection dropped, 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. +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 the error is reported as ``DatabricksUnityMCPError`` without it. Use task retries for diff --git a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py index 13039b46943f4..be929beaca400 100644 --- a/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -52,7 +52,7 @@ ) from e if TYPE_CHECKING: - from collections.abc import AsyncGenerator, AsyncIterator, Sequence + from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Sequence from pydantic_ai._run_context import RunContext from pydantic_ai.toolsets.abstract import ToolsetTool @@ -88,6 +88,9 @@ (-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" _CLIENT_ERROR_PREFIXES = ( "Unexpected content type:", "Failed to parse JSON response:", @@ -131,8 +134,8 @@ def _parse_retry_after(value: str | None) -> float | None: return seconds if seconds >= 0 else None -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.""" +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: @@ -140,13 +143,24 @@ def _find(exc: BaseException, types: type[BaseException] | tuple[type[BaseExcept if id(current) in seen: continue seen.add(id(current)) - if isinstance(current, types): - return 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) - return None + + +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_client_error(error: McpError) -> bool: @@ -281,6 +295,9 @@ class DatabricksUnityMCPToolset(AirflowToolset): :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. @@ -292,10 +309,14 @@ def __init__( *, 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() # A templated name is checked once it has been rendered, when the toolset connects. @@ -334,7 +355,7 @@ def _get_server(self) -> Any: auth=auth, httpx_client_factory=_build_proxied_client_factory(proxy) if proxy else None, ) - toolset = PydanticAIMCPToolset(transport) + 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: @@ -392,6 +413,15 @@ def _translate_error( return DatabricksUnityMCPTransportError( f"Could not reach Unity Gateway for MCP Service {service!r}: {transport_error}.{ambiguous}" ) + if _is_handshake_timeout(error): + if during_tool_call: + return DatabricksUnityMCPTransportError( + f"Timed out waiting for Unity Gateway to answer MCP Service {service!r}.{ambiguous}" + ) + 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." + ) return None @contextlib.asynccontextmanager diff --git a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py index aaa73a81bc413..8c80cd96fcc13 100644 --- a/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py +++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py @@ -66,6 +66,9 @@ class _StubGateway(BaseHTTPRequestHandler): 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 ``initialize``; ``release`` ends any wait early on teardown. + initialize_delay: float = 0 + release = threading.Event() requests: list[tuple[str, str | None, str | None]] = [] def do_POST(self): @@ -79,6 +82,7 @@ def do_POST(self): if "id" not in body: return self._reply(self.notification_status or 202) if method == "initialize": + self.release.wait(self.initialize_delay) result = { "protocolVersion": body["params"]["protocolVersion"], "capabilities": {"tools": {}}, @@ -138,11 +142,14 @@ def gateway(): _StubGateway.service_reply = {} _StubGateway.notification_status = None _StubGateway.tool_call_barrier = None + _StubGateway.initialize_delay = 0 + _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() @@ -468,6 +475,33 @@ def test_unreachable_gateway_is_a_transport_error(self, create_connection_withou _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) +class TestHandshakeTimeout: + def test_slow_handshake_within_init_timeout_connects(self, gateway, gateway_conn): + _StubGateway.initialize_delay = 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.initialize_delay = 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 TestFrameworkNeutralTools: def test_are_not_supported(self): with pytest.raises(NotImplementedError, match="framework-neutral"):