Add a Databricks toolset for Unity Gateway MCP Services

docs/spelling_wordlist.txt

@@ -1078,6 +1078,7 @@ Maxime
 MaxRuntimeInSeconds
 mb
 MCP
+mcp
 md
 mem
 memcached

providers/databricks/docs/index.rst

@@ -36,6 +36,7 @@
 
     Connection types <connections/databricks>
     Operators <operators/index>
+    Toolsets <toolsets/index>
     Plugins <plugins/index>
 
 .. toctree::
@@ -133,6 +134,7 @@ You can install such cross-provider dependencies when installing from PyPI. For
 Dependent package                                                                                               Extra
 ==============================================================================================================  ===============
 `apache-airflow-providers-amazon <https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_            ``amazon``
+`apache-airflow-providers-common-ai <https://airflow.apache.org/docs/apache-airflow-providers-common-ai>`_      ``common.ai``
 `apache-airflow-providers-google <https://airflow.apache.org/docs/apache-airflow-providers-google>`_            ``google``
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_  ``openlineage``
 ==============================================================================================================  ===============
@@ -154,6 +156,7 @@ Extra               Dependencies
 ``avro``            ``fastavro>=1.9.0; python_version<"3.14"``, ``fastavro>=1.10.0; python_version>="3.12" and python_version<"3.14"``, ``fastavro>=1.12.1; python_version>="3.14"``
 ``amazon``          ``apache-airflow-providers-amazon>=9.22.0``
 ``azure-identity``  ``azure-identity>=1.25.3``
+``common.ai``       ``apache-airflow-providers-common-ai[mcp]>=0.11.0``
 ``fab``             ``apache-airflow-providers-fab>=2.2.0``
 ``google``          ``apache-airflow-providers-google>=10.24.0``
 ``sdk``             ``databricks-sdk==0.10.0``

providers/databricks/docs/toolsets/index.rst

@@ -0,0 +1,33 @@
+ .. 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.
+
+
+Databricks Toolsets
+===================
+
+Toolsets give agents from the :doc:`Common AI provider <apache-airflow-providers-common-ai:index>`
+tools backed by Databricks. They need the ``common.ai`` extra:
+
+.. code-block:: bash
+
+    pip install 'apache-airflow-providers-databricks[common.ai]'
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *

providers/databricks/docs/toolsets/unity_mcp.rst

@@ -0,0 +1,111 @@
+ .. 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, or the connection dropped.
+    * - ``DatabricksUnityMCPError``
+      - Any other gateway error, or a failure to get a token from the connection.
+
+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
+calls that are safe to repeat.

providers/databricks/provider.yaml

@@ -144,6 +144,9 @@ integrations:
     how-to-guide:
       - /docs/apache-airflow-providers-databricks/operators/workflow.rst
     tags: [service]
+  - integration-name: Databricks Unity Gateway
+    external-doc-url: https://docs.databricks.com/aws/en/unity-gateway/concepts
+    tags: [service]
 
 operators:
   - integration-name: Databricks
@@ -195,6 +198,11 @@ sensors:
       - airflow.providers.databricks.sensors.databricks_sql
       - airflow.providers.databricks.sensors.databricks_partition
 
+toolsets:
+  - integration-name: Databricks Unity Gateway
+    python-modules:
+      - airflow.providers.databricks.toolsets.unity_mcp
+
 connection-types:
   - hook-class-name: airflow.providers.databricks.hooks.databricks.DatabricksHook
     hook-name: "Databricks"

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
+]
 "fab" = [
     "apache-airflow-providers-fab>=2.2.0"
 ]
@@ -113,6 +116,7 @@ dev = [
     "apache-airflow-task-sdk",
     "apache-airflow-devel-common",
     "apache-airflow-providers-amazon",
+    "apache-airflow-providers-common-ai",
     "apache-airflow-providers-common-compat",
     "apache-airflow-providers-common-sql",
     "apache-airflow-providers-google",
@@ -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]",
     "apache-airflow-providers-common-sql[pandas,polars]",
     "apache-airflow-providers-fab",
     "apache-airflow-providers-databricks[sqlalchemy]",

providers/databricks/src/airflow/providers/databricks/exceptions.py

@@ -46,3 +46,34 @@ class DatabricksApiError(AirflowException):
     def __init__(self, message: str, *, http_status_code: int | None = None) -> None:
         super().__init__(message)
         self.http_status_code = http_status_code
+
+
+class DatabricksUnityMCPError(DatabricksApiError):
+    """Raised when a call to a Unity Gateway MCP Service fails."""
+
+
+class DatabricksUnityMCPAccessDeniedError(DatabricksUnityMCPError):
+    """
+    Raised when the gateway will not let the caller invoke the service.
+
+    The service may not exist, the caller may lack a privilege on it, or its credentials may be invalid.
+    """
+
+
+class DatabricksUnityMCPThrottledError(DatabricksUnityMCPError):
+    """Raised when the gateway rate-limits the caller."""
+
+    def __init__(
+        self, message: str, *, http_status_code: int | None = None, retry_after: float | None = None
+    ) -> None:
+        super().__init__(message, http_status_code=http_status_code)
+        self.retry_after = retry_after
+
+
+class DatabricksUnityMCPTransportError(DatabricksUnityMCPError):
+    """
+    Raised when the gateway cannot be reached or the connection drops.
+
+    When this interrupts a tool call, the tool may or may not have run, so the call is not
+    retried automatically.
+    """

providers/databricks/src/airflow/providers/databricks/get_provider_info.py

@@ -67,6 +67,11 @@ def get_provider_info():
                 "how-to-guide": ["/docs/apache-airflow-providers-databricks/operators/workflow.rst"],
                 "tags": ["service"],
             },
+            {
+                "integration-name": "Databricks Unity Gateway",
+                "external-doc-url": "https://docs.databricks.com/aws/en/unity-gateway/concepts",
+                "tags": ["service"],
+            },
         ],
         "operators": [
             {

providers/databricks/src/airflow/providers/databricks/toolsets/__init__.py

@@ -0,0 +1,17 @@
+#
+# 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.