firasbouzazi commented on code in PR #74198: URL: https://github.com/apache/airflow/pull/74198#discussion_r4323962019
########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,601 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP Service.""" + +from __future__ import annotations + +import asyncio +import contextlib +import ipaddress +import json +import logging +import re +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, mask_secret +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook + +try: + # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the common.ai extra installs. + import httpx2 + from fastmcp.client.transports import StreamableHttpTransport + from fastmcp.exceptions import McpError + from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + + from airflow.providers.common.ai.utils.toolset_base import AirflowToolset +except ImportError as e: + raise AirflowOptionalProviderFeatureException( + "DatabricksUnityMCPToolset needs the 'common.ai' extra of the databricks provider: " + "pip install 'apache-airflow-providers-databricks[common.ai]'" + ) from e + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence + + from pydantic_ai._run_context import RunContext + from pydantic_ai.toolsets.abstract import ToolsetTool + + from airflow.providers.common.ai.tools import AirflowTool + +log = logging.getLogger(__name__) + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + +# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the caller may not invoke +# the service, including when no service has that name. +_NOT_AUTHORIZED_CODE = -32007 + +# Errors the MCP client makes up when the gateway gave no JSON-RPC answer, such as an HTTP error +# with a plain body or a response that never arrived (see mcp.client.streamable_http). Every other +# MCP error is the server's own answer. Matched on code and message, the only things they carry. +_CLIENT_ERRORS = { + (-32603, "Server returned an error response"), + (-32601, "Not Found"), + (-32600, "Session terminated"), + (-32000, "Connection closed"), +} +# What fastmcp raises, in place of the TimeoutError it loses, when the MCP handshake outlasts +# init_timeout (see fastmcp.client.client). +_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session" +# What 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", + "server answered a request with 202 Accepted", + "Redirect to ", +) + + +def validate_service_name(service_name: str) -> None: + """ + Raise ``ValueError`` unless ``service_name`` is a three-level ``catalog.schema.service`` name. + + Each part may contain only ASCII letters, digits and underscores. + """ + if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): + raise ValueError( + f"Invalid Unity Gateway MCP Service name {service_name!r}: expected " + "'catalog.schema.service', each part made of ASCII letters, digits or '_'." + ) + + +def _is_loopback(host: str) -> bool: + if host == "localhost": + return True + try: + return ipaddress.ip_address(host.strip("[]")).is_loopback + except ValueError: + return False + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + try: + seconds = float(value) + except ValueError: + # The HTTP-date form is not worth parsing: the caller only uses this as a hint. + return None + return seconds if seconds >= 0 else None + + +def _walk_exceptions(exc: BaseException) -> Iterator[BaseException]: + """Yield ``exc``, its causes and the exceptions inside any exception groups, each once.""" + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop(0) + if id(current) in seen: + continue + seen.add(id(current)) + yield current + # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio uses the backport on 3.10. + if isinstance(grouped := getattr(current, "exceptions", None), (list, tuple)): + pending.extend(e for e in grouped if isinstance(e, BaseException)) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + + +def _describe(error: BaseException) -> str: + return str(error) or type(error).__name__ + + +def _find(exc: BaseException, types: type[BaseException] | tuple[type[BaseException], ...]) -> Any: + """Return the first exception of ``types`` in ``exc``, its causes and any exception groups.""" + return next((e for e in _walk_exceptions(exc) if isinstance(e, types)), None) + + +def _is_handshake_timeout(error: BaseException) -> bool: + return any( + isinstance(e, TimeoutError) + or (isinstance(e, RuntimeError) and str(e).endswith(_HANDSHAKE_TIMEOUT_MESSAGE)) + for e in _walk_exceptions(error) + ) + + +def _is_request_timeout(error: McpError | None) -> bool: + return ( + error is not None + and error.error.code == _REQUEST_TIMEOUT_CODE + and _REQUEST_TIMEOUT_MESSAGE.fullmatch(error.error.message) is not None + ) + + +def _is_client_error(error: McpError) -> bool: + code, message = error.error.code, error.error.message + return (code, message) in _CLIENT_ERRORS or message.startswith(_CLIENT_ERROR_PREFIXES) + + +class _TokenUnavailableError(Exception): + """Raised, for this module only, when no token could be fetched for a gateway request.""" + + +@dataclass(eq=False) +class _Operation: + """ + One toolset operation (connecting, listing tools or calling a tool) and the gateway errors seen while it ran. + + The MCP client does not say which HTTP response made a call fail, so the error statuses are + only attributed to the operation when no other operation of the toolset ran at the same time. + A request that got no response at all ends the MCP session, which fails every operation + running then, so those failures are kept whatever the overlap. + """ + + overlapped: bool = False + errors: list[tuple[int, float | None]] = field(default_factory=list) + failures: list[Exception] = field(default_factory=list) + + @property + def http_error(self) -> tuple[int, float | None] | None: + """Return the status and ``Retry-After`` of the last error response it got, if it is known.""" + return self.errors[-1] if self.errors and not self.overlapped else None + + +class _DatabricksTokenAuth(httpx2.Auth): + """ + Authenticate each gateway request with a token from the Databricks connection. + + Asking the hook on every request, rather than once, lets OAuth tokens refresh during a long + agent run and on reconnection; the hook caches tokens until they are about to expire. + """ + + def __init__(self, hook: DatabricksHook) -> None: + self._hook = hook + self._token_lock = threading.Lock() + self._masked_token: str | None = None + + def fetch_token(self) -> str: + with self._token_lock: + token = self._hook._get_token(raise_error=False) + if not token: + raise ValueError( + f"Connection {self._hook.databricks_conn_id!r} has no token-based authentication " + "configured. This toolset sends a bearer token: use a personal access token, " + "service principal OAuth, Azure AD, or workload identity federation. Username and " + "password authentication is not supported." + ) + return token + + def fetch_and_mask_token(self) -> None: + """Fetch a token and mask it; only call it through ``AirflowToolset.run_blocking``.""" + self._mask(self.fetch_token()) + + def _mask(self, token: str) -> None: + # A personal access token is masked when the connection is fetched; mask minted + # OAuth tokens too, so they never reach task logs. + mask_secret(token) + self._masked_token = token + + async def async_auth_flow( + self, request: httpx2.Request + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + try: + # 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 + # resolved when the toolset connects, so this 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.fetch_token) Review Comment: Done in 3056dc34: before 3.2 the token fetch goes through run_blocking; 3.2+ keeps the toolset's own lock. A test checks every fetch holds the lock before 3.2. -- 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]
