Skip to content

databricks-solutions/databricks-exec-code-mcp

Repository files navigation

Databricks MCP Code Execution Template

This template enables AI-assisted development in Databricks by leveraging the Databricks Command Execution API through an MCP server. Test code directly on clusters, then deploy with Databricks Asset Bundles (DABs).

🎯 What This Does

This template enables AI assistants to:

  • ✅ Run and test code directly on Databricks clusters via the Command Execution API
  • ✅ Auto-select clusters - no need to specify a cluster ID, automatically connects to a running cluster
  • ✅ Create Databricks Asset Bundles (DABs) projects
  • ✅ Deploy pipelines to multiple environments (dev/staging/prod)
  • ✅ All from natural language prompts!

Just describe what you want → AI builds, tests the code on Databricks, and deploys the complete pipeline.

💡 Smart Cluster Selection: If no cluster_id is provided, the MCP server automatically finds a running cluster in your workspace (prioritizing clusters with "interactive", "general", or "all-purpose" in the name).


⚙️ Prerequisites - Set Up Environment Variables

Before running the MCP server, configure your Databricks credentials as environment variables:

# Set your Databricks workspace URL (no trailing slash)
export DATABRICKS_HOST=https://your-workspace.cloud.databricks.com

# Set your Databricks Personal Access Token
export DATABRICKS_TOKEN=dapi_your_token_here

To get your Personal Access Token (PAT):

  1. Go to your Databricks workspace
  2. Click your profile icon → SettingsDeveloper
  3. Click Manage under Access Tokens → Generate new token
  4. Copy the token (it won't be shown again!)

💡 Tip: Add these exports to your ~/.zshrc or ~/.bashrc to persist them across terminal sessions.


🚀 Quick Start

1. Clone This Repository

cd databricks-exec-code-mcp

2. Install Requirements

python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
pip install -r requirements.txt

3. Configure Your AI Client

Add the MCP server to your configuration. Use the path to your cloned repo:

For Cursor (.cursor/mcp.json in your project, or global ~/.cursor/mcp.json):

{
  "mcpServers": {
    "databricks": {
      "command": "/path/to/databricks-exec-code-mcp/.venv/bin/python",
      "args": ["/path/to/databricks-exec-code-mcp/mcp_tools/tools.py"]
    }
  }
}

For Claude Code

cd /path/to/your/project

# Register the MCP server
claude mcp add-json databricks '{"command":"/path/to/.venv/bin/python","args":["/path/to/mcp_tools/tools.py"]}'    

💡 Tip: Replace /path/to/databricks-exec-code-mcp with the actual path to your cloned repo.

Restart your AI client after configuration.

4. Install Skills

Install the Databricks skills to teach your AI assistant how to work with Databricks:

# Install for both Cursor and Claude Code (from cloned repo)
./install_skills.sh --all

# Or install for a specific client
./install_skills.sh --cursor
./install_skills.sh --claude 

This installs skills to:

  • Cursor: .cursor/rules/
  • Claude Code: .claude/skills/

5. Start Building!

Just describe what you want in natural language:

Data Engineering:

"Build a Data Engineering pipeline using Medallion Architecture on the NYC Taxi dataset and deploy it with DABs"

Machine Learning:

"Train a classification model on the Titanic dataset, register it to Unity Catalog, and deploy as a DAB job"

Quick Test:

"Run a SQL query to show the top 10 tables in my catalog"


📁 What Gets Generated

The AI will create a complete DABs project:

your-project/
├── databricks.yml              # DABs configuration
├── resources/
│   └── training_job.yml        # Databricks job definition
├── src/<project>/
│   └── notebooks/
│       ├── 01_data_prep.py
│       ├── 02_training.py
│       └── 03_validation.py
└── tests/                      # Unit tests (optional)

🌟 Features

Feature Description
Direct Cluster Execution Test code on Databricks clusters via Databricks Execution API
DABs Packaging Production-ready bundle deployment
Multi-Environment Support for dev/staging/prod targets
Unity Catalog Models and data registered to UC for governance
MLflow Tracking Experiment tracking and model versioning

📚 Resources


📜 License

© 2025 Databricks, Inc. All rights reserved. The source in this project is provided subject to the Databricks License.

Third-Party Licenses

Package License Copyright
FastAPI MIT License Copyright (c) 2018 Sebastián Ramírez
requests Apache License 2.0 Copyright 2019 Kenneth Reitz
uvicorn BSD 3-Clause License Copyright © 2017-present, Encode OSS Ltd
python-dotenv BSD 3-Clause License Copyright (c) 2014, Saurabh Kumar

About

No description, website, or topics provided.

Resources

License

Security policy

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 3

  •  
  •  
  •