Author(s): Vinay Ram Gazula

Originally published on Towards AI.

In the rapidly evolving world of Agentic AI, the biggest bottleneck isn’t the model’s intelligence — it’s the “glue” code required to connect that model to your data.

Every time I wanted to build an agent that could query my Databricks Lakehouse, I found myself rewriting the same tool definitions, schema fetchers, and SQL execution wrappers. I was hardcoding specific tools into my LangGraph agents, making them brittle and hard to scale.

That changed when I implemented the Model Context Protocol (MCP).

In this post, I’ll walk you through exactly how I built a custom Databricks MCP Server. This server acts as a standardized gateway, allowing any MCP-compliant agent (like Claude Desktop or a custom LangGraph bot) to discover Unity Catalog metadata, inspect table lineage, and execute Spark SQL queries — all without hardcoding a single tool into the agent itself.

What is the Model Context Protocol (MCP)?

Before we dive into the Python code, let’s clarify what we are building.

MCP is an open standard that decouples AI Models from Tools. Instead of manually defining function schemas (like OpenAI function calling JSONs) inside your chatbot code, you build a lightweight “server” that exposes these capabilities. The agent simply connects to the server and “learns” what tools are available.

Think of it as a USB-C port for AI tools: standardizing how models connect to data and execution environments.

The architecture: An AI Agent connects via MCP Protocol to our custom Databricks Server, which proxies requests to the Unity Catalog and Serverless Compute.

The Setup (Python & FastMCP)

I chose Python for this implementation because the Databricks SDK for Python is mature and feature-rich. To handle the MCP protocol details without getting bogged down in low-level transport logic, I used the mcp library and its FastMCP class.

Project Structure:

databricks-mcp-server/
├── databricks_mcp_server/
│ ├── tools.py # The core tool logic
│ ├── unitycatalog.py # Wrapper for Metadata API
│ ├── warehouse.py # Wrapper for SQL Execution API
│ ├── lineage.py # Table Lineage Logic
│ └── utils.py # Markdown formatting helpers
├── server.py # Entry point
└── requirements.txt

In server.py, setting up the server is surprisingly simple. I configured it to run over stdio (for local usage) or SSE (Server-Sent Events) for remote deployment.

from mcp.server.fastmcp import FastMCP
from databricks_mcp_server.tools import (
 fetch_schemas_in_catalog,
)

# Initialize FastMCP
mcp = FastMCP(
 name="databricks-mcp-server",
 log_level="INFO",
 json_response=True,
 host="0.0.0.0",
 port=8000,
)

# Register Tools
mcp.add_tool(fetch_schemas_in_catalog)

if __name__ == "__main__":
 mcp.run(transport="streamable-http")

The “LLM Translation” Layer

One of the most critical aspects of this project was formatting. API responses from Unity Catalog are massive JSON objects containing IDs, timestamps, and other noise. Feeding raw JSON to an LLM wastes tokens and confuses the model.

I implemented a helper in utils.py to convert these responses into clean Markdown.

# A helper to format table metadata into LLM-friendly Markdown
def format_table_info(table_info: List[TableInfo], lineage_info: List[LineageInfo], extended: bool = False) -> str:
 markdown_output = ""
 for table in table_info:
 markdown_output += f"### Table: {table.full_name}\n"
 markdown_output += f"- **Type**: {table.table_type}\n"
 markdown_output += f"- **Comment**: {table.comment}\n"
 
 # Format columns with types and nullability
 markdown_output += "| Column | Type | Nullable | Comment |\n"
 markdown_output += "| --- | --- | --- | --- |\n"
 for col in table.columns:
 comment = col.comment if col.comment else ""
 markdown_output += f"| {col.name} | {col.type_text} | {col.type_boolean} | {comment} |\n"
 
 if extended and lineage_info:
 # ... lineage formatting logic ...
 return markdown_output

Why Markdown? LLMs are trained heavily on Markdown text. By converting schema definitions into structured text with headers and bullet points, we dramatically improve the model’s ability to “understand” the table structure.

Implementing the Tools

The server currently exposes four primary tools designed to guide an agent through a logical “Discovery → Drill-Down → Inspection → Execution” workflow. This hierarchical approach prevents the agent from getting overwhelmed by too much context at once.

1. High-Level Discovery: fetch_schemas_in_catalog

This is the entry point. When an agent enters a new environment, it often doesn’t know what data is available.

• Function: fetch_schemas_in_catalog(catalog: str)
• Purpose: Lists all schemas within a specific catalog.
• Agent Use Case: “What schemas are available in the ‘main’ catalog?”

2. Targeted Exploration: fetch_tables_in_schema

Once the agent identifies a relevant schema (e.g., finance), it needs to see the assets inside.

• Function: fetch_tables_in_schema(catalog: str, schema: str)
• Purpose: Returns a list of tables for a specific schema.
• Agent Use Case: “List all tables in the ‘finance’ schema so I can see if we have revenue data.”

3. Deep Inspection & Lineage: fetch_table_info

This tool is the powerhouse. Before writing a query, an agent must understand column names, data types, and relationships. Crucially, I added Lineage support here.

• Function: fetch_table_info(table_names: List[str])
• Purpose: Returns detailed metadata (columns, types, comments) and upstream/downstream lineage dependencies.
• Agent Use Case: “Get me the schema for finance.revenue_report. Also, where does this data come from?"

4. Safe Execution: execute_spark_sql_query

Finally, the agent needs “hands” to retrieve the actual data.

• Function: execute_spark_sql_query(query: str)
• Purpose: Executes a read-only Spark SQL query via the Databricks SQL Statement Execution API.
• Safety: Explicitly restricted to SELECT statements to prevent accidental data modification.

Testing with MCP Inspector

Before deploying, I needed to verify that the tool definitions (JSON schemas) were correct and that the server responded properly. The MCP Inspector tool (by Anthropic) was invaluable here.

It provides a web interface to interact with your MCP server, view the exact JSON payloads sent to the tools, and debug errors in real-time.

Guardrails & Security

Building an agent that can query your Data Lakehouse sounds risky. That’s why Guardrails are non-negotiable.

Least Privilege Access

I designed the SQL tool to be explicitly for retrieval. While the Python code includes docstring warnings, the real enforcement happens at the credential level. The MCP server runs with a Service Principal identity that has strictly scoped permissions.

Leveraging Unity Catalog (UC)

This is the biggest advantage of building on Databricks. I don’t need to implement custom row-level security or access control lists in my Python code.

• Identity: The MCP server operates as a specific Service Principal.
• Governance: If that Service Principal doesn’t have SELECT permission on a table containing PII, the query fails at the Databricks layer. Unity Catalog handles the masking, filtering, and access denial.

By relying on UC, we ensure that the agent inherits the same robust governance as any human user.

Real-World Application: The AgenticLakehouse

This MCP server isn’t just a proof of concept; it is the “hands” of a larger project I call AgenticLakehouse.

AgenticLakehouse is a multi-agent system built with LangGraph where a “Router” agent intelligently delegates tasks. Here is how the agent leverages all four tools in a real conversation:

1. User Query: “Who are our top customers in the retail sector?”
2. Discovery: The agent calls fetch_schemas_in_catalog("main") to see available domains. It spots a retail schema.
3. Drill-Down: The agent calls fetch_tables_in_schema("main", "retail") and identifies tables like customers and orders.
4. Inspection: The agent calls fetch_table_info(["main.retail.customers", "main.retail.orders"]). It learns the column names for joining (e.g., cust_id) and confirms the data types.
5. Execution: Finally, the agent constructs a valid SQL query and calls execute_spark_sql_query to retrieve and rank the customers.

By separating the tools (MCP) from the reasoning (LangGraph), I’ve built a modular system where the agent can dynamically explore the Lakehouse structure rather than relying on hardcoded assumptions.

Read more about the AgenticLakehouse here

Future Roadmap

This is just version 1.0. The MCP standard is flexible, and I plan to expand the server’s capabilities:

• MLflow Integration: Tools to list_experiments and fetch_best_run so agents can assist with MLOps.
• Jobs API: Tools to debug failed workflows by fetching run logs.
• System Prompts: Adding standard system prompts to the MCP server to guide the LLM on how to interpret Databricks-specific nuances.

Resources

• Databricks MCP Server Repo: https://github.com/vinay-ram1999/databricks-mcp-server
• AgenticLakehouse Repo: https://github.com/vinay-ram1999/AgenticLakehouse

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI