kaxil commented on code in PR #74198: URL: https://github.com/apache/airflow/pull/74198#discussion_r4317373054
########## providers/databricks/docs/toolsets/unity_mcp.rst: ########## @@ -0,0 +1,120 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/toolset:DatabricksUnityMCPToolset: + +Unity Gateway MCP Services +========================== + +Use :class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset` to give an +agent the tools of a Unity Gateway MCP Service: either one Databricks provides for workspace tools +and SaaS applications, such as ``system.ai.google_calendar``, or an external MCP server registered as +an MCP Service in Unity Catalog. The toolset works with +:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, ``@task.agent``, and the LangChain +bridge of the Common AI provider. + +The Dag names the service by its three-level Unity Catalog name, ``catalog.schema.service``, and the +:ref:`Databricks connection <howto/connection:databricks>` to use. The toolset builds the service URL, +``https://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``, from the connection's +host, so neither the gateway URL nor a token appears in Dag code, and the connection's token is only +sent to that workspace, over HTTPS. A connection whose ``schema`` is ``http`` is rejected unless its +host is a loopback address. Requests go through the proxy in the connection's ``proxies`` extra, if +set. Service names may contain only ASCII letters, digits and ``_`` in each part; any other name is +rejected before a request is made. + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_unity_mcp.py + :language: python + :start-after: [START howto_toolset_databricks_unity_mcp] + :end-before: [END howto_toolset_databricks_unity_mcp] + +The service name and connection ID are templated when the toolset is passed to ``AgentOperator`` or +``@task.agent``. + +Caller identity and authentication +---------------------------------- + +MCP Services need a workspace enabled for Unity Catalog, in a region where Model Serving is +supported. + +The gateway runs every tool call as the identity of the connection's credentials. That identity needs +``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and schema +(``EXECUTE`` alone is not enough), and an assignment to the workspace. It needs no privilege on the +Unity Catalog connection behind the service. On the built-in ``system.ai`` services, account users +hold these privileges by default (see `MCP Services +<https://docs.databricks.com/aws/en/agents/mcp-tools/mcp-services>`__). The gateway exposes only the +tools selected for the service, and the service's policies apply. Grant the identity only the services +the agent should use. + +Built-in services that act on a user's own data, such as ``system.ai.google_calendar`` or +``system.ai.gmail``, need that identity to complete a one-time OAuth login first, for example by +opening the service in Catalog Explorer and clicking **Login**. + +The toolset sends the connection's token as a bearer token, so the connection must use one of these +authentication modes of the Databricks connection: + +* a personal access token (the identity is the token's user or service principal); +* service principal OAuth (``service_principal_oauth``); +* Azure AD: a service principal, a managed identity, or ``DefaultAzureCredential``; +* workload identity federation (Kubernetes, AWS IAM, or a supplied token provider). + +For service principal OAuth, the OAuth secret must allow the ``all-apis`` scope, which the connection +requests; a secret restricted to narrower scopes fails (see `OAuth for service principals +<https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m>`__). Username and password +authentication is not supported. A token is fetched from the connection for +every request, so OAuth and Azure AD tokens are refreshed during a long agent run and when the toolset +reconnects. Tokens minted for the toolset are masked in task logs. + +Errors and retries +------------------ + +When the MCP server answers a tool call with an error, such as invalid arguments, the error goes to the +model, which can correct the call and try again. + +When the gateway gives no answer of its own, a tool call that was already sent may or may not have +run, and repeating a tool that changes data could apply the change twice. The toolset never retries +it: the task fails with one of these exceptions from :mod:`airflow.providers.databricks.exceptions`, +which are also raised for failures while connecting and listing tools: + +.. list-table:: + :header-rows: 1 + + * - Exception + - Cause + * - ``DatabricksUnityMCPAccessDeniedError`` + - The gateway will not let the identity invoke the service: no service has that name, the + identity lacks ``EXECUTE`` on the service or ``USE CATALOG`` / ``USE SCHEMA`` on its parents, + or the credentials are invalid. A missing service is reported like a missing privilege. + * - ``DatabricksUnityMCPThrottledError`` + - HTTP 429. Its ``retry_after`` attribute holds the ``Retry-After`` delay in seconds, when given. + * - ``DatabricksUnityMCPTransportError`` + - The gateway could not be reached, the connection dropped, a request got no answer in time, + or the MCP handshake did not finish within ``init_timeout``. + * - ``DatabricksUnityMCPError`` + - Any other gateway error, or a failure to get a token from the connection after the toolset + connected. + +A connection that cannot be used at all, such as one without a host or without token-based +authentication, raises ``ValueError`` when the toolset connects. + +Unity Gateway forwards a response only once the MCP server has finished sending it, so connecting +to a server that keeps its response stream open takes as long as the server keeps it open. The toolset +waits ``init_timeout`` seconds (30 by default) for the handshake; raise it for slow services. + +The exceptions carry the HTTP status in ``http_status_code`` when it is known. When the agent runs +several tool calls of the toolset at once, the status of a failed call cannot be told apart from the +others', so the error is reported as ``DatabricksUnityMCPError`` without it. Use task retries for Review Comment: This reads broader than what happens. With overlapping calls only the status-based classes are lost (401/403 and 429); transport errors, timeouts and a -32007 denial keep their class. ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,521 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +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 + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What the MCP client raises when a request outlasts its read timeout (see +# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the message is matched too. +_REQUEST_TIMEOUT_CODE = -32001 +_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out") +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream 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_request_timeout(error: McpError | None) -> bool: + return ( + error is not None + and error.error.code == _REQUEST_TIMEOUT_CODE + and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None + ) + + +def _is_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) Review Comment: `get_token()` calls `mask_secret` on every request, and before Airflow 3.2 that is a `SUPERVISOR_COMMS.send` with no lock on the socket (`_thread_lock` only arrived in 3.2.0). The supervisor answers each `MaskSecret` with a frame, so from `to_thread` outside `run_blocking`, two parallel tool calls, or this plus another toolset's connection lookup, can read each other's responses on 3.0/3.1, which both providers still support. So the comment above isn't quite right: fetching a token does talk to the supervisor. Could it mask only when the token changes, and do that one call through `AirflowToolset.run_blocking`? ########## providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py: ########## @@ -0,0 +1,576 @@ +# 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 functools +import importlib +import inspect +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 mcp_types import jsonrpc +from pydantic_ai import ModelRetry, RunContext +from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + +from airflow.models import Connection +from airflow.providers.common.ai.utils import toolset_base +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook +from airflow.providers.databricks.toolsets import unity_mcp +from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + +CONN_ID = "databricks_unity" +SERVICE = "main.tools.genie" + +# What Unity Gateway answers, at HTTP 403, for a service the caller may not invoke or that does not exist. +NOT_AUTHORIZED = {"code": -32007, "message": "Not authorized to invoke MCP service."} +INVALID_PARAMS = {"code": -32602, "message": "Invalid params: 'text' is required"} + + +class _StubGateway(BaseHTTPRequestHandler): + """ + A Unity Gateway stand-in that serves one MCP Service with a single ``echo`` tool. + + A call of ``echo`` whose text is ``"reply <status>"`` gets that HTTP status with an empty body, + and ``"reply <status> <error>"`` gets it with the JSON-RPC error named by ``<error>``. + """ + + protocol_version = "HTTP/1.1" + errors = {"not_authorized": NOT_AUTHORIZED, "invalid_params": INVALID_PARAMS} + # Per-test behaviour, set by the ``gateway`` fixture: a reply for every request to a service, + # a status for notifications only, and a barrier tool calls wait at, so that they overlap. + service_reply: dict[str, tuple[int, dict[str, str], str | None]] = {} + notification_status: int | None = None + tool_call_barrier: threading.Barrier | None = None + # Seconds to hold the answer to each MCP method; ``release`` ends any wait early on teardown. + delays: dict[str, float] = {} + release = threading.Event() + requests: list[tuple[str, str | None, str | None]] = [] + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"])) or b"{}") + method = body.get("method") + self.requests.append((self.path, self.headers.get("Authorization"), method)) + service = self.path.rsplit("/", 1)[-1] + if service in self.service_reply: + status, headers, error = self.service_reply[service] + return self._reply_error(status, body.get("id"), error, headers) + if "id" not in body: + return self._reply(self.notification_status or 202) + self.release.wait(self.delays.get(method, 0)) + if method == "initialize": + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "stub", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "echo", + "description": "Echo the text back.", + "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}}, + } + ] + } + elif method == "tools/call": + if self.tool_call_barrier: + self.tool_call_barrier.wait() + if "text" not in body["params"]["arguments"]: + return self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "error": INVALID_PARAMS}) + text = body["params"]["arguments"]["text"] + if text.startswith("reply "): + status, _, error = text.removeprefix("reply ").partition(" ") + return self._reply_error(int(status), body["id"], error or None) + result = {"content": [{"type": "text", "text": f"echo: {text}"}], "isError": False} + else: + result = {} + self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "result": result}) + + def do_GET(self): + self._reply(405) + + def do_DELETE(self): + self._reply(200) + + def _reply_error(self, status, request_id, error, headers=None): + payload = {"jsonrpc": "2.0", "id": request_id, "error": self.errors[error]} if error else None + self._reply(status, payload, headers) + + def _reply(self, status, payload=None, headers=None): + data = json.dumps(payload).encode() if payload is not None else b"" + self.send_response(status) + for name, value in (headers or {}).items(): + self.send_header(name, value) + if payload is not None: + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args): + pass + + [email protected] +def gateway(): + _StubGateway.service_reply = {} + _StubGateway.notification_status = None + _StubGateway.tool_call_barrier = None + _StubGateway.delays = {} + _StubGateway.release = threading.Event() + _StubGateway.requests = [] + server = ThreadingHTTPServer(("127.0.0.1", 0), _StubGateway) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server + _StubGateway.release.set() + server.shutdown() + server.server_close() + + [email protected] +def gateway_conn(gateway, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=CONN_ID, + conn_type="databricks", + host="127.0.0.1", + port=gateway.server_address[1], + schema="http", + password="dapi-token", + ) + ) + + +def _run(coro): + return asyncio.run(coro) + + +def _ctx(): + ctx = mock.MagicMock(spec=RunContext) + ctx.max_retries = 1 + return ctx + + +async def _list_and_call(toolset, text="hi"): + async with toolset: + tools = await toolset.get_tools(_ctx()) + result = await toolset.execute_tool("echo", {"text": text}, ctx=_ctx(), tool=tools["echo"]) + return tools, result + + +async def _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): Review Comment: With per-operation error lists, the 400 on `notifications/initialized` lands on the connect operation, which succeeds, so this passes whatever `_is_jsonrpc_request` returns (making it always `True` keeps the suite green). A direct test with a POST request, a POST notification and a GET would bind it. ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,521 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +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: Review Comment: `airflow.providers.common.compat.sdk` already exports `mask_secret` with this same fallback chain, and line 32 imports from it already. ########## providers/databricks/docs/toolsets/unity_mcp.rst: ########## @@ -0,0 +1,120 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/toolset:DatabricksUnityMCPToolset: + +Unity Gateway MCP Services +========================== + +Use :class:`~airflow.providers.databricks.toolsets.unity_mcp.DatabricksUnityMCPToolset` to give an +agent the tools of a Unity Gateway MCP Service: either one Databricks provides for workspace tools +and SaaS applications, such as ``system.ai.google_calendar``, or an external MCP server registered as +an MCP Service in Unity Catalog. The toolset works with +:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, ``@task.agent``, and the LangChain +bridge of the Common AI provider. + +The Dag names the service by its three-level Unity Catalog name, ``catalog.schema.service``, and the +:ref:`Databricks connection <howto/connection:databricks>` to use. The toolset builds the service URL, +``https://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``, from the connection's +host, so neither the gateway URL nor a token appears in Dag code, and the connection's token is only +sent to that workspace, over HTTPS. A connection whose ``schema`` is ``http`` is rejected unless its +host is a loopback address. Requests go through the proxy in the connection's ``proxies`` extra, if +set. Service names may contain only ASCII letters, digits and ``_`` in each part; any other name is +rejected before a request is made. + +.. exampleinclude:: /../../databricks/tests/system/databricks/example_databricks_unity_mcp.py + :language: python + :start-after: [START howto_toolset_databricks_unity_mcp] + :end-before: [END howto_toolset_databricks_unity_mcp] + +The service name and connection ID are templated when the toolset is passed to ``AgentOperator`` or +``@task.agent``. + +Caller identity and authentication +---------------------------------- + +MCP Services need a workspace enabled for Unity Catalog, in a region where Model Serving is +supported. + +The gateway runs every tool call as the identity of the connection's credentials. That identity needs +``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its parent catalog and schema +(``EXECUTE`` alone is not enough), and an assignment to the workspace. It needs no privilege on the +Unity Catalog connection behind the service. On the built-in ``system.ai`` services, account users +hold these privileges by default (see `MCP Services +<https://docs.databricks.com/aws/en/agents/mcp-tools/mcp-services>`__). The gateway exposes only the +tools selected for the service, and the service's policies apply. Grant the identity only the services +the agent should use. + +Built-in services that act on a user's own data, such as ``system.ai.google_calendar`` or +``system.ai.gmail``, need that identity to complete a one-time OAuth login first, for example by +opening the service in Catalog Explorer and clicking **Login**. + +The toolset sends the connection's token as a bearer token, so the connection must use one of these +authentication modes of the Databricks connection: + +* a personal access token (the identity is the token's user or service principal); +* service principal OAuth (``service_principal_oauth``); +* Azure AD: a service principal, a managed identity, or ``DefaultAzureCredential``; +* workload identity federation (Kubernetes, AWS IAM, or a supplied token provider). + +For service principal OAuth, the OAuth secret must allow the ``all-apis`` scope, which the connection +requests; a secret restricted to narrower scopes fails (see `OAuth for service principals +<https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m>`__). Username and password +authentication is not supported. A token is fetched from the connection for +every request, so OAuth and Azure AD tokens are refreshed during a long agent run and when the toolset +reconnects. Tokens minted for the toolset are masked in task logs. + +Errors and retries +------------------ + +When the MCP server answers a tool call with an error, such as invalid arguments, the error goes to the +model, which can correct the call and try again. + +When the gateway gives no answer of its own, a tool call that was already sent may or may not have +run, and repeating a tool that changes data could apply the change twice. The toolset never retries +it: the task fails with one of these exceptions from :mod:`airflow.providers.databricks.exceptions`, +which are also raised for failures while connecting and listing tools: + +.. list-table:: + :header-rows: 1 + + * - Exception + - Cause + * - ``DatabricksUnityMCPAccessDeniedError`` + - The gateway will not let the identity invoke the service: no service has that name, the + identity lacks ``EXECUTE`` on the service or ``USE CATALOG`` / ``USE SCHEMA`` on its parents, + or the credentials are invalid. A missing service is reported like a missing privilege. + * - ``DatabricksUnityMCPThrottledError`` + - HTTP 429. Its ``retry_after`` attribute holds the ``Retry-After`` delay in seconds, when given. Review Comment: `_parse_retry_after` returns `None` for the HTTP-date form (the test asserts that), so "when given as a number of seconds"? ########## providers/databricks/tests/unit/databricks/toolsets/test_unity_mcp.py: ########## @@ -0,0 +1,576 @@ +# 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 functools +import importlib +import inspect +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 mcp_types import jsonrpc +from pydantic_ai import ModelRetry, RunContext +from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + +from airflow.models import Connection +from airflow.providers.common.ai.utils import toolset_base +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook +from airflow.providers.databricks.toolsets import unity_mcp +from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + +CONN_ID = "databricks_unity" +SERVICE = "main.tools.genie" + +# What Unity Gateway answers, at HTTP 403, for a service the caller may not invoke or that does not exist. +NOT_AUTHORIZED = {"code": -32007, "message": "Not authorized to invoke MCP service."} +INVALID_PARAMS = {"code": -32602, "message": "Invalid params: 'text' is required"} + + +class _StubGateway(BaseHTTPRequestHandler): + """ + A Unity Gateway stand-in that serves one MCP Service with a single ``echo`` tool. + + A call of ``echo`` whose text is ``"reply <status>"`` gets that HTTP status with an empty body, + and ``"reply <status> <error>"`` gets it with the JSON-RPC error named by ``<error>``. + """ + + protocol_version = "HTTP/1.1" + errors = {"not_authorized": NOT_AUTHORIZED, "invalid_params": INVALID_PARAMS} + # Per-test behaviour, set by the ``gateway`` fixture: a reply for every request to a service, + # a status for notifications only, and a barrier tool calls wait at, so that they overlap. + service_reply: dict[str, tuple[int, dict[str, str], str | None]] = {} + notification_status: int | None = None + tool_call_barrier: threading.Barrier | None = None + # Seconds to hold the answer to each MCP method; ``release`` ends any wait early on teardown. + delays: dict[str, float] = {} + release = threading.Event() + requests: list[tuple[str, str | None, str | None]] = [] + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"])) or b"{}") + method = body.get("method") + self.requests.append((self.path, self.headers.get("Authorization"), method)) + service = self.path.rsplit("/", 1)[-1] + if service in self.service_reply: + status, headers, error = self.service_reply[service] + return self._reply_error(status, body.get("id"), error, headers) + if "id" not in body: + return self._reply(self.notification_status or 202) + self.release.wait(self.delays.get(method, 0)) + if method == "initialize": + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "stub", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "echo", + "description": "Echo the text back.", + "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}}, + } + ] + } + elif method == "tools/call": + if self.tool_call_barrier: + self.tool_call_barrier.wait() + if "text" not in body["params"]["arguments"]: + return self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "error": INVALID_PARAMS}) + text = body["params"]["arguments"]["text"] + if text.startswith("reply "): + status, _, error = text.removeprefix("reply ").partition(" ") + return self._reply_error(int(status), body["id"], error or None) + result = {"content": [{"type": "text", "text": f"echo: {text}"}], "isError": False} + else: + result = {} + self._reply(200, {"jsonrpc": "2.0", "id": body["id"], "result": result}) + + def do_GET(self): + self._reply(405) + + def do_DELETE(self): + self._reply(200) + + def _reply_error(self, status, request_id, error, headers=None): + payload = {"jsonrpc": "2.0", "id": request_id, "error": self.errors[error]} if error else None + self._reply(status, payload, headers) + + def _reply(self, status, payload=None, headers=None): + data = json.dumps(payload).encode() if payload is not None else b"" + self.send_response(status) + for name, value in (headers or {}).items(): + self.send_header(name, value) + if payload is not None: + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args): + pass + + [email protected] +def gateway(): + _StubGateway.service_reply = {} + _StubGateway.notification_status = None + _StubGateway.tool_call_barrier = None + _StubGateway.delays = {} + _StubGateway.release = threading.Event() + _StubGateway.requests = [] + server = ThreadingHTTPServer(("127.0.0.1", 0), _StubGateway) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server + _StubGateway.release.set() + server.shutdown() + server.server_close() + + [email protected] +def gateway_conn(gateway, create_connection_without_db): + create_connection_without_db( + Connection( + conn_id=CONN_ID, + conn_type="databricks", + host="127.0.0.1", + port=gateway.server_address[1], + schema="http", + password="dapi-token", + ) + ) + + +def _run(coro): + return asyncio.run(coro) + + +def _ctx(): + ctx = mock.MagicMock(spec=RunContext) + ctx.max_retries = 1 + return ctx + + +async def _list_and_call(toolset, text="hi"): + async with toolset: + tools = await toolset.get_tools(_ctx()) + result = await toolset.execute_tool("echo", {"text": text}, ctx=_ctx(), tool=tools["echo"]) + return tools, result + + +async def _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")) Review Comment: Nothing pins the -32007 check in `_translate_error`. Every denied-access test runs on a call that doesn't overlap, where `status in (401, 403)` already decides it, so dropping the code check keeps the suite green. That check is the only thing keeping a denied call that overlaps another from going back to the model as a retry. Adding a `"reply 403 not_authorized"` + `"hi"` pair here and asserting `DatabricksUnityMCPAccessDeniedError` would cover it. ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,521 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +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 + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What the MCP client raises when a request outlasts its read timeout (see +# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the message is matched too. +_REQUEST_TIMEOUT_CODE = -32001 +_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out") +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream ended without a response", Review Comment: mcp 2.2.0 also makes up `(-32000, "SSE stream ended and reconnection attempts were exhausted")` (`mcp/client/streamable_http.py:533`) when a tool call's SSE stream drops after an event with an id and the Last-Event-ID reconnects fail. That doesn't match this prefix or `_CLIENT_ERRORS`, so `_translate_error` returns `None` and the model gets a `ModelRetry` for a call that may already have run. I don't know whether Unity Gateway sends event ids today, but the client path doesn't depend on it, and a `"SSE stream ended"` prefix covers both messages. `TestClientErrorMessages` checks that each listed string still exists, not the other way round. Collecting the messages passed to `_resolve_abandoned_request(` and `ErrorData(` in the client source and asserting `_is_client_error` matches each one would catch the next one in CI. ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,521 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +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 + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What the MCP client raises when a request outlasts its read timeout (see +# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the message is matched too. +_REQUEST_TIMEOUT_CODE = -32001 +_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out") +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream 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_request_timeout(error: McpError | None) -> bool: + return ( + error is not None + and error.error.code == _REQUEST_TIMEOUT_CODE + and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None + ) + + +def _is_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://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``. The + workspace host and the credentials both come from the Databricks connection, so Dag code holds + neither a gateway URL nor a token, and the token is only ever sent to that workspace, over + HTTPS. The connection's ``proxies`` extra applies to gateway requests. + + The gateway runs every tool call as the identity of the connection's credentials (a user's + personal access token, a service principal, or an Azure AD / federated identity). That identity + needs ``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and + schema, and an assignment to the workspace. It sees only the tools selected for the service, and + the service's policies apply. + + Tokens are fetched from the connection for each request, so OAuth tokens refresh during a long + agent run and on reconnection. An error the server returns for a tool call reaches the model, + which can correct its arguments and try again. When the gateway gives no answer of its own, a + tool call may or may not have run, so it is never retried; it fails the task with + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`, as do failures + while connecting and listing tools. + + .. code-block:: python + + from airflow.providers.common.ai.operators.agent import AgentOperator + from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + + AgentOperator( + task_id="ask_mcp_service", + prompt="Which tools do you have?", + llm_conn_id="pydanticai_default", + toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp", databricks_conn_id="databricks")], + ) + + :param service_name: Three-level name of the MCP Service, ``catalog.schema.service``. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param databricks_conn_id: Databricks connection whose host and credentials are used. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param tool_prefix: Optional prefix prepended to tool names. + :param init_timeout: Seconds to wait for the MCP handshake when connecting. Unity Gateway forwards + a response only once the MCP server has finished sending it, so a server that keeps its + response stream open needs more than pydantic-ai's default of 5 seconds. + """ + + # Rendered, on a copy, by AgentOperator. + agent_template_fields: Sequence[str] = ("_databricks_conn_id", "_service_name") + + def __init__( + self, + service_name: str, + *, + databricks_conn_id: str = DatabricksHook.default_conn_name, + tool_prefix: str | None = None, + init_timeout: float = 30, + ) -> None: + if not init_timeout > 0: + raise ValueError(f"init_timeout must be a positive number of seconds, got {init_timeout!r}.") + self._databricks_conn_id = databricks_conn_id + self._service_name = service_name + self._tool_prefix = tool_prefix + self._init_timeout = init_timeout + self._server: Any = None + self._operations: set[_Operation] = set() + # 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() Review Comment: A failure here isn't wrapped. With service principal OAuth and a token endpoint that hangs, the `requests` error carries a `TimeoutError` in its context, so `_is_handshake_timeout` reports "did not complete the MCP handshake ... Raise init_timeout", though the gateway was never contacted and the wait is the hook's 3 x 10s token timeout. A 401 from the token endpoint comes out as the hook's raw `AirflowException`. Wrapping it like `async_auth_flow` does (re-raise `ValueError`, wrap the rest in `_TokenUnavailableError`) would give the "Could not get a token" error, and the docs row could drop "after the toolset connected". ########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,521 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +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 + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What the MCP client raises when a request outlasts its read timeout (see +# mcp.shared.jsonrpc_dispatcher). A server may use the same code, so the message is matched too. +_REQUEST_TIMEOUT_CODE = -32001 +_REQUEST_TIMEOUT_MESSAGE = re.compile(r"Request '[^']*' timed out") +_CLIENT_ERROR_PREFIXES = ( + "Unexpected content type:", + "Failed to parse JSON response:", + "Failed to parse SSE message:", + "SSE stream 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_request_timeout(error: McpError | None) -> bool: + return ( + error is not None + and error.error.code == _REQUEST_TIMEOUT_CODE + and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None + ) + + +def _is_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://<workspace host>/ai-gateway/mcp-services/<catalog.schema.service>``. The + workspace host and the credentials both come from the Databricks connection, so Dag code holds + neither a gateway URL nor a token, and the token is only ever sent to that workspace, over + HTTPS. The connection's ``proxies`` extra applies to gateway requests. + + The gateway runs every tool call as the identity of the connection's credentials (a user's + personal access token, a service principal, or an Azure AD / federated identity). That identity + needs ``EXECUTE`` on the MCP Service, ``USE CATALOG`` and ``USE SCHEMA`` on its catalog and + schema, and an assignment to the workspace. It sees only the tools selected for the service, and + the service's policies apply. + + Tokens are fetched from the connection for each request, so OAuth tokens refresh during a long + agent run and on reconnection. An error the server returns for a tool call reaches the model, + which can correct its arguments and try again. When the gateway gives no answer of its own, a + tool call may or may not have run, so it is never retried; it fails the task with + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPAccessDeniedError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPThrottledError`, + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPTransportError`, or + :class:`~airflow.providers.databricks.exceptions.DatabricksUnityMCPError`, as do failures + while connecting and listing tools. + + .. code-block:: python + + from airflow.providers.common.ai.operators.agent import AgentOperator + from airflow.providers.databricks.toolsets.unity_mcp import DatabricksUnityMCPToolset + + AgentOperator( + task_id="ask_mcp_service", + prompt="Which tools do you have?", + llm_conn_id="pydanticai_default", + toolsets=[DatabricksUnityMCPToolset("main.default.my_mcp", databricks_conn_id="databricks")], + ) + + :param service_name: Three-level name of the MCP Service, ``catalog.schema.service``. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param databricks_conn_id: Databricks connection whose host and credentials are used. Templated + when the toolset is passed to ``AgentOperator`` / ``@task.agent``. + :param tool_prefix: Optional prefix prepended to tool names. + :param init_timeout: Seconds to wait for the MCP handshake when connecting. Unity Gateway forwards + a response only once the MCP server has finished sending it, so a server that keeps its + response stream open needs more than pydantic-ai's default of 5 seconds. + """ + + # Rendered, on a copy, by AgentOperator. + agent_template_fields: Sequence[str] = ("_databricks_conn_id", "_service_name") + + def __init__( + self, + service_name: str, + *, + databricks_conn_id: str = DatabricksHook.default_conn_name, + tool_prefix: str | None = None, + init_timeout: float = 30, + ) -> None: + if not init_timeout > 0: + raise ValueError(f"init_timeout must be a positive number of seconds, got {init_timeout!r}.") + self._databricks_conn_id = databricks_conn_id + self._service_name = service_name + self._tool_prefix = tool_prefix + self._init_timeout = init_timeout + self._server: Any = None + self._operations: set[_Operation] = set() + # 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 (transport_error := _find(error, httpx2.TransportError)) is not None: Review Comment: Once the session is open, the POST runs in the transport's own task group, so an `httpx2` error there never reaches the waiting call's exception chain. A socket dropped mid tool call comes out as `DatabricksUnityMCPError("...an error without an answer (Connection closed)")`, not the `DatabricksUnityMCPTransportError` the docs table and the exception docstring promise. A token refresh failing mid-run gives the same message rather than "Could not get a token". Not retrying is still right; it's the class and message that are off. Recording these on the running operations from `async_auth_flow` (like `errors`) or narrowing the docs would fix it. `test_token_failure_during_the_run_is_reported` fails the fetch while connecting, so it doesn't cover this path. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
