"""MCP Server for Databricks Proxy."""
import os
import sys
from typing import Optional
from mcp.server.fastmcp import FastMCP
from .auth import start_oauth_flow, DEFAULT_SCOPES
from .client import DatabricksMCPProxy
class AppState:
"""Global application state."""
proxy: Optional[DatabricksMCPProxy] = None
host: Optional[str] = None
app_url: Optional[str] = None
scopes: str = DEFAULT_SCOPES
authenticated: bool = False
state = AppState()
# Create MCP server
mcp = FastMCP("databricks-mcp-proxy")
@mcp.tool()
def authenticate() -> str:
"""
Authenticate with Databricks using OAuth U2M flow.
Opens a browser for authorization.
Uses DATABRICKS_HOST and DATABRICKS_APP_URL from app.yaml or environment.
"""
try:
host = state.host or os.environ.get("DATABRICKS_HOST")
app_url = state.app_url or os.environ.get("DATABRICKS_APP_URL")
scopes = state.scopes or os.environ.get("DATABRICKS_SCOPES", DEFAULT_SCOPES)
if not host:
return "Error: DATABRICKS_HOST not configured. Set it in app.yaml or environment."
if not app_url:
return "Error: DATABRICKS_APP_URL not configured. Set it in app.yaml or environment."
print(f"Starting OAuth flow for {host}...", file=sys.stderr)
access_token = start_oauth_flow(host, scopes)
state.proxy = DatabricksMCPProxy(host, app_url, access_token)
state.proxy.connect()
state.proxy.discover_tools()
state.authenticated = True
tool_names = [t.name for t in state.proxy.tools]
return f"Authenticated successfully!\n\nAvailable tools ({len(tool_names)}):\n" + "\n".join(f" - {name}" for name in tool_names)
except Exception as e:
state.authenticated = False
return f"Authentication failed: {e}"
@mcp.tool()
def list_databricks_tools() -> str:
"""
List all available tools on the remote Databricks MCP server.
Must authenticate first.
"""
if not state.authenticated or not state.proxy:
return "Not authenticated. Call 'authenticate' first."
if not state.proxy.tools:
return "No tools available."
lines = [f"Available tools ({len(state.proxy.tools)}):\n"]
for tool in state.proxy.tools:
lines.append(f"**{tool.name}**")
lines.append(f" {tool.description}")
if tool.input_schema.get("properties"):
lines.append(f" Parameters: {list(tool.input_schema['properties'].keys())}")
lines.append("")
return "\n".join(lines)
@mcp.tool()
def call_databricks_tool(tool_name: str, arguments: dict = {}) -> str:
"""
Call a tool on the remote Databricks MCP server.
Args:
tool_name: Name of the tool to call (use list_databricks_tools to see available tools)
arguments: Arguments to pass to the tool
"""
if not state.authenticated or not state.proxy:
return "Not authenticated. Call 'authenticate' first."
try:
result = state.proxy.call_tool(tool_name, arguments)
if hasattr(result, 'content') and result.content:
texts = [c.text for c in result.content if hasattr(c, 'text')]
if texts:
return "\n".join(texts)
return str(result)
except Exception as e:
return f"Error calling tool '{tool_name}': {e}"
