Add Databricks migration source

.env.example

@@ -18,12 +18,13 @@ ANTHROPIC_API_KEY=sk-ant-your-key-here
 
 
 # ── Agent Provider / Model ────────────────────────────────────
-# The three pre-built migration agents (Postgres, Snowflake, ClickHouse OSS)
-# each default to whichever API key is set above. Override globally or
-# per-source. Use LibreChat provider casing exactly:
+# The five pre-built migration agents (Postgres, Snowflake, BigQuery,
+# ClickHouse OSS, Databricks) each default to whichever API key is set
+# above. Override globally or per-source. Use LibreChat provider casing
+# exactly:
 #   anthropic | openAI | google | bedrock | azureOpenAI
 #
-# Global default (applies to all three agents unless overridden):
+# Global default (applies to all five agents unless overridden):
 # AGENT_PROVIDER=anthropic
 # AGENT_MODEL=claude-sonnet-4-6
 #
@@ -36,6 +37,8 @@ ANTHROPIC_API_KEY=sk-ant-your-key-here
 # AGENT_MODEL_BIGQUERY=claude-sonnet-4-6
 # AGENT_PROVIDER_CLICKHOUSE_OSS=anthropic
 # AGENT_MODEL_CLICKHOUSE_OSS=claude-haiku-4-5-20251001
+# AGENT_PROVIDER_DATABRICKS=anthropic
+# AGENT_MODEL_DATABRICKS=claude-sonnet-4-6
 #
 # Changing any of these after first start requires `make reset-agent`.
 # (Partners can also edit each agent's model directly in the LibreChat UI.)
@@ -202,3 +205,22 @@ DATASET_SIZE=medium
 # STAGING_GCS_KEY_FILE=./secrets/gcp-key.json  # optional — defaults to BIGQUERY_KEY_FILE
 # STAGING_GCS_ACCESS_KEY_ID=<hmac-access-id>
 # STAGING_GCS_SECRET_ACCESS_KEY=<hmac-secret>
+
+# ── Databricks Source (only for the databricks migration source) ──
+# Requires a Databricks workspace with Unity Catalog and a SQL warehouse.
+# Provision everything with `make databricks-provision` (existing
+# workspace) or `make databricks-provision-workspace` (creates a new
+# serverless workspace), or set these by hand and run
+# `make databricks-setup` for the workload only.
+# See sources/databricks/GUIDE.md.
+#
+# DATABRICKS_HOST=https://dbc-xxxxxxxx-xxxx.cloud.databricks.com
+# DATABRICKS_HTTP_PATH=/sql/1.0/warehouses/xxxxxxxxxxxxxxxx
+# DATABRICKS_TOKEN=dapi................................
+# DATABRICKS_NAMESPACE=migration_demo.tpch     # <catalog>.<schema>
+#
+# S3 staging for large Databricks → ClickHouse Cloud migrations reuses the
+# STAGING_S3_* variables above. On Databricks it additionally needs a Unity
+# Catalog external location over the bucket with WRITE FILES granted —
+# `make databricks-provision` creates one when enable_s3_staging=true.
+# Without it, migrations still work via the direct batch path.

.gitignore

@@ -22,6 +22,9 @@ venv/
 **/.terraform/
 **/.terraform.lock.hcl
 
+# Generated by `make databricks-provision-workspace` — contains an OAuth secret
+sources/databricks/terraform/demo/workspace.auto.tfvars.json
+
 # Service-account JSON keys (BigQuery, GCS staging). `secrets/` is
 # created on `make setup` with a placeholder gcp-key.json. The real key
 # replaces that file but must never be committed.

Makefile

@@ -1,4 +1,4 @@
-.PHONY: setup up up-snowflake up-bigquery down reset reset-agent health logs pull diagram snowflake-setup snowflake-provision bigquery-provision tpch-data tpch-load-bigquery tpch-load-postgres tpch-load-clickhouse-oss migration-status
+.PHONY: setup up up-snowflake up-bigquery up-databricks down reset reset-agent health logs pull diagram snowflake-setup snowflake-provision bigquery-provision databricks-setup databricks-provision databricks-provision-workspace tpch-data tpch-load-bigquery tpch-load-postgres tpch-load-clickhouse-oss migration-status
 
 setup:
 	@echo "Setting up MigrationRoom..."
@@ -15,7 +15,8 @@ setup:
 	@mkdir -p secrets && [ -f secrets/gcp-key.json ] || echo '{}' > secrets/gcp-key.json
 	@# Seed a runtime librechat.yaml so `docker compose up` works even
 	@# before the user has run one of the up* targets. Defaults to empty
-	@# profiles, which strips snowflake-source and bigquery-source.
+	@# profiles, which strips snowflake-source, bigquery-source, and
+	@# databricks-mcp.
 	@COMPOSE_PROFILES="" bash scripts/build-librechat-runtime.sh
 	@echo "✅ Setup complete. Run: make up"
 
@@ -87,6 +88,57 @@ bigquery-provision:
 	@echo ""
 	@echo "Capture the .env block with: cd sources/bigquery/terraform && terraform output -raw env_block"
 
+up-databricks: export COMPOSE_PROFILES := databricks
+up-databricks:
+	@echo "Regenerating librechat.runtime.yaml for active profiles: databricks"
+	@bash scripts/build-librechat-runtime.sh
+	@echo "Pulling images..."
+	docker compose pull
+	@echo "Building custom containers..."
+	docker compose build
+	@echo "Starting services (including databricks-mcp)..."
+	docker compose up -d
+	@echo ""
+	@echo "Container status:"
+	@docker compose ps --format "table {{.Name}}\t{{.Status}}\t{{.Ports}}"
+	@echo ""
+	@echo "If databricks-mcp shows unhealthy, check: docker compose logs databricks-mcp"
+	@echo "(DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN in .env must be set.)"
+
+databricks-setup:
+	@echo "Installing setup dependencies (databricks-sql-connector)…"
+	@python3 -m pip install --quiet -r sources/databricks/scripts/requirements.txt
+	@echo "Setting up migration_demo.tpch workload in Databricks…"
+	@set -a; [ -f .env ] && . ./.env; set +a; \
+	  python3 sources/databricks/scripts/setup_workload.py
+
+databricks-provision:
+	@echo "Provisioning the Databricks demo objects with Terraform…"
+	cd sources/databricks/terraform/demo && terraform init && terraform apply
+	@echo ""
+	@echo "Capture the .env block with:"
+	@echo "  cd sources/databricks/terraform/demo && terraform output -raw env_block"
+
+databricks-provision-workspace:
+	@echo "Phase 1/2 — creating a serverless Databricks workspace…"
+	cd sources/databricks/terraform/workspace && terraform init && terraform apply
+	@echo ""
+	@echo "Phase 2/2 — provisioning the demo objects into the new workspace…"
+	@# Hand the new workspace URL and the account SP's OAuth credentials to
+	@# the demo module. Terraform auto-loads *.auto.tfvars.json, so no
+	@# copy-paste step. The merge logic lives in a script, not an inline
+	@# heredoc: a heredoc spanning multiple Makefile recipe lines only
+	@# works under GNU Make's `.ONESHELL` (added in 3.82) — plain `make`
+	@# on macOS is still 3.81, which runs each recipe line in its own
+	@# shell and silently breaks a multi-line heredoc.
+	cd sources/databricks/terraform/workspace && terraform output -json > /tmp/mr-dbx-workspace-out.json
+	python3 sources/databricks/scripts/merge_workspace_tfvars.py /tmp/mr-dbx-workspace-out.json
+	cd sources/databricks/terraform/demo && terraform init && terraform apply
+	@rm -f /tmp/mr-dbx-workspace-out.json
+	@echo ""
+	@echo "Capture the .env block with:"
+	@echo "  cd sources/databricks/terraform/demo && terraform output -raw env_block"
+
 # Shared TPC-H workload. BigQuery is the first loader; future sources
 # get sibling targets (tpch-load-postgres, tpch-load-clickhouse-oss).
 # The Snowflake source keeps `snowflake-setup` — different mechanics
@@ -133,7 +185,7 @@ tpch-load-clickhouse-oss: tpch-data
 	  python3 workloads/tpch/clickhouse-oss/load.py
 
 down:
-	docker compose --profile snowflake --profile bigquery down
+	docker compose --profile snowflake --profile bigquery --profile databricks down
 
 reset:
 	@bash scripts/reset.sh

README.md

@@ -7,8 +7,8 @@ migration end-to-end in under an hour.**
 
 MigrationRoom is a self-contained Docker Compose playground that
 turns the messy reality of database migration into a six-click
-workflow. Pick a source (PostgreSQL, Snowflake, BigQuery, or
-ClickHouse OSS), click each step on the dashboard, and watch an LLM
+workflow. Pick a source (PostgreSQL, Snowflake, BigQuery, Databricks,
+or ClickHouse OSS), click each step on the dashboard, and watch an LLM
 agent with live MCP connections do the work: introspect the source,
 design a ClickHouse target schema, move the data, validate row
 counts, rewrite analytical queries, and benchmark source vs target —
@@ -35,7 +35,7 @@ all on a real ClickHouse Cloud service you control.
   not its training-data approximation of them. The skills are pulled
   in as a git submodule so updates ship with `git submodule update
   --remote`.
-- **Source-agnostic in the same shape.** Four sources, one dashboard,
+- **Source-agnostic in the same shape.** Five sources, one dashboard,
   one set of six step buttons. The agent and source MCP swap behind
   the scenes when you change the source dropdown — no setup juggling.
 - **MCP-native.** Every database connection is exposed through an MCP
@@ -58,6 +58,7 @@ all on a real ClickHouse Cloud service you control.
 | **ClickHouse OSS → ClickHouse Cloud** | Web analytics platform + TPC-H option | [sources/clickhouse-oss/GUIDE.md](sources/clickhouse-oss/GUIDE.md) |
 | **Snowflake → ClickHouse Cloud** | TPC-H + Snowflake-specific augmentations (VARIANT, TIMESTAMP_TZ, Stream, Dynamic Table, Clustering Key) | [sources/snowflake/GUIDE.md](sources/snowflake/GUIDE.md) |
 | **BigQuery → ClickHouse Cloud** | TPC-H + BigQuery-specific augmentations (STRUCT, ARRAY<STRUCT>, partitioned + clustered tables, materialized view) | [sources/bigquery/GUIDE.md](sources/bigquery/GUIDE.md) |
+| **Databricks → ClickHouse Cloud** | TPC-H + Databricks-specific augmentations (VARIANT, ARRAY&lt;STRUCT&gt;, MAP, generated column, liquid clustering, deletion vectors, materialized view) | [sources/databricks/GUIDE.md](sources/databricks/GUIDE.md) |
 
 ## Prerequisites
 
@@ -99,10 +100,10 @@ self-signed certificate warning). Sign in with
 > total. PostgreSQL: `docker compose logs postgres -f`. ClickHouse
 > OSS: `docker compose logs clickhouse-oss -f`.
 
-For Snowflake or BigQuery, run `make up-snowflake` or `make
-up-bigquery` instead — both source MCPs are profile-gated because
-they need account credentials in `.env` (see the per-source GUIDE for
-the setup walkthrough).
+For Snowflake, BigQuery, or Databricks, run `make up-snowflake`, `make
+up-bigquery`, or `make up-databricks` instead — all three source MCPs
+are profile-gated because they need account credentials in `.env` (see
+the per-source GUIDE for the setup walkthrough).
 
 ## Using the MigrationRoom dashboard
 
@@ -120,16 +121,17 @@ Three controls that determine which agent runs and what data it
 operates on. Change them before clicking any step button.
 
 - **Source dropdown** — `Postgres` / `Snowflake` / `BigQuery` /
-  `ClickHouse OSS`. Switching the source **auto-switches the
-  LibreChat agent** in the right pane to the matching pre-built agent
-  (e.g. picking BigQuery selects the `BigQuery → ClickHouse Cloud`
-  agent with its MCPs and system prompt). No agent toggling needed in
-  LibreChat.
+  `Databricks` / `ClickHouse OSS`. Switching the source
+  **auto-switches the LibreChat agent** in the right pane to the
+  matching pre-built agent (e.g. picking BigQuery selects the
+  `BigQuery → ClickHouse Cloud` agent with its MCPs and system
+  prompt). No agent toggling needed in LibreChat.
 - **Source database dropdown** — the actual database / dataset /
   schema name on the source. For bundled workloads: `ecommerce`
   (Postgres) or `analytics` (ClickHouse OSS) or `migration_demo`
-  (Snowflake / BigQuery). For TPC-H loaded via `make tpch-load-*`:
-  `tpch`.
+  (Snowflake / BigQuery / Databricks, e.g. `migration_demo.tpch` for
+  Databricks' `catalog.schema` pair). For TPC-H loaded via `make
+  tpch-load-*`: `tpch`.
 - **Edit · N OLAP button** — opens an editor for the analytical
   queries that drive step 1's `ORDER BY` design, step 4's query
   rewrite, and step 5's benchmark. `N` is how many queries are
@@ -209,6 +211,7 @@ reads `KPI DASHBOARD · LIVE · STEP N — <step name>`.
   [Postgres](sources/postgres/GUIDE.md) ·
   [Snowflake](sources/snowflake/GUIDE.md) ·
   [BigQuery](sources/bigquery/GUIDE.md) ·
+  [Databricks](sources/databricks/GUIDE.md) ·
   [ClickHouse OSS](sources/clickhouse-oss/GUIDE.md).
 - **Add your own source database** —
   [docs/adding-a-source.md](docs/adding-a-source.md) explains the
@@ -226,9 +229,10 @@ make setup               # first-time setup (submodules + agent skills + .env)
 make up                  # start the playground (Postgres + ClickHouse OSS sources)
 make up-snowflake        # also start the Snowflake source MCP (needs SNOWFLAKE_* in .env)
 make up-bigquery         # also start the BigQuery source MCP (needs BIGQUERY_* in .env)
+make up-databricks       # also start the Databricks source MCP (needs DATABRICKS_* in .env)
 make down                # stop without removing data
 make reset               # destroy volumes and start fresh
-make reset-agent         # delete + recreate the four pre-built agents (after model/provider changes)
+make reset-agent         # delete + recreate the five pre-built agents (after model/provider changes)
 ```
 
 The full command list lives in the [Makefile](Makefile); see the

docker-compose.yml

@@ -229,6 +229,40 @@ services:
     networks:
       - playground-net
 
+  # ── Databricks Source MCP ───────────────────────────────────
+  # Purpose-built read-only MCP (docker/databricks-mcp) over a Databricks
+  # SQL warehouse. Databricks publishes no introspect-and-SELECT MCP for
+  # warehouses — their `databricks-mcp` package is an OAuth helper for the
+  # hosted UC-functions / vector-search / Genie servers — so this is ours,
+  # in the same spirit as clickhousectl-mcp. Because we author the tool
+  # schemas, no Gemini shim is needed (unlike snowflake-source).
+  #
+  # Profile-gated: partners using Databricks set DATABRICKS_* in .env then
+  # run `make up-databricks` (or `docker compose --profile databricks up -d`).
+  databricks-mcp:
+    profiles: ["databricks"]
+    build: ./docker/databricks-mcp
+    environment:
+      DATABRICKS_HOST: ${DATABRICKS_HOST:-}
+      DATABRICKS_HTTP_PATH: ${DATABRICKS_HTTP_PATH:-}
+      DATABRICKS_TOKEN: ${DATABRICKS_TOKEN:-}
+      DATABRICKS_NAMESPACE: ${DATABRICKS_NAMESPACE:-}
+    ports:
+      - "8008:8000"
+    # Probe /sse with Python stdlib — the image is python:3.12-slim and has
+    # no curl/wget. The connection to Databricks is lazy (opened per tool
+    # call), so this passes even with invalid credentials; that's
+    # deliberate, so a credential typo surfaces as a clear tool error in
+    # chat rather than a container that won't start.
+    healthcheck:
+      test: ["CMD", "python", "-c", "import urllib.request, sys; sys.exit(0 if urllib.request.urlopen('http://localhost:8000/sse', timeout=2).status == 200 else 1)"]
+      interval: 5s
+      timeout: 5s
+      retries: 20
+      start_period: 10s
+    networks:
+      - playground-net
+
   # ── Migration Runner ────────────────────────────────────────
   # Executes Python migration scripts in-chat (via run_python MCP tool)
   # so the agent can drive end-to-end migrations without asking the
@@ -329,8 +363,10 @@ services:
     # `docker compose ps` instead of a silent runtime bug.
     #
     # snowflake-source / snowflake-source-shim are profile-gated to
-    # `snowflake`; Compose silently skips depends_on entries for services
-    # excluded by the active profile set, so leaving them here is safe.
+    # `snowflake`, bigquery-source to `bigquery`, and databricks-mcp to
+    # `databricks`; Compose silently skips depends_on entries for
+    # services excluded by the active profile set, so leaving them
+    # here is safe.
     depends_on:
       mongodb:
         condition: service_healthy
@@ -353,6 +389,9 @@ services:
         # check we can express.
         condition: service_started
         required: false
+      databricks-mcp:
+        condition: service_healthy
+        required: false
     healthcheck:
       test: ["CMD-SHELL", "wget -qO- http://127.0.0.1:3080/health || exit 1"]
       interval: 10s
@@ -377,9 +416,9 @@ services:
       MONGO_URI: "mongodb://mongodb:27017/LibreChat"
       LIBRECHAT_URL: "http://librechat:3080"
       # Forwarded from the host's $COMPOSE_PROFILES so the agent loop knows
-      # which optional sources (snowflake, bigquery) the partner actually
-      # turned on. Agents for inactive sources are skipped and any stale
-      # row from a previous profile is deleted.
+      # which optional sources (snowflake, bigquery, databricks) the partner
+      # actually turned on. Agents for inactive sources are skipped and any
+      # stale row from a previous profile is deleted.
       ACTIVE_PROFILES: ${COMPOSE_PROFILES:-}
     entrypoint: ["/bin/bash", "-c"]
     command:
@@ -436,6 +475,7 @@ services:
           "postgres|Postgres → ClickHouse Cloud|POSTGRES"
           "snowflake|Snowflake → ClickHouse Cloud|SNOWFLAKE"
           "bigquery|BigQuery → ClickHouse Cloud|BIGQUERY"
+          "databricks|Databricks → ClickHouse Cloud|DATABRICKS"
           "clickhouse-oss|ClickHouse OSS → ClickHouse Cloud|CLICKHOUSE_OSS"
         )
         SHARED_MCPS=(clickhousectl clickhouse-docs migration-runner)
@@ -447,6 +487,7 @@ services:
         declare -A OPTIONAL_PROFILE=(
           [snowflake]=snowflake
           [bigquery]=bigquery
+          [databricks]=databricks
         )
         NORM_PROFILES=",$$(echo "$${ACTIVE_PROFILES:-}" | tr ' ' ',' | tr -s ','),"
 

docker/databricks-mcp/Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+COPY requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r /app/requirements.txt
+
+COPY sql_guard.py server.py /app/
+
+EXPOSE 8000
+CMD ["python", "server.py"]

docker/databricks-mcp/requirements.txt

@@ -0,0 +1,3 @@
+mcp[cli]>=1.0
+databricks-sql-connector>=4.0
+sqlglot>=30.0