Databricks MCP: throw a better error for incompatible workspace clients (#263)

Signed-off-by: Bryan Qiu <bryan.qiu@databricks.com>

Author: bbqiu

databricks_mcp/src/databricks_mcp/mcp.py

@@ -23,6 +23,28 @@
 
 logger = logging.getLogger(__name__)
 
+
+def _is_databricks_apps_url(url: str) -> bool:
+    """Check if the URL is hosted on Databricks Apps."""
+    parsed = urlparse(url)
+    return parsed.netloc.endswith(".databricksapps.com")
+
+
+def _is_oauth_auth(workspace_client: WorkspaceClient) -> bool:
+    """Check if the workspace client is using OAuth authentication.
+
+    Uses the SDK's oauth_token() method to determine if OAuth is available.
+    This is more resilient than checking auth_type directly, as it handles
+    various non-OAuth auth types (pat, runtime, etc.).
+    """
+    try:
+        workspace_client.config.oauth_token()
+        return True
+    except ValueError:
+        # oauth_token() raises ValueError when not using OAuth-based auth
+        return False
+
+
 # MCP URL types
 UC_FUNCTIONS_MCP = "uc_functions_mcp"
 VECTOR_SEARCH_MCP = "vector_search_mcp"
@@ -123,6 +145,15 @@ def __init__(self, server_url: str, workspace_client: Optional[WorkspaceClient]
         self.client = workspace_client or WorkspaceClient()
         self.server_url = server_url
 
+        # Early detection: error if using non-OAuth auth with Databricks Apps
+        if _is_databricks_apps_url(server_url) and not _is_oauth_auth(self.client):
+            raise ValueError(
+                "OAuth authentication is required for MCP servers hosted on Databricks Apps. "
+                "Your current authentication method is not supported. "
+                "Please use OAuth authentication instead. "
+                "For more information: https://docs.databricks.com/aws/en/generative-ai/mcp/custom-mcp"
+            )
+
     def _get_databricks_managed_mcp_url_type(self) -> str:
         """Determine the MCP URL type based on the path."""
         path = urlparse(self.server_url).path

databricks_mcp/src/databricks_mcp/oauth_provider.py

@@ -1,4 +1,5 @@
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors.platform import PermissionDenied
 from mcp.client.auth import OAuthClientProvider, TokenStorage
 from mcp.shared.auth import OAuthToken
 
@@ -47,6 +48,17 @@ class DatabricksOAuthClientProvider(OAuthClientProvider):
 
     def __init__(self, workspace_client: WorkspaceClient):
         self.workspace_client = workspace_client
+
+        # Pre-flight check: verify the workspace client has basic permissions
+        try:
+            workspace_client.current_user.me()
+        except PermissionDenied as e:
+            raise PermissionError(
+                f"The workspace client does not have permission to access the Databricks workspace. "
+                f"Please ensure the service principal or user has the required permissions to call Databricks APIs. "
+                f"Original error: {e}"
+            ) from e
+
         self.databricks_token_storage = DatabricksTokenStorage(workspace_client)
 
         super().__init__(

databricks_mcp/tests/integration_tests/test_mcp_client.py

@@ -0,0 +1,35 @@
+from databricks.sdk import WorkspaceClient
+
+from databricks_mcp import DatabricksMCPClient
+
+# Replace with your deployed app URL
+mcp_server_url = "https://mcp-chloe-test-6051921418418893.staging.aws.databricksapps.com/mcp"
+
+workspace_client = WorkspaceClient(
+    host="https://e2-dogfood.staging.cloud.databricks.com",
+    token="<token>",
+)
+
+# print(workspace_client.current_user.me())
+
+mcp_client = DatabricksMCPClient(server_url=mcp_server_url, workspace_client=workspace_client)
+
+# List available tools
+tools = mcp_client.list_tools()
+# print(f"Available tools: {tools}")
+
+from databricks.sdk import WorkspaceClient
+
+from databricks_mcp import DatabricksMCPClient
+
+ws_client = WorkspaceClient(
+    host="https://e2-dogfood.staging.cloud.databricks.com/",
+    client_id="<client-id>",
+    client_secret="<client-secret>",
+)
+
+mcp_client = DatabricksMCPClient(
+    server_url=mcp_server_url,
+    workspace_client=ws_client,
+)
+# print(mcp_client.list_tools())

databricks_mcp/tests/unit_tests/test_mcp.py

@@ -11,6 +11,8 @@
     UC_FUNCTIONS_MCP,
     VECTOR_SEARCH_MCP,
     DatabricksMCPClient,
+    _is_databricks_apps_url,
+    _is_oauth_auth,
 )
 
 
@@ -131,6 +133,7 @@ async def test_get_tools_async(self):
         with (
             patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
             patch("databricks_mcp.mcp.ClientSession") as mock_session_class,
+            patch("databricks_mcp.mcp.DatabricksOAuthClientProvider"),
         ):
             mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), None)
             mock_session_class.return_value.__aenter__.return_value = mock_session
@@ -156,6 +159,7 @@ async def test_call_tools_async(self):
         with (
             patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
             patch("databricks_mcp.mcp.ClientSession") as mock_session_class,
+            patch("databricks_mcp.mcp.DatabricksOAuthClientProvider"),
         ):
             mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), None)
             mock_session_class.return_value.__aenter__.return_value = mock_session
@@ -434,6 +438,7 @@ def test_error_decorator_managed_server_reraises_original(self):
                 client, "_get_databricks_managed_mcp_url_type", return_value=UC_FUNCTIONS_MCP
             ),
             patch("databricks_mcp.mcp.streamablehttp_client") as mock_client,
+            patch("databricks_mcp.mcp.DatabricksOAuthClientProvider"),
             patch("requests.request") as mock_request,
         ):
             mock_client.side_effect = original_error
