"""Databricks MCP Client wrapper."""
import sys
from dataclasses import dataclass
from typing import Any, List, Optional
from concurrent.futures import ThreadPoolExecutor
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient
@dataclass
class ToolInfo:
"""Information about a remote tool."""
name: str
description: str
input_schema: dict
class DatabricksMCPProxy:
"""Proxy client for Databricks MCP server."""
def __init__(self, host: str, app_url: str, access_token: str):
self.host = host
self.app_url = app_url.rstrip("/")
self.mcp_url = f"{self.app_url}/mcp"
self.access_token = access_token
self._workspace_client: Optional[WorkspaceClient] = None
self._mcp_client: Optional[DatabricksMCPClient] = None
self._tools: List[ToolInfo] = []
self._connected = False
def connect(self):
"""Establish connection to the Databricks MCP server."""
print(f"Creating WorkspaceClient for {self.host}...", file=sys.stderr)
self._workspace_client = WorkspaceClient(host=self.host, token=self.access_token)
print("WorkspaceClient created", file=sys.stderr)
print(f"Connecting to MCP server at {self.mcp_url}...", file=sys.stderr)
self._mcp_client = DatabricksMCPClient(server_url=self.mcp_url, workspace_client=self._workspace_client)
print("MCP client connected", file=sys.stderr)
self._connected = True
@property
def is_connected(self) -> bool:
return self._connected
def discover_tools(self) -> List[ToolInfo]:
"""Discover available tools from the remote MCP server."""
if not self._mcp_client:
raise RuntimeError("Not connected. Call connect() first.")
print("Discovering remote tools...", file=sys.stderr)
with ThreadPoolExecutor() as executor:
remote_tools = executor.submit(self._mcp_client.list_tools).result()
self._tools = [
ToolInfo(
name=t.name,
description=t.description or "",
input_schema=t.inputSchema if hasattr(t, 'inputSchema') else {"type": "object", "properties": {}}
)
for t in remote_tools
]
print(f"Discovered {len(self._tools)} tools", file=sys.stderr)
for tool in self._tools:
print(f" - {tool.name}", file=sys.stderr)
return self._tools
@property
def tools(self) -> List[ToolInfo]:
return self._tools
def call_tool(self, name: str, arguments: dict) -> Any:
"""Call a tool on the remote MCP server."""
if not self._mcp_client:
raise RuntimeError("Not connected. Call connect() first.")
with ThreadPoolExecutor() as executor:
return executor.submit(self._mcp_client.call_tool, name, arguments or {}).result()
