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 7aed991add904..450c25cbf94b6 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.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/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..43f243b1089da --- /dev/null +++ b/providers/databricks/docs/toolsets/unity_mcp.rst @@ -0,0 +1,116 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/toolset:DatabricksUnityMCPToolset: + +Unity Gateway MCP Services +========================== + +Use :class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset` to give an +agent the tools of a Unity Gateway MCP Service: either one Databricks provides for workspace tools +and SaaS applications, such as ``system.ai.google_calendar``, or an external MCP server registered as +an MCP Service in Unity Catalog. The toolset works with +:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, ``@task.agent``, and the LangChain +bridge of the Common AI provider. + +The Dag names the service by its three-level Unity Catalog name, ``catalog.schema.service``, and the +:ref:`Databricks connection ` 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, over HTTPS. A connection whose ``schema`` is ``http`` is rejected unless its +host is a loopback address. Requests go through the proxy in the connection's ``proxies`` extra, if +set. Service names may contain only ASCII letters, digits and ``_`` in each part; any other name is +rejected before a request is made. + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_unity_mcp.py + :language: python + :start-after: [START howto_toolset_databricks_unity_mcp] + :end-before: [END howto_toolset_databricks_unity_mcp] + +The service name and connection ID are templated when the toolset is passed to ``AgentOperator`` or +``@task.agent``. + +Caller identity and authentication +---------------------------------- + +MCP Services need a workspace enabled for Unity Catalog, in a region where Model Serving is +supported. + +The gateway runs every tool call as the identity of the connection's credentials. That identity needs +``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and schema +(``EXECUTE`` alone is not enough), and an assignment to the workspace. It needs no privilege on the +Unity Catalog connection behind the service. On the built-in ``system.ai`` services, account users +hold these privileges by default (see `MCP Services +`__). The gateway exposes only the +tools selected for the service, and the service's policies apply. Grant the identity only the services +the agent should use. + +Built-in services that act on a user's own data, such as ``system.ai.google_calendar`` or +``system.ai.gmail``, need that identity to complete a one-time OAuth login first, for example by +opening the service in Catalog Explorer and clicking **Login**. + +The toolset sends the connection's token as a bearer token, so the connection must use one of these +authentication modes of the Databricks connection: + +* a personal access token (the identity is the token's user or service principal); +* service principal OAuth (``service_principal_oauth``); +* Azure AD: a service principal, a managed identity, or ``DefaultAzureCredential``; +* workload identity federation (Kubernetes, AWS IAM, or a supplied token provider). + +For service principal OAuth, the OAuth secret must allow the ``all-apis`` scope, which the connection +requests; a secret restricted to narrower scopes fails (see `OAuth for service principals +`__). Username and password +authentication is not supported. A token is fetched from the connection for +every request, so OAuth and Azure AD tokens are refreshed during a long agent run and when the toolset +reconnects. Tokens minted for the toolset are masked in task logs. + +Errors and retries +------------------ + +When the MCP server answers a tool call with an error, such as invalid arguments, the error goes to the +model, which can correct the call and try again. + +When the gateway gives no answer of its own, a tool call that was already sent may or may not have +run, and repeating a tool that changes data could apply the change twice. The toolset never retries +it: the task fails with one of these exceptions from :mod:`airflow.providers.databricks.exceptions`, +which are also raised for failures while connecting and listing tools: + +.. list-table:: + :header-rows: 1 + + * - Exception + - Cause + * - ``DatabricksUnityMCPAccessDeniedError`` + - The gateway will not let the identity invoke the service: no service has that name, the + identity lacks ``EXECUTE`` on the service or ``USE CATALOG`` / ``USE SCHEMA`` on its parents, + or the credentials are invalid. A missing service is reported like a missing privilege. + * - ``DatabricksUnityMCPThrottledError`` + - HTTP 429. Its ``retry_after`` attribute holds the ``Retry-After`` delay in seconds, when given. + * - ``DatabricksUnityMCPTransportError`` + - 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 +calls that are safe to repeat. diff --git a/providers/databricks/provider.yaml b/providers/databricks/provider.yaml index 2b66f546bf2a6..becb146b7f1f1 100644 --- a/providers/databricks/provider.yaml +++ b/providers/databricks/provider.yaml @@ -144,6 +144,9 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-databricks/operators/workflow.rst tags: [service] + - integration-name: Databricks Unity Gateway + external-doc-url: https://docs.databricks.com/aws/en/unity-gateway/concepts + tags: [service] operators: - integration-name: Databricks @@ -195,6 +198,11 @@ sensors: - airflow.providers.databricks.sensors.databricks_sql - airflow.providers.databricks.sensors.databricks_partition +toolsets: + - integration-name: Databricks Unity Gateway + python-modules: + - airflow.providers.databricks.toolsets.unity_mcp + connection-types: - hook-class-name: airflow.providers.databricks.hooks.databricks.DatabricksHook hook-name: "Databricks" diff --git a/providers/databricks/pyproject.toml b/providers/databricks/pyproject.toml index 28910f667f5a8..6df8b3152b96e 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.11.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", @@ -122,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/providers/databricks/src/airflow/providers/databricks/exceptions.py b/providers/databricks/src/airflow/providers/databricks/exceptions.py index a0f52a6077c06..75487794e9151 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 Gateway MCP Service fails.""" + + +class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError): + """ + Raised when the gateway will not let the caller invoke the service. + + The service may not exist, the caller may lack a privilege on it, or its credentials may be invalid. + """ + + +class DatabricksUnityMCPThrottledError(DatabricksUnityMCPError): + """Raised when the gateway rate-limits the caller.""" + + def __init__( + self, message: str, *, http_status_code: int | None = None, retry_after: float | None = None + ) -> None: + super().__init__(message, http_status_code=http_status_code) + self.retry_after = retry_after + + +class DatabricksUnityMCPTransportError(DatabricksUnityMCPError): + """ + Raised when the gateway cannot be reached or the connection drops. + + When this interrupts a tool call, the tool may or may not have run, so the call is not + retried automatically. + """ diff --git a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py index a218ce3546417..7b384389452a2 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,11 @@ def get_provider_info(): "how-to-guide": ["/docs/apache-airflow-providers-databricks/operators/workflow.rst"], "tags": ["service"], }, + { + "integration-name": "Databricks Unity Gateway", + "external-doc-url": "https://docs.databricks.com/aws/en/unity-gateway/concepts", + "tags": ["service"], + }, ], "operators": [ { diff --git a/providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py b/providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py new file mode 100644 index 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..be929beaca400 --- /dev/null +++ b/providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py @@ -0,0 +1,494 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import 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, + 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.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, Iterator, Sequence + + 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: + 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}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +_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 and underscores. + """ + if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): + raise ValueError( + f"Invalid Unity Gateway MCP Service name {service_name!r}: expected " + "'catalog.schema.service', each part made of ASCII letters, digits or '_'." + ) + + +def _is_loopback(host: str) -> bool: + if host == "localhost": + return True + try: + return ipaddress.ip_address(host.strip("[]")).is_loopback + except ValueError: + return False + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + try: + seconds = float(value) + except ValueError: + # The HTTP-date form is not worth parsing: the caller only uses this as a hint. + return None + return seconds if seconds >= 0 else None + + +def _walk_exceptions(exc: BaseException) -> Iterator[BaseException]: + """Yield ``exc``, its causes and the exceptions inside any exception groups, each once.""" + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop(0) + if id(current) in seen: + continue + seen.add(id(current)) + yield current + # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio uses the backport on 3.10. + if isinstance(grouped := getattr(current, "exceptions", None), (list, tuple)): + pending.extend(e for e in grouped if isinstance(e, BaseException)) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + + +def _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: + 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. + """ + + def __init__(self, hook: DatabricksHook, operations: set[_Operation]) -> None: + self._hook = hook + self._operations = operations + self._token_lock = threading.Lock() + + def get_token(self) -> str: + with self._token_lock: + token = self._hook._get_token(raise_error=False) + if not token: + raise ValueError( + f"Connection {self._hook.databricks_conn_id!r} has no token-based authentication " + "configured. This toolset sends a bearer token: use a personal access token, " + "service principal OAuth, Azure AD, or workload identity federation. Username and " + "password authentication is not supported." + ) + # 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 + + 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 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: + raise _TokenUnavailableError(str(e)) from e + request.headers["Authorization"] = f"Bearer {token}" + response = yield request + 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)) + + +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, over + HTTPS. The connection's ``proxies`` extra applies to gateway requests. + + The gateway runs every tool call as the identity of the connection's credentials (a user's + personal access token, a service principal, or an Azure AD / federated identity). That identity + needs ``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and + schema, and an assignment to the workspace. It sees only the tools selected for the service, and + the service's policies apply. + + Tokens are fetched from the connection for each request, so OAuth tokens refresh during a long + agent run and on reconnection. An error the server returns for a tool call reaches the model, + which can correct its arguments and try again. When the gateway gives no answer of its own, a + tool call may or may not have run, so it is never retried; it fails the task with + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`, as do failures + while connecting and listing tools. + + .. code-block:: python + + from airflow.providers.common.ai.operators.agent import AgentOperator + from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + + AgentOperator( + task_id="ask_mcp_service", + prompt="Which tools do you have?", + llm_conn_id="pydanticai_default", + toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp", databricks_conn_id="databricks")], + ) + + :param service_name: Three-level name of the MCP Service, ``catalog.schema.service``. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param databricks_conn_id: Databricks connection whose host and credentials are used. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param tool_prefix: Optional prefix prepended to tool names. + :param init_timeout: Seconds to wait for the MCP handshake when connecting. Unity Gateway forwards + a response only once the MCP server has finished sending it, so a server that keeps its + response stream open needs more than pydantic-ai's default of 5 seconds. + """ + + # Rendered, on a copy, by AgentOperator. + agent_template_fields: Sequence[str] = ("_databricks_conn_id", "_service_name") + + def __init__( + self, + service_name: str, + *, + databricks_conn_id: str = DatabricksHook.default_conn_name, + tool_prefix: str | None = None, + init_timeout: float = 30, + ) -> None: + if not init_timeout > 0: + raise ValueError(f"init_timeout must be a positive number of seconds, got {init_timeout!r}.") + self._databricks_conn_id = databricks_conn_id + self._service_name = service_name + self._tool_prefix = tool_prefix + self._init_timeout = init_timeout + self._server: Any = None + self._operations: set[_Operation] = set() + # A templated name is checked once it has been rendered, when the toolset connects. + if "{{" not in service_name: + validate_service_name(service_name) + + @property + def id(self) -> str: + return f"databricks-unity-mcp-{self._databricks_conn_id}-{self._service_name}" + + def get_service_url(self, hook: DatabricksHook) -> str: + """Return the gateway URL of the MCP Service on the connection's workspace.""" + validate_service_name(self._service_name) + if not hook.host: + raise ValueError(f"Connection {self._databricks_conn_id!r} has no workspace host.") + url = hook._endpoint_url(f"{GATEWAY_MCP_SERVICES_PATH}/{self._service_name}") + parts = urlsplit(url) + if parts.scheme != "https" and not _is_loopback(parts.hostname or ""): + raise ValueError( + f"Connection {self._databricks_conn_id!r} uses the {parts.scheme!r} scheme. Unity Gateway " + "requests carry a bearer token, so they are only sent over HTTPS." + ) + return url + + def _get_server(self) -> Any: + hook = DatabricksHook(self._databricks_conn_id, caller=type(self).__name__) + url = self.get_service_url(hook) + auth = _DatabricksTokenAuth(hook, 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, init_timeout=self._init_timeout) + return toolset.prefixed(self._tool_prefix) if self._tool_prefix else toolset + + async def _resolve_server(self) -> Any: + if self._server is None: + # Resolving the connection talks to the supervisor, so it takes the blocking-call lock. + self._server = await self.run_blocking(self._get_server) + return self._server + + def _translate_error( + self, error: Exception, operation: _Operation, *, during_tool_call: bool + ) -> Exception | None: + """Return the provider exception for a gateway failure, or ``None`` to re-raise ``error`` as is.""" + service = self._service_name + 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}" + ) + 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}" + f"{f' (HTTP {status})' if status else ''}. Either no service has that name, or the " + f"identity of connection {self._databricks_conn_id!r} lacks EXECUTE on the service or " + "USE CATALOG and USE SCHEMA on its catalog and schema, or its credentials are not valid.", + http_status_code=status, + ) + if status == 429: + hint = f" Retry after {retry_after:g} seconds." if retry_after is not None else "" + return DatabricksUnityMCPThrottledError( + f"Unity Gateway rate-limited calls to MCP Service {service!r} (HTTP 429).{hint}", + http_status_code=status, + retry_after=retry_after, + ) + ambiguous = ( + " The tool call may or may not have run, so it was not retried." if during_tool_call else "" + ) + if mcp_error is not None: + if not _is_client_error(mcp_error): + # The server's own answer: for a tool call, it reaches the model as a retry. + if during_tool_call: + return None + return DatabricksUnityMCPError( + f"MCP Service {service!r} returned an error: {mcp_error.error.message} " + f"(code {mcp_error.error.code}).", + http_status_code=status, + ) + what = f"HTTP {status}" if status else f"an error without an answer ({mcp_error.error.message})" + return DatabricksUnityMCPError( + f"Unity Gateway returned {what} for MCP Service {service!r}.{ambiguous}", + http_status_code=status, + ) + 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}" + ) + 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 + async def _gateway_operation(self, *, during_tool_call: bool = False) -> AsyncIterator[None]: + operation = _Operation(overlapped=bool(self._operations)) + for other in self._operations: + other.overlapped = True + self._operations.add(operation) + try: + yield + except Exception as e: + translated = self._translate_error(e, operation, during_tool_call=during_tool_call) + if translated is None: + raise + raise translated from e + finally: + self._operations.discard(operation) + + async def __aenter__(self) -> DatabricksUnityMCPToolset: + 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_operation(): + return await (await self._resolve_server()).get_tools(ctx) + + async def execute_tool( + self, + name: str, + tool_args: dict[str, Any], + *, + ctx: RunContext[Any], + tool: ToolsetTool[Any], + ) -> Any: + async with self._gateway_operation(during_tool_call=True): + return await (await self._resolve_server()).call_tool(name, tool_args, ctx, tool) + + def airflow_tools(self) -> list[AirflowTool]: + """ + Not supported: use the agent framework's own MCP client instead. + + An MCP session belongs to the event loop that opened it, and the framework-neutral + tools run outside the Pydantic AI run that manages it, so every call would reconnect. + """ + raise NotImplementedError( + "DatabricksUnityMCPToolset works in Pydantic AI agents and through the LangChain bridge, " + "not through the framework-neutral tools." + ) + + +def _build_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/system/databricks/example_databricks_unity_mcp.py b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py new file mode 100644 index 0000000000000..02a5e5235db12 --- /dev/null +++ b/providers/databricks/tests/system/databricks/example_databricks_unity_mcp.py @@ -0,0 +1,69 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Example Dag giving an agent the tools of a Unity Gateway MCP Service. + +Needs a Databricks connection whose identity has EXECUTE on the MCP Service and USE CATALOG and +USE SCHEMA on its catalog and schema, an LLM connection for the agent, and these environment +variables: + +- ``UNITY_MCP_SERVICE``: the service's three-level name, ``catalog.schema.service``. +- ``UNITY_MCP_READ_ONLY_TOOL``: the name of one of its tools that only reads, which the agent calls + once. +""" + +from __future__ import annotations + +import os +from datetime import datetime + +from airflow.providers.common.ai.operators.agent import AgentOperator +from airflow.providers.common.compat.sdk import DAG +from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + +DAG_ID = "example_databricks_unity_mcp" +UNITY_MCP_SERVICE = os.environ.get("UNITY_MCP_SERVICE", "main.default.my_mcp_service") +UNITY_MCP_READ_ONLY_TOOL = os.environ.get("UNITY_MCP_READ_ONLY_TOOL", "my_read_only_tool") + +with DAG( + dag_id=DAG_ID, + schedule=None, + start_date=datetime(2021, 1, 1), + tags=["example"], + catchup=False, +) as dag: + # [START howto_toolset_databricks_unity_mcp] + ask_unity_mcp_service = AgentOperator( + task_id="ask_unity_mcp_service", + prompt=( + "List the names of the tools you have. Then call the " + f"{UNITY_MCP_READ_ONLY_TOOL} tool once and summarize what it returned." + ), + llm_conn_id="pydanticai_default", + toolsets=[ + DatabricksUnityMCPToolset(UNITY_MCP_SERVICE, databricks_conn_id="databricks_default"), + ], + ) + # [END howto_toolset_databricks_unity_mcp] + + from tests_common.test_utils.watcher import watcher + + list(dag.tasks) >> watcher() + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +test_run = get_test_run(dag) diff --git a/providers/databricks/tests/unit/databricks/toolsets/__init__.py b/providers/databricks/tests/unit/databricks/toolsets/__init__.py new file mode 100644 index 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..8c80cd96fcc13 --- /dev/null +++ b/providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py @@ -0,0 +1,517 @@ +# 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 ModelRetry, 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, + DatabricksUnityMCPError, + 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" + +# What Unity Gateway answers, at HTTP 403, for a service the caller may not invoke or that does not exist. +NOT_AUTHORIZED = {"code": -32007, "message": "Not authorized to invoke MCP service."} +INVALID_PARAMS = {"code": -32602, "message": "Invalid params: 'text' is required"} + + +class _StubGateway(BaseHTTPRequestHandler): + """ + A Unity Gateway stand-in that serves one MCP Service with a single ``echo`` tool. + + A call of ``echo`` whose text is ``"reply "`` gets that HTTP status with an empty body, + and ``"reply "`` gets it with the JSON-RPC error named by ````. + """ + + protocol_version = "HTTP/1.1" + errors = {"not_authorized": NOT_AUTHORIZED, "invalid_params": INVALID_PARAMS} + # Per-test behaviour, set by the ``gateway`` fixture: a reply for every request to a service, + # a status for notifications only, and a barrier tool calls wait at, so that they overlap. + service_reply: dict[str, tuple[int, dict[str, str], str | None]] = {} + notification_status: int | None = None + tool_call_barrier: threading.Barrier | None = None + # Seconds to hold the answer to ``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): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"])) or b"{}") + method = body.get("method") + self.requests.append((self.path, self.headers.get("Authorization"), method)) + service = self.path.rsplit("/", 1)[-1] + if service in self.service_reply: + status, headers, error = self.service_reply[service] + return self._reply_error(status, body.get("id"), error, headers) + if "id" not in body: + return self._reply(self.notification_status or 202) + if method == "initialize": + self.release.wait(self.initialize_delay) + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "stub", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "echo", + "description": "Echo the text back.", + "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}}, + } + ] + } + elif method == "tools/call": + if self.tool_call_barrier: + self.tool_call_barrier.wait() + if "text" not in body["params"]["arguments"]: + return self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "error": INVALID_PARAMS}) + text = body["params"]["arguments"]["text"] + if text.startswith("reply "): + status, _, error = text.removeprefix("reply ").partition(" ") + return self._reply_error(int(status), body["id"], error or None) + result = {"content": [{"type": "text", "text": f"echo: {text}"}], "isError": False} + else: + result = {} + self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "result": result}) + + def do_GET(self): + self._reply(405) + + def do_DELETE(self): + self._reply(200) + + def _reply_error(self, status, request_id, error, headers=None): + payload = {"jsonrpc": "2.0", "id": request_id, "error": self.errors[error]} if error else None + self._reply(status, payload, headers) + + def _reply(self, status, payload=None, headers=None): + data = json.dumps(payload).encode() if payload is not None else b"" + self.send_response(status) + for name, value in (headers or {}).items(): + self.send_header(name, value) + if payload is not None: + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args): + pass + + +@pytest.fixture +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() + + +@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 + + +async def _call_concurrently(*texts): + """Call ``echo`` once per text, all at the same time, and return each call's result or exception.""" + toolset = DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID) + async with toolset: + tools = await toolset.get_tools(_ctx()) + return await asyncio.gather( + *(toolset.execute_tool("echo", {"text": t}, ctx=_ctx(), tool=tools["echo"]) for t in texts), + return_exceptions=True, + ) + + +class TestServiceName: + @pytest.mark.parametrize( + "name", + [ + "main.tools", + "main.tools.genie.extra", + "main.tools.genie/../../api/2.0/jobs", + "main.tools.genie?x=1", + "main.tools.genie#frag", + "main.tools.ge nie", + "main.tools.my-genie", + "main..genie", + "evil.example.com/main.tools.genie", + "", + ], + ) + def test_rejects_names_that_are_not_three_safe_parts(self, name): + with pytest.raises(ValueError, match="catalog.schema.service"): + DatabricksUnityMCPToolset(name) + + def test_templated_name_is_validated_once_rendered(self, gateway_conn): + toolset = DatabricksUnityMCPToolset("{{ params.service }}", databricks_conn_id=CONN_ID) + toolset._service_name = "main.tools/../../api" + + with pytest.raises(ValueError, match="catalog.schema.service"): + toolset._get_server() + + def test_id_names_connection_and_service(self): + assert ( + DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID).id + == f"databricks-unity-mcp-{CONN_ID}-{SERVICE}" + ) + + def test_connection_and_service_are_agent_template_fields(self): + assert set(DatabricksUnityMCPToolset.agent_template_fields) == { + "_databricks_conn_id", + "_service_name", + } + + +class TestServiceUrl: + @pytest.mark.parametrize( + ("host", "extra", "expected"), + [ + ("xx.cloud.databricks.com", None, "https://xx.cloud.databricks.com"), + ("https://xx.cloud.databricks.com/", None, "https://xx.cloud.databricks.com"), + (None, {"host": "https://yy.azuredatabricks.net"}, "https://yy.azuredatabricks.net"), + ], + ) + def test_url_is_built_on_the_connection_host(self, create_connection_without_db, host, extra, expected): + create_connection_without_db( + Connection(conn_id=CONN_ID, conn_type="databricks", host=host, password="t", extra=extra) + ) + toolset = DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID) + + assert ( + toolset.get_service_url(DatabricksHook(CONN_ID)) + == f"{expected}/ai-gateway/mcp-services/{SERVICE}" + ) + + def test_connection_without_host_is_rejected(self, create_connection_without_db): + create_connection_without_db(Connection(conn_id=CONN_ID, conn_type="databricks", password="t")) + toolset = DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID) + + with pytest.raises(ValueError, match="no workspace host"): + toolset.get_service_url(DatabricksHook(CONN_ID)) + + @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "[::1]"]) + def test_plain_http_is_allowed_to_loopback_only(self, create_connection_without_db, host): + create_connection_without_db( + Connection(conn_id=CONN_ID, conn_type="databricks", host=host, schema="http", password="t") + ) + toolset = DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID) + + assert toolset.get_service_url(DatabricksHook(CONN_ID)).startswith(f"http://{host}/") + + def test_plain_http_to_a_workspace_is_rejected(self, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=CONN_ID, + conn_type="databricks", + host="xx.cloud.databricks.com", + schema="http", + password="t", + ) + ) + toolset = DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID) + + with pytest.raises(ValueError, match="only sent over HTTPS"): + toolset.get_service_url(DatabricksHook(CONN_ID)) + + def test_requests_go_through_the_connection_proxy(self, gateway, create_connection_without_db): + # The stub gateway doubles as the proxy: a proxied request names the full URL in its path. + create_connection_without_db( + Connection( + conn_id=CONN_ID, + conn_type="databricks", + host="127.0.0.2", + port=1, + schema="http", + password="dapi-token", + extra={"proxies": {"http": f"http://127.0.0.1:{gateway.server_address[1]}"}}, + ) + ) + + _, result = _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + assert result == "echo: hi" + assert {path for path, _, _ in _StubGateway.requests} == { + f"http://127.0.0.2:1/ai-gateway/mcp-services/{SERVICE}" + } + + +class TestAuthentication: + def test_tools_are_discovered_and_called_with_the_connection_token(self, gateway, gateway_conn): + tools, result = _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + assert list(tools) == ["echo"] + assert result == "echo: hi" + assert _StubGateway.requests + assert {path for path, _, _ in _StubGateway.requests} == {f"/ai-gateway/mcp-services/{SERVICE}"} + assert {auth for _, auth, _ in _StubGateway.requests} == {"Bearer dapi-token"} + + @mock.patch.object(DatabricksHook, "_get_token", autospec=True) + def test_token_is_fetched_for_every_request_so_it_can_refresh(self, get_token, gateway, gateway_conn): + tokens = (f"oauth-{i}" for i in range(1000)) + get_token.side_effect = lambda *_, **__: next(tokens) + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + sent = [auth for _, auth, _ in _StubGateway.requests] + # One token is fetched up front to fail fast, then one per request. + assert sent == [f"Bearer oauth-{i}" for i in range(1, len(sent) + 1)] + + 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))) + + 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="sends a bearer token"): + DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID)._get_server() + + @mock.patch.object(DatabricksHook, "_get_token", autospec=True) + def test_token_failure_during_the_run_is_reported(self, get_token, gateway, gateway_conn): + # The up-front check and the first request get a token; the next request's fetch fails. + get_token.side_effect = ["dapi-token", "dapi-token", OSError("token endpoint unreachable")] + + with pytest.raises(DatabricksUnityMCPError, match="Could not get a token") as err: + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + assert "token endpoint unreachable" in str(err.value) + + +class TestErrors: + @pytest.mark.parametrize( + ("status", "headers", "error", "expected", "retry_after"), + [ + (401, {}, None, DatabricksUnityMCPAccessDeniedError, None), + (403, {}, None, DatabricksUnityMCPAccessDeniedError, None), + (403, {}, "not_authorized", DatabricksUnityMCPAccessDeniedError, None), + (404, {}, None, DatabricksUnityMCPError, None), + (429, {"Retry-After": "12"}, None, DatabricksUnityMCPThrottledError, 12.0), + ( + 429, + {"Retry-After": "Wed, 21 Oct 2026 07:28:00 GMT"}, + None, + DatabricksUnityMCPThrottledError, + None, + ), + ], + ) + def test_gateway_errors_while_connecting_are_translated( + self, gateway, gateway_conn, status, headers, error, expected, retry_after + ): + _StubGateway.service_reply[SERVICE] = (status, headers, error) + + with pytest.raises(expected) as err: + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + assert type(err.value) is expected + assert err.value.http_status_code == status + if expected is DatabricksUnityMCPThrottledError: + assert err.value.retry_after == retry_after + + def test_missing_service_is_reported_as_access_denied(self, gateway, gateway_conn): + _StubGateway.service_reply[SERVICE] = (403, {}, "not_authorized") + + with pytest.raises(DatabricksUnityMCPAccessDeniedError, match="no service has that name"): + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + def test_server_error_during_a_tool_call_says_the_call_may_have_run(self, gateway, gateway_conn): + with pytest.raises(DatabricksUnityMCPError, match="may or may not have run") as err: + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID), "reply 503")) + + assert type(err.value) is DatabricksUnityMCPError + assert err.value.http_status_code == 503 + # The failed call was sent once and not retried. + assert [m for _, _, m in _StubGateway.requests].count("tools/call") == 1 + + def test_denied_tool_call_is_not_left_to_the_model(self, gateway, gateway_conn): + with pytest.raises(DatabricksUnityMCPAccessDeniedError): + _run( + _list_and_call( + DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID), "reply 403 not_authorized" + ) + ) + + @pytest.mark.parametrize( + "text", ["reply 200 invalid_params", "reply 400 invalid_params", "reply 500 invalid_params"] + ) + def test_json_rpc_error_from_the_server_is_left_to_the_model(self, gateway, gateway_conn, text): + # A server may answer bad arguments with a JSON-RPC error at an HTTP error status; it is + # still the server's answer, so the model can correct the call. + with pytest.raises(ModelRetry, match="'text' is required"): + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID), text)) + + def test_concurrent_calls_are_each_judged_by_their_own_failure(self, gateway, gateway_conn): + # Calls run in parallel, so another call's response must not change how a failure is + # reported: a server error stays a non-retried failure however the other call ends. + _StubGateway.tool_call_barrier = threading.Barrier(2, timeout=10) + for _ in range(10): + failed, succeeded = _run(_call_concurrently("reply 503", "hi")) + assert type(failed) is DatabricksUnityMCPError + assert "may or may not have run" in str(failed) + assert succeeded == "echo: hi" + + failed, throttled = _run(_call_concurrently("reply 503", "reply 429")) + # Neither status can be attributed to its call, so neither is reported as throttling. + assert type(failed) is DatabricksUnityMCPError + assert type(throttled) is DatabricksUnityMCPError + assert failed.http_status_code is None + + invalid, failed = _run(_call_concurrently("reply 400 invalid_params", "reply 503")) + assert isinstance(invalid, ModelRetry) + assert type(failed) is DatabricksUnityMCPError + + 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( + conn_id=CONN_ID, conn_type="databricks", host="127.0.0.1", port=1, schema="http", password="t" + ) + ) + + with pytest.raises(DatabricksUnityMCPTransportError, match="Could not reach Unity Gateway"): + _run(_list_and_call(DatabricksUnityMCPToolset(SERVICE, databricks_conn_id=CONN_ID))) + + +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"): + DatabricksUnityMCPToolset(SERVICE).airflow_tools() + + +class TestOptionalDependency: + def test_import_without_common_ai_names_the_extra(self): + module = "airflow.providers.databricks.toolsets.unity_mcp" + with mock.patch.dict(sys.modules, {"airflow.providers.common.ai.utils.toolset_base": None}): + sys.modules.pop(module) + with pytest.raises(AirflowOptionalProviderFeatureException, match=r"databricks\[common.ai\]"): + importlib.import_module(module) diff --git a/uv.lock b/uv.lock index 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" },