databricks-dbldatagen-mcp 0.1.0

Databricks MCP Server for Synthetic Data Generation with dbldatagen

Databricks dbldatagen MCP Server

A Model Context Protocol (MCP) server for generating synthetic test data using dbldatagen on Databricks. Enables AI assistants to analyze source tables, generate realistic synthetic data, run SQL queries, and manage notebooks — all through natural language.

Features

• Schema Analysis — Inspect column types, nullable flags, metadata, and detect primary keys
• Data Profiling — Deep profiling including distributions, cardinality, null ratios, and pattern detection
• Synthetic Data Generation — Content-aware generation using dbldatagen DataAnalyzer with preserved columns, fixed values, and schema casting
• SQL Execution — Run any SQL query on Databricks (SELECT, DESCRIBE, CREATE, etc.)
• Notebook Operations — Import, sync, and export notebooks to/from Databricks workspace
• Windows Support — Full Windows compatibility with optimized async handling

Architecture

┌────────────────────────────────────────────────────────────┐
│                   AI Assistant (VS Code)                   │
└───────────────────────────┬────────────────────────────────┘
                            │ MCP Protocol (stdio)
                            ▼
┌────────────────────────────────────────────────────────────┐
│            databricks-dbldatagen-mcp (FastMCP)             │
│                                                            │
│  tools/generate_data.py ────┐                              │
│  tools/analyze_schema.py ───┤                              │
│  tools/profile.py ──────────┼──► @mcp.tool decorators      │
│  tools/sql.py ──────────────┤                              │
│  tools/notebook_ops.py ─────┘                              │
│                                                            │
│  core/analyzer.py ──────────── DataProfiler                │
│  auth.py ───────────────────── Authentication & caching    │
│  identity.py ───────────────── User-agent tagging          │
│  app.py ────────────────────── FastMCP instance + patches  │
│  server.py ─────────────────── Entry point                 │
└───────────────────────────┬────────────────────────────────┘
                            │ Databricks SDK / REST API
                            ▼
                  ┌───────────────────┐
                  │    Databricks     │
                  │    Workspace      │
                  └───────────────────┘

Quick Start

1. Install

git clone <repo-url>
cd databricks-dbldatagen-mcp

# Create venv and install
python -m venv .venv
.venv\Scripts\activate    # Windows
pip install -e .

2. Configure authentication

Create a .env file in the project root:

DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
DATABRICKS_TOKEN=your-token
DATABRICKS_CLUSTER_ID=your-cluster-id

3. Add to VS Code

Create .vscode/mcp.json:

{
  "servers": {
    "databricks-dbldatagen": {
      "type": "stdio",
      "command": "${workspaceFolder}\\.venv\\Scripts\\python.exe",
      "args": ["-m", "databricks_dbldatagen_mcp.server"],
      "env": {
        "PYTHONPATH": "${workspaceFolder}",
        "FASTMCP_LOG_LEVEL": "WARNING"
      }
    }
  }
}

4. Use

Open VS Code Copilot Chat and start prompting:

• "Analyze the schema of catalog.schema.my_table"
• "Generate 5000 rows of synthetic data from source_table into target_table"
• "Run SELECT * FROM my_table LIMIT 10"

Available Tools

Tool	Description
generate_data	Generate synthetic data matching source table structure and distributions
analyze_schema	Get table schema — columns, types, nullable flags, primary keys, row count
profile_table_data	Deep profile — distributions, cardinality, null ratios, patterns per column
sql	Execute any SQL query on Databricks
run_job	Run a Databricks notebook as a one-time job and return output
import_notebook	Download a notebook from Databricks to local
sync_notebook	Push local notebook changes back to Databricks
export_notebook	Upload a local notebook to Databricks workspace

See TOOLS.md for detailed usage with all parameters and examples.

Environment Variables

Variable	Description	Required
DATABRICKS_HOST	Workspace URL	Yes
DATABRICKS_TOKEN	Personal access token	Yes*
DATABRICKS_CLUSTER_ID	Cluster ID for generation/SQL	Yes
DATABRICKS_CONFIG_PROFILE	Profile from ~/.databrickscfg	Alternative
DATABRICKS_CLIENT_ID	OAuth client ID	For OAuth
DATABRICKS_CLIENT_SECRET	OAuth client secret	For OAuth

* Required unless using OAuth or config profile.

Development

pip install -e ".[dev]"
pytest tests/ -v
black databricks_dbldatagen_mcp/
ruff check databricks_dbldatagen_mcp/ --fix

License

MIT License — see LICENSE for details.