shahar1 commented on code in PR #74198:
URL: https://github.com/apache/airflow/pull/74198#discussion_r4310330713
##########
providers/databricks/pyproject.toml:
##########
@@ -88,6 +88,9 @@ dependencies = [
"azure-identity" = [
"azure-identity>=1.25.3",
]
+"common.ai" = [
+ "apache-airflow-providers-common-ai[mcp]>=0.11.0", # use next version
Review Comment:
Dependency range allows broken FastMCP versions. The declared dependencies
can install FastMCP 3.3.0–3.4.9, which are incompatible with the HTTP and MCP
versions used by this code. The lockfile pins 4.0.5, so CI doesn’t catch the
problem. Raise the minimum FastMCP version to 4.0.0.
##########
providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py:
##########
@@ -0,0 +1,494 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP
Service."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import ipaddress
+import json
+import re
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlsplit
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.exceptions import (
+ DatabricksUnityMCPAccessDeniedError,
+ DatabricksUnityMCPError,
+ DatabricksUnityMCPThrottledError,
+ DatabricksUnityMCPTransportError,
+)
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+
+try:
+ # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the
common.ai extra installs.
+ import httpx2
+ from fastmcp.client.transports import StreamableHttpTransport
+ from fastmcp.exceptions import McpError
+ from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset
+
+ from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "DatabricksUnityMCPToolset needs the 'common.ai' extra of the
databricks provider: "
+ "pip install 'apache-airflow-providers-databricks[common.ai]'"
+ ) from e
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncGenerator, AsyncIterator, Iterator,
Sequence
+
+ from pydantic_ai._run_context import RunContext
+ from pydantic_ai.toolsets.abstract import ToolsetTool
+
+ from airflow.providers.common.ai.tools import AirflowTool
+ from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+ try:
+ from airflow.sdk.log import mask_secret
+ except ImportError:
+ try:
+ from airflow.sdk.execution_time.secrets_masker import mask_secret
+ except ImportError:
+ from airflow.utils.log.secrets_masker import mask_secret
+
+GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services"
+
+# Restricting each part to these characters keeps the name from adding a path,
query or
+# fragment to the URL, so the request can only reach the service path on the
connection's host.
+_SERVICE_NAME_PART = r"[A-Za-z0-9_]+"
+_SERVICE_NAME =
re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}")
+
+# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the
caller may not invoke
+# the service, including when no service has that name.
+_NOT_AUTHORIZED_CODE = -32007
+
+# Errors the MCP client makes up when the gateway gave no JSON-RPC answer,
such as an HTTP error
+# with a plain body or a response that never arrived (see
mcp.client.streamable_http). Every other
+# MCP error is the server's own answer. Matched on code and message, the only
things they carry.
+_CLIENT_ERRORS = {
+ (-32603, "Server returned an error response"),
+ (-32601, "Not Found"),
+ (-32600, "Session terminated"),
+ (-32000, "Connection closed"),
+}
+# What fastmcp raises, in place of the TimeoutError it loses, when the MCP
handshake outlasts
+# init_timeout (see fastmcp.client.client).
+_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session"
+_CLIENT_ERROR_PREFIXES = (
+ "Unexpected content type:",
+ "Failed to parse JSON response:",
+ "Failed to parse SSE message:",
+ "SSE stream ended without a response",
+ "server answered a request with 202 Accepted",
+ "Redirect to ",
+)
+
+
+def validate_service_name(service_name: str) -> None:
+ """
+ Raise ``ValueError`` unless ``service_name`` is a three-level
``catalog.schema.service`` name.
+
+ Each part may contain only ASCII letters, digits and underscores.
+ """
+ if not isinstance(service_name, str) or not
_SERVICE_NAME.fullmatch(service_name):
+ raise ValueError(
+ f"Invalid Unity Gateway MCP Service name {service_name!r}:
expected "
+ "'catalog.schema.service', each part made of ASCII letters, digits
or '_'."
+ )
+
+
+def _is_loopback(host: str) -> bool:
+ if host == "localhost":
+ return True
+ try:
+ return ipaddress.ip_address(host.strip("[]")).is_loopback
+ except ValueError:
+ return False
+
+
+def _parse_retry_after(value: str | None) -> float | None:
+ if value is None:
+ return None
+ try:
+ seconds = float(value)
+ except ValueError:
+ # The HTTP-date form is not worth parsing: the caller only uses this
as a hint.
+ return None
+ return seconds if seconds >= 0 else None
+
+
+def _walk_exceptions(exc: BaseException) -> Iterator[BaseException]:
+ """Yield ``exc``, its causes and the exceptions inside any exception
groups, each once."""
+ seen: set[int] = set()
+ pending: list[BaseException] = [exc]
+ while pending:
+ current = pending.pop(0)
+ if id(current) in seen:
+ continue
+ seen.add(id(current))
+ yield current
+ # Duck-typed: the ExceptionGroup builtin is Python 3.11+, and anyio
uses the backport on 3.10.
+ if isinstance(grouped := getattr(current, "exceptions", None), (list,
tuple)):
+ pending.extend(e for e in grouped if isinstance(e, BaseException))
+ pending.extend(e for e in (current.__cause__, current.__context__) if
e is not None)
+
+
+def _find(exc: BaseException, types: type[BaseException] |
tuple[type[BaseException], ...]) -> Any:
+ """Return the first exception of ``types`` in ``exc``, its causes and any
exception groups."""
+ return next((e for e in _walk_exceptions(exc) if isinstance(e, types)),
None)
+
+
+def _is_handshake_timeout(error: BaseException) -> bool:
+ return any(
+ isinstance(e, TimeoutError)
+ or (isinstance(e, RuntimeError) and
str(e).endswith(_HANDSHAKE_TIMEOUT_MESSAGE))
+ for e in _walk_exceptions(error)
+ )
+
+
+def _is_client_error(error: McpError) -> bool:
+ code, message = error.error.code, error.error.message
+ return (code, message) in _CLIENT_ERRORS or
message.startswith(_CLIENT_ERROR_PREFIXES)
+
+
+class _TokenUnavailableError(Exception):
+ """Raised, for this module only, when no token could be fetched for a
gateway request."""
+
+
+@dataclass(eq=False)
+class _Operation:
+ """
+ One toolset operation (connecting, listing tools or calling a tool) and
the gateway errors seen while it ran.
+
+ The MCP client does not say which HTTP response made a call fail, so the
errors are only
+ attributed to the operation when no other operation of the toolset ran at
the same time.
+ """
+
+ overlapped: bool = False
+ errors: list[tuple[int, float | None]] = field(default_factory=list)
+
+ @property
+ def http_error(self) -> tuple[int, float | None] | None:
+ """Return the status and ``Retry-After`` of the last error response it
got, if it is known."""
+ return self.errors[-1] if self.errors and not self.overlapped else None
+
+
+class _DatabricksTokenAuth(httpx2.Auth):
+ """
+ Authenticate each gateway request with a token from the Databricks
connection.
+
+ Asking the hook on every request, rather than once, lets OAuth tokens
refresh during a long
+ agent run and on reconnection; the hook caches tokens until they are about
to expire.
+ """
+
+ def __init__(self, hook: DatabricksHook, operations: set[_Operation]) ->
None:
+ self._hook = hook
+ self._operations = operations
+ self._token_lock = threading.Lock()
+
+ def get_token(self) -> str:
+ with self._token_lock:
+ token = self._hook._get_token(raise_error=False)
+ if not token:
+ raise ValueError(
+ f"Connection {self._hook.databricks_conn_id!r} has no
token-based authentication "
+ "configured. This toolset sends a bearer token: use a personal
access token, "
+ "service principal OAuth, Azure AD, or workload identity
federation. Username and "
+ "password authentication is not supported."
+ )
+ # A personal access token is masked when the connection is fetched;
mask minted
+ # OAuth tokens too, so they never reach task logs.
+ mask_secret(token)
+ return token
+
+ async def async_auth_flow(
+ self, request: httpx2.Request
+ ) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+ try:
+ # Not AirflowToolset.run_blocking: its lock is shared by every
toolset, so a long SQL
+ # query in another toolset would hold up each request here. The
connection is already
+ # resolved when the toolset connects, so fetching a token only
calls the token
+ # endpoint, and the hook is this toolset's own, so a lock of its
own is enough.
+ token = await asyncio.to_thread(self.get_token)
+ except Exception as e:
+ raise _TokenUnavailableError(str(e)) from e
+ request.headers["Authorization"] = f"Bearer {token}"
+ response = yield request
+ if response.status_code >= 400 and _is_jsonrpc_request(request):
+ retry_after =
_parse_retry_after(response.headers.get("Retry-After"))
+ for operation in self._operations:
+ operation.errors.append((response.status_code, retry_after))
+
+
+def _is_jsonrpc_request(request: httpx2.Request) -> bool:
+ """
+ Whether ``request`` sends a JSON-RPC request, which expects an answer.
+
+ Not a notification or a GET for server events: the MCP client tolerates
errors on those, so
+ they say nothing about why an operation failed.
+ """
+ if request.method != "POST":
+ return False
+ try:
+ return "id" in json.loads(request.content)
+ except (ValueError, TypeError):
+ return False
+
+
+class DatabricksUnityMCPToolset(AirflowToolset):
+ """
+ Give an agent the tools of a Unity Gateway MCP Service, authenticated as
the connection's identity.
+
+ The service is named by its three-level Unity Catalog name,
``catalog.schema.service``, and
+ reached at ``https://<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)
Review Comment:
Timeouts are classified incorrectly. The code checks for an MCP error before
checking for a network timeout. FastMCP wraps some timeouts as MCP errors, so
they get treated as service errors. During a tool call, that path can return
`None` and let the model continue, even though the documented behavior is to
retry. Check for transport errors first; that change made the local suite pass
all 49 tests.
##########
providers/databricks/pyproject.toml:
##########
@@ -122,6 +126,7 @@ dev = [
"deltalake>=1.1.3,!=1.3.0",
"apache-airflow-providers-fab>=2.2.0",
"apache-airflow-providers-microsoft-azure",
+ "apache-airflow-providers-common-ai[mcp]",
Review Comment:
This dependency appears twice (it should be only as part of `[ccommon.ai]`)
##########
providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py:
##########
@@ -0,0 +1,494 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP
Service."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import ipaddress
+import json
+import re
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlsplit
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.exceptions import (
+ DatabricksUnityMCPAccessDeniedError,
+ DatabricksUnityMCPError,
+ DatabricksUnityMCPThrottledError,
+ DatabricksUnityMCPTransportError,
+)
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+
+try:
+ # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the
common.ai extra installs.
+ import httpx2
+ from fastmcp.client.transports import StreamableHttpTransport
+ from fastmcp.exceptions import McpError
+ from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset
+
+ from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "DatabricksUnityMCPToolset needs the 'common.ai' extra of the
databricks provider: "
+ "pip install 'apache-airflow-providers-databricks[common.ai]'"
+ ) from e
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncGenerator, AsyncIterator, Iterator,
Sequence
+
+ from pydantic_ai._run_context import RunContext
+ from pydantic_ai.toolsets.abstract import ToolsetTool
+
+ from airflow.providers.common.ai.tools import AirflowTool
+ from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+ try:
+ from airflow.sdk.log import mask_secret
+ except ImportError:
+ try:
+ from airflow.sdk.execution_time.secrets_masker import mask_secret
+ except ImportError:
+ from airflow.utils.log.secrets_masker import mask_secret
+
+GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services"
+
+# Restricting each part to these characters keeps the name from adding a path,
query or
+# fragment to the URL, so the request can only reach the service path on the
connection's host.
+_SERVICE_NAME_PART = r"[A-Za-z0-9_]+"
+_SERVICE_NAME =
re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}")
+
+# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the
caller may not invoke
+# the service, including when no service has that name.
+_NOT_AUTHORIZED_CODE = -32007
+
+# Errors the MCP client makes up when the gateway gave no JSON-RPC answer,
such as an HTTP error
+# with a plain body or a response that never arrived (see
mcp.client.streamable_http). Every other
+# MCP error is the server's own answer. Matched on code and message, the only
things they carry.
+_CLIENT_ERRORS = {
Review Comment:
Error classification depends on private message text. The code recognizes
errors by matching internal MCP/FastMCP message strings. Those strings could
change in a dependency update. Add a test that alerts maintainers if the
expected messages change.
##########
providers/databricks/src/airflow/providers/databricks/toolsets/unity_mcp.py:
##########
@@ -0,0 +1,494 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Toolset that gives a common.ai agent the tools of a Unity Gateway MCP
Service."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import ipaddress
+import json
+import re
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlsplit
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.exceptions import (
+ DatabricksUnityMCPAccessDeniedError,
+ DatabricksUnityMCPError,
+ DatabricksUnityMCPThrottledError,
+ DatabricksUnityMCPTransportError,
+)
+from airflow.providers.databricks.hooks.databricks import DatabricksHook
+
+try:
+ # httpx2 and fastmcp come with pydantic-ai's mcp extra, which the
common.ai extra installs.
+ import httpx2
+ from fastmcp.client.transports import StreamableHttpTransport
+ from fastmcp.exceptions import McpError
+ from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset
+
+ from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "DatabricksUnityMCPToolset needs the 'common.ai' extra of the
databricks provider: "
+ "pip install 'apache-airflow-providers-databricks[common.ai]'"
+ ) from e
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncGenerator, AsyncIterator, Iterator,
Sequence
+
+ from pydantic_ai._run_context import RunContext
+ from pydantic_ai.toolsets.abstract import ToolsetTool
+
+ from airflow.providers.common.ai.tools import AirflowTool
+ from airflow.sdk.execution_time.secrets_masker import mask_secret
+else:
+ try:
+ from airflow.sdk.log import mask_secret
+ except ImportError:
+ try:
+ from airflow.sdk.execution_time.secrets_masker import mask_secret
+ except ImportError:
+ from airflow.utils.log.secrets_masker import mask_secret
+
+GATEWAY_MCP_SERVICES_PATH = "ai-gateway/mcp-services"
+
+# Restricting each part to these characters keeps the name from adding a path,
query or
+# fragment to the URL, so the request can only reach the service path on the
connection's host.
+_SERVICE_NAME_PART = r"[A-Za-z0-9_]+"
+_SERVICE_NAME =
re.compile(rf"{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}\.{_SERVICE_NAME_PART}")
+
+# The JSON-RPC error code Unity Gateway answers with, at HTTP 403, when the
caller may not invoke
+# the service, including when no service has that name.
+_NOT_AUTHORIZED_CODE = -32007
+
+# Errors the MCP client makes up when the gateway gave no JSON-RPC answer,
such as an HTTP error
+# with a plain body or a response that never arrived (see
mcp.client.streamable_http). Every other
+# MCP error is the server's own answer. Matched on code and message, the only
things they carry.
+_CLIENT_ERRORS = {
+ (-32603, "Server returned an error response"),
+ (-32601, "Not Found"),
+ (-32600, "Session terminated"),
+ (-32000, "Connection closed"),
+}
+# What fastmcp raises, in place of the TimeoutError it loses, when the MCP
handshake outlasts
+# init_timeout (see fastmcp.client.client).
+_HANDSHAKE_TIMEOUT_MESSAGE = "Failed to initialize server session"
+_CLIENT_ERROR_PREFIXES = (
+ "Unexpected content type:",
+ "Failed to parse JSON response:",
+ "Failed to parse SSE message:",
+ "SSE stream ended without a response",
+ "server answered a request with 202 Accepted",
+ "Redirect to ",
+)
+
+
+def validate_service_name(service_name: str) -> None:
+ """
+ Raise ``ValueError`` unless ``service_name`` is a three-level
``catalog.schema.service`` name.
+
Review Comment:
nit —
[unity_mcp.rst:107](vscode-webview://08iqfetk3mhe481jb4fir51tvalvm437hili4m3ses6rqk9spgck/providers/databricks/docs/toolsets/unity_mcp.rst#L107):
up-front token failure raises ValueError, not DatabricksUnityMCPError as the
table says.
##########
providers/databricks/pyproject.toml:
##########
@@ -88,6 +88,9 @@ dependencies = [
"azure-identity" = [
"azure-identity>=1.25.3",
]
+"common.ai" = [
+ "apache-airflow-providers-common-ai[mcp]>=0.11.0", # use next version
Review Comment:
`# use next version` marker misused: common.ai untouched;
run_blocking/agent_template_fields shipped in common-ai 1.0.0 (2026-10-06).
Keep the floor >=1.0.0 (without the marker).
--
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]
