from pathlib import Path
from mcp.server.fastmcp import FastMCP
from fastapi import FastAPI
from fastapi.responses import FileResponse
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksOAuthClientProvider
from mcp import ClientSession
import asyncio
from typing import Any,Dict
import os
#######Workspace client#########
client = WorkspaceClient(
host=os.environ["DATABRICKS_HOST"]
)
############databricks auth###########
#auth=DatabricksOAuthClientProvider(client)
STATIC_DIR = Path(__file__).parent / "static"
# Create an MCP server
mcp = FastMCP("Custom MCP Server on Databricks Apps")
# Add an addition tool
@mcp.tool()
def add(a: int, b: int) -> int:
"""Add two numbers"""
return a + b
@mcp.tool("get_nyc_taxi_sample")
async def get_nyc_taxi_sample() -> dict:
"""
Fetch sample NYC taxi trip data from Databricks table `samples.nyctaxi.trips`.
NOTE:
- This tool takes **no input arguments**.
- It returns 5 sample rows with column names.
"""
# Execute the SQL query
stmt = client.statement_execution.execute_statement(
warehouse_id="56bb8419b9dfb530",
statement="SELECT * FROM samples.nyctaxi.trips LIMIT 5"
)
# Wait for completion and get metadata
status = client.statement_execution.get_statement(stmt.statement_id)
columns = [col.name for col in status.manifest.schema.columns]
rows = []
chunk_index = 0
# Fetch all chunks
while True:
try:
chunk = client.statement_execution.get_statement_result_chunk_n(
statement_id=stmt.statement_id,
chunk_index=chunk_index
)
except Exception as e:
if "out of bounds" in str(e):
break
else:
raise e
rows.extend(chunk.data_array)
chunk_index += 1
return {"columns": columns, "rows": rows}
# Add a dynamic greeting resource
@mcp.resource("greeting://{name}")
def get_greeting(name: str) -> str:
"""Get a personalized greeting"""
return f"Hello, {name}!"
mcp_app = mcp.streamable_http_app()
app = FastAPI(
lifespan=lambda _: mcp.session_manager.run(),
)
@app.get("/", include_in_schema=False)
async def serve_index():
return FileResponse(STATIC_DIR / "index.html")
app.mount("/", mcp_app)
