MCP Support in Workspace Client

databricks/sdk/service/mcp.py

@@ -0,0 +1,27 @@
+from mcp.client.auth import OAuthClientProvider, TokenStorage
+from mcp.shared.auth import OAuthToken
+
+
+class DatabricksTokenStorage(TokenStorage):
+    def __init__(self, config):
+        self.config = config
+
+    async def get_tokens(self) -> OAuthToken | None:
+        headers = self.config.authenticate()
+        token = headers["Authorization"].split("Bearer ")[1]
+        return OAuthToken(access_token=token, expires_in=60)
+
+
+class MCP:
+    def __init__(self, config):
+        self._config = config
+        self.databricks_token_storage = DatabricksTokenStorage(config)
+
+    def get_oauth_provider(self):
+        return OAuthClientProvider(
+            server_url="",
+            client_metadata=None,
+            storage=self.databricks_token_storage,
+            redirect_handler=None,
+            callback_handler=None,
+        )

pyproject.toml

@@ -53,6 +53,8 @@ dev = [
     'langchain-openai; python_version > "3.7"',
     "httpx",
     "build", # some integration tests depend on the databricks-sdk-py wheel
+    "mcp>=1.9.1",
+    "pytest-asyncio"
 ]
 notebook = [
     "ipython>=8,<10",
@@ -63,6 +65,9 @@ openai = [
     'langchain-openai; python_version > "3.7"',
     "httpx",
 ]
+mcp = [
+    "mcp>=1.9.1"
+]
 
 [tool.setuptools.dynamic]
 version = { attr = "databricks.sdk.version.__version__" }

tests/test_mcp.py

@@ -0,0 +1,20 @@
+import time
+
+import httpx
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_mcp_oauth_provider(monkeypatch):
+    monkeypatch.setattr(time, "time", lambda: 100)
+    from databricks.sdk import WorkspaceClient
+
+    monkeypatch.setenv("DATABRICKS_HOST", "test_host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "test_token")
+
+    w = WorkspaceClient()
+    mcp_oauth_provider = w.mcp.get_oauth_provider()
+
+    request = httpx.Request("GET", "https://example.com")
+    response = await anext(mcp_oauth_provider.async_auth_flow(request))
+    assert response.headers["Authorization"] == "Bearer test_token"