Repository navigation
Databricks: add Unity Gateway MCP integration for common.ai #74193
Description
Description
Add a Databricks-provider adapter that resolves a Unity MCP service and uses the existing generic MCP toolset machinery. Keep Databricks authentication and service addressing in the Databricks provider, with an optional dependency on common.ai.
Use case/motivation
A DAG author should be able to give an in-worker agent governed tools using a Databricks connection and a service name, without copying bearer tokens or gateway URLs into DAG code. One MCP integration can expose services such as UC functions, Genie and AI Search where those services are configured and authorized.
Proposed scope and acceptance
- Reuse existing Databricks connection authentication; document the caller identity and supported auth modes.
- Validate service names and construct URLs within the configured workspace; never forward credentials to an arbitrary caller-supplied host.
- Refresh credentials across long-lived calls/reconnections; retain generic transport behavior in common.ai.
- Translate access denial, missing service, throttling and transport errors distinctly. Do not blindly retry mutating tool calls after ambiguous delivery.
- Keep common.ai optional: importing unrelated Databricks operators must continue to work without AI dependencies.
- Test authentication refresh, URL construction, denied access, optional-dependency behavior and tool discovery.
- Provide an opt-in system test: AgentOperator → configured Unity MCP service → list tools → harmless tool call. Record actual live validation separately from mock tests.
Start with invocation of an existing service. Service administration/CRUD and SDK modernization are separate follow-ups.
Related issues
Independent of #74191; does not block or expand that PR.
References: https://docs.databricks.com/aws/en/unity-gateway/concepts and https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/index.html
Activity
garvit-arora commentedon Oct 10, 2026
@zozo123 I would like to take this on, starting with the bounded adapter slice: resolve an existing Unity Gateway MCP service from a Databricks connection, validate the service name and workspace URL, refresh credentials, and hand the transport to the existing common.ai MCP toolset. I will keep service administration and SDK modernization out of the PR, preserve common.ai as optional, and include mocked coverage for auth refresh, errors, URL validation and tool discovery. I am already working on #74194's native Genie hook/toolset; this would complement it through the managed MCP route rather than overlap it. Could you assign #74193 to @garvit-arora?
Add a Databricks toolset for Unity AI Gateway MCP Services
garvit-arora commentedon Oct 11, 2026
Correction: I missed #74198 when I posted my assignment follow-up. That PR already covers the Unity Gateway MCP toolset and is actively being reviewed, so please disregard my requests to take this slice and publish a parallel PR. I will follow the existing PR's review instead. Apologies for the duplicate ping, @zozo123.