firasbouzazi commented on code in PR #74198: URL: https://github.com/apache/airflow/pull/74198#discussion_r4323958493
########## 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", Review Comment: Changed to "SSE stream ". Also checked mcp 2.3.0's client source: all 12 errors it makes up are matched. -- 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]
