firasbouzazi commented on code in PR #74198: URL: https://github.com/apache/airflow/pull/74198#discussion_r4239613510
########## providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py: ########## @@ -0,0 +1,339 @@ +# 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 re +import threading +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException +from airflow.providers.databricks.exceptions import ( + DatabricksUnityMCPAccessDeniedError, + DatabricksUnityMCPError, + DatabricksUnityMCPServiceNotFoundError, + DatabricksUnityMCPThrottledError, + DatabricksUnityMCPTransportError, +) +from airflow.providers.databricks.hooks.databricks import DatabricksHook + +try: + import httpx2 + from fastmcp.client.transports import StreamableHttpTransport + from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset + + from airflow.providers.common.ai.toolsets.mcp import MCPToolset +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, Sequence + + from pydantic_ai._run_context import RunContext + from pydantic_ai.toolsets.abstract import ToolsetTool + + from airflow.sdk.execution_time.secrets_masker import mask_secret +else: + try: + from airflow.sdk.log import mask_secret + except ImportError: + try: + from airflow.sdk.execution_time.secrets_masker import mask_secret + except ImportError: + from airflow.utils.log.secrets_masker import mask_secret + +GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services" + +# Restricting each part to these characters keeps the name from adding a path, query or +# fragment to the URL, so the request can only reach the service path on the connection's host. +_SERVICE_NAME_PART = r"[A-Za-z0-9_-]+" +_SERVICE_NAME = re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}") + + +def validate_service_name(service_name: str) -> None: + """ + Raise ``ValueError`` unless ``service_name`` is a three-level ``catalog.schema.service`` name. + + Each part may contain only ASCII letters, digits, underscores and hyphens. + """ + if not isinstance(service_name, str) or not _SERVICE_NAME.fullmatch(service_name): + raise ValueError( + f"Invalid Unity Gateway MCP Service name {service_name!r}: expected " + "'catalog.schema.service', each part made of ASCII letters, digits, '_' or '-'." + ) + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + try: + seconds = float(value) + except ValueError: + # The HTTP-date form is not worth parsing: the caller only uses this as a hint. + return None + return seconds if seconds >= 0 else None + + +def _find_transport_error(exc: BaseException) -> httpx2.TransportError | None: + """Return the network error behind ``exc``, looking through causes and exception groups.""" + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, httpx2.TransportError): + return current + # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio uses the backport on 3.10. + if isinstance(grouped := getattr(current, "exceptions", None), (list, tuple)): + pending.extend(e for e in grouped if isinstance(e, BaseException)) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + return None + + +class _DatabricksTokenAuth(httpx2.Auth): + """ + Authenticate each gateway request with a token from the Databricks connection. + + Asking the hook on every request, rather than once, lets OAuth tokens refresh during a long + agent run and on reconnection; the hook caches tokens until they are about to expire. + + The MCP client reports every HTTP error from the gateway as the same generic error, so the + status and ``Retry-After`` of the last error response, and any failure to get a token, are + kept here for :class:`DatabricksUnityMCPToolset` to report what went wrong. With concurrent + calls on one toolset, the error reported for a failed call can be another call's. + """ + + def __init__(self, hook: DatabricksHook) -> None: + self._hook = hook + self.error_status: int | None = None + self.retry_after: float | None = None + self.token_error: Exception | None = None + 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 + + def take_error(self) -> tuple[int | None, float | None, Exception | None]: + """Return the last recorded error and forget it.""" + error = (self.error_status, self.retry_after, self.token_error) + self.error_status = self.retry_after = self.token_error = None + return error + + def _record(self, response: httpx2.Response) -> None: Review Comment: there's no shared error state any more: the fatal-vs-retry decision comes from the call's own exception chain. A client-synthesized `McpError` on `tools/call` raises `DatabricksUnityMCPError` and is never retried. The HTTP status is still recorded, but only for JSON-RPC requests, and it's only used for the message and exception type when the operation ran alone; with parallel calls it's left out. `test_concurrent_calls_are_each_judged_by_their_own_failure` holds two calls at a barrier so they really overlap (503 + 200, 503 + 429, 400-with-JSON-RPC + 503). -- 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]