@@ -442,3 +447,57 @@ def test_error_decorator_managed_server_reraises_original(self):
                 client.list_tools()
 
             mock_request.assert_not_called()
+
+
+class TestIsDatabricksAppsUrl:
+    """Test cases for _is_databricks_apps_url helper function."""
+
+    @pytest.mark.parametrize(
+        "url,expected",
+        [
+            ("https://my-app.staging.aws.databricksapps.com/mcp", True),
+            ("https://my-app.prod.azure.databricksapps.com/mcp", True),
+            ("https://my-app.databricksapps.com", True),
+            ("https://test.cloud.databricks.com/api/2.0/mcp/functions/a/b", False),
+            ("https://custom-server.example.com/mcp", False),
+            ("https://databricksapps.com.evil.com/mcp", False),
+            ("https://notdatabricksapps.com/mcp", False),
+        ],
+    )
+    def test_is_databricks_apps_url(self, url, expected):
+        assert _is_databricks_apps_url(url) == expected
+
+
+class TestIsOauthAuth:
+    @pytest.mark.parametrize(
+        "side_effect,expected",
+        [
+            (None, True),  # oauth_token succeeds
+            (ValueError("not available"), False),  # oauth_token raises
+        ],
+    )
+    def test_is_oauth_auth(self, side_effect, expected):
+        mock_client = MagicMock(spec=WorkspaceClient)
+        if side_effect:
+            mock_client.config.oauth_token.side_effect = side_effect
+        assert _is_oauth_auth(mock_client) is expected
+
+
+class TestDatabricksMCPClientOAuthValidation:
+    @pytest.mark.parametrize("auth_type", ["pat", "runtime"])
+    def test_raises_error_for_non_oauth_with_databricks_apps(self, auth_type):
+        mock_client = MagicMock(spec=WorkspaceClient)
+        mock_client.config.oauth_token.side_effect = ValueError(f"not available for {auth_type}")
+        with pytest.raises(ValueError, match="OAuth authentication is required"):
+            DatabricksMCPClient("https://my-app.databricksapps.com/mcp", mock_client)
+
+    def test_allows_oauth_with_databricks_apps(self):
+        mock_client = MagicMock(spec=WorkspaceClient)
+        client = DatabricksMCPClient("https://my-app.databricksapps.com/mcp", mock_client)
+        assert client.server_url == "https://my-app.databricksapps.com/mcp"
+
+    def test_allows_non_oauth_with_non_databricks_apps(self):
+        mock_client = MagicMock(spec=WorkspaceClient)
+        mock_client.config.oauth_token.side_effect = ValueError("not available")
+        client = DatabricksMCPClient("https://test.com/api/2.0/mcp/functions/a/b", mock_client)
+        assert client.server_url == "https://test.com/api/2.0/mcp/functions/a/b"

databricks_mcp/tests/unit_tests/test_oauth_provider.py

@@ -1,34 +1,54 @@
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors.platform import PermissionDenied
 
 from databricks_mcp import DatabricksOAuthClientProvider
 
 
 @pytest.mark.asyncio
 async def test_oauth_provider():
     workspace_client = WorkspaceClient(host="https://test-databricks.com", token="test-token")
-    provider = DatabricksOAuthClientProvider(workspace_client=workspace_client)
-    oauth_token = await provider.context.storage.get_tokens()
-    assert oauth_token.access_token == "test-token"
-    assert oauth_token.expires_in == 60
-    assert oauth_token.token_type.lower() == "bearer"
+    with patch.object(workspace_client.current_user, "me", return_value=MagicMock()):
+        provider = DatabricksOAuthClientProvider(workspace_client=workspace_client)
+        oauth_token = await provider.context.storage.get_tokens()
+        assert oauth_token.access_token == "test-token"
+        assert oauth_token.expires_in == 60
+        assert oauth_token.token_type.lower() == "bearer"
 
 
 @pytest.mark.asyncio
 async def test_authenticate_raises_exception():
     workspace_client = WorkspaceClient(host="https://test-databricks.com", token="test-token")
 
+    with patch.object(workspace_client.current_user, "me", return_value=MagicMock()):
+        with patch.object(
+            workspace_client.config, "authenticate", return_value={"Authorization": "Basic abc123"}
+        ):
+            with pytest.raises(
+                ValueError, match="Invalid authentication token format. Expected Bearer token."
+            ):
+                provider = DatabricksOAuthClientProvider(workspace_client=workspace_client)
+
+                oauth_token = await provider.context.storage.get_tokens()
+                assert oauth_token.access_token == "test-token"
+                assert oauth_token.expires_in == 60
+                assert oauth_token.token_type.lower() == "bearer"
+
+
+def test_preflight_check_raises_permission_denied():
+    workspace_client = WorkspaceClient(host="https://test-databricks.com", token="test-token")
+
     with patch.object(
-        workspace_client.config, "authenticate", return_value={"Authorization": "Basic abc123"}
+        workspace_client.current_user,
+        "me",
+        side_effect=PermissionDenied(
+            "This API is disabled for users without the workspace-access entitlement."
+        ),
     ):
         with pytest.raises(
-            ValueError, match="Invalid authentication token format. Expected Bearer token."
+            PermissionError,
+            match="The workspace client does not have permission to access the Databricks workspace",
         ):
-            provider = DatabricksOAuthClientProvider(workspace_client=workspace_client)
-
-            oauth_token = await provider.context.storage.get_tokens()
-            assert oauth_token.access_token == "test-token"
-            assert oauth_token.expires_in == 60
-            assert oauth_token.token_type.lower() == "bearer"
+            DatabricksOAuthClientProvider(workspace_client=workspace_client)