knustx/databricks-mcp-server

Databricks MCP Server

A Model Context Protocol (MCP) server that provides seamless integration with Databricks Unity Catalog. This server enables AI assistants to interact with your Databricks workspace, query metadata, sample data, and perform various Unity Catalog operations.

Features

• Unity Catalog Integration: Browse catalogs, schemas, and tables
• Metadata Querying: Get detailed information about tables, columns, and properties
• Data Sampling: Sample data from tables for analysis
• SQL Query Execution: Run SQL queries against your Databricks warehouses
• Table Search: Search for tables by name or metadata
• Data Discovery: Advanced search and filtering capabilities
• Data Quality Insights: Basic data quality analysis
• Lineage Information: Table lineage tracking (when available)

Installation

Prerequisites

• Python 3.8 or higher
• Databricks workspace access
• Databricks personal access token

Install from Source

git clone <repository-url>
cd databricks-mcp-server
pip install -e .

Install Development Dependencies

pip install -e ".[dev]"

Configuration

Environment Variables

Set the following environment variables:

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"
export DATABRICKS_WAREHOUSE_ID="your-warehouse-id"  # Optional but recommended
export LOG_LEVEL="INFO"  # Optional

Configuration File

Alternatively, create a config.json file:

{
  "databricks_host": "https://your-workspace.cloud.databricks.com",
  "databricks_token": "your-personal-access-token", 
  "databricks_warehouse_id": "your-warehouse-id",
  "log_level": "INFO"
}

Usage

Running the Server

# Run directly
python -m databricks_mcp_server.server

# Or use the installed command
databricks-mcp-server

MCP Client Integration

The server implements the Model Context Protocol and can be used with any MCP-compatible client. Here's an example configuration for Claude Desktop:

{
  "mcpServers": {
    "databricks": {
      "command": "databricks-mcp-server",
      "env": {
        "DATABRICKS_HOST": "https://your-workspace.cloud.databricks.com",
        "DATABRICKS_TOKEN": "your-token"
      }
    }
  }
}

Available Tools

Catalog Operations

• list_catalogs: List all Unity Catalog catalogs
• list_schemas: List schemas in a catalog
• list_tables: List tables in a schema

Table Operations

• describe_table: Get detailed table information including columns and metadata
• sample_table: Sample data from a table (configurable limit)
• search_tables: Search for tables by name or metadata

Query Operations

• execute_query: Execute SQL queries against Databricks warehouses
• get_table_lineage: Get lineage information for tables

Resources

The server exposes Databricks resources through URIs:

• databricks://catalog/{catalog_name}: Catalog information
• databricks://catalog/{catalog_name}/{schema_name}: Schema information
• databricks://catalog/{catalog_name}/{schema_name}/{table_name}: Table information

Examples

Basic Usage

from databricks_mcp_server import DatabricksClient

# Initialize client
client = await DatabricksClient.create()

# List catalogs
catalogs = await client.list_catalogs()
print(f"Found {len(catalogs)} catalogs")

# Get table info
table_info = await client.describe_table("main", "default", "my_table")
print(f"Table has {len(table_info.columns)} columns")

# Sample data
sample = await client.sample_table("main", "default", "my_table", limit=5)
print(f"Sampled {sample.row_count} rows")

Advanced Data Discovery

from databricks_mcp_server import UnityCatalogManager

# Initialize manager
manager = UnityCatalogManager(client)

# Discover tables with patterns
results = await manager.discover_data(
    search_patterns=["customer", "user"],
    catalogs=["main", "analytics"],
    include_metadata=True
)

print(f"Found {results.total_tables} matching tables")

Development

Running Tests

pytest

Code Formatting

black src/ tests/
isort src/ tests/

Type Checking

mypy src/

Troubleshooting

Common Issues

1. Authentication Error: Verify your DATABRICKS_TOKEN is valid and has appropriate permissions
2. Connection Error: Check that DATABRICKS_HOST is correct and accessible
3. No Warehouses: Ensure you have at least one SQL warehouse running in your workspace

Debugging

Enable debug logging:

export LOG_LEVEL=DEBUG
databricks-mcp-server

Configuration Validation

Use the built-in validation:

from databricks_mcp_server.utils import validate_databricks_config

validation = validate_databricks_config()
if not validation["valid"]:
    print("Configuration errors:", validation["errors"])

Security Considerations

• Never commit access tokens to version control
• Use environment variables or secure configuration management
• Limit token permissions to minimum required scope
• Consider using service principals for production deployments

Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests
5. Run the test suite
6. Submit a pull request

License

MIT License - see LICENSE file for details.

Support

For issues and questions:

1. Check the troubleshooting section
2. Search existing issues
3. Create a new issue with detailed information

Changelog

v0.1.0

• Initial release
• Basic Unity Catalog integration
• Table metadata and sampling
• SQL query execution
• MCP server implementation