diff --git a/.env.example b/.env.example index 6f971a9..23eb6ed 100644 --- a/.env.example +++ b/.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= # STAGING_GCS_SECRET_ACCESS_KEY= + +# ── 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 # . +# +# 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. diff --git a/.gitignore b/.gitignore index 485e470..d2701a6 100644 --- a/.gitignore +++ b/.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. diff --git a/Makefile b/Makefile index 63fdb74..967b172 100644 --- a/Makefile +++ b/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 diff --git a/README.md b/README.md index 6bea4e6..011d93f 100644 --- a/README.md +++ b/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, partitioned + clustered tables, materialized view) | [sources/bigquery/GUIDE.md](sources/bigquery/GUIDE.md) | +| **Databricks → ClickHouse Cloud** | TPC-H + Databricks-specific augmentations (VARIANT, ARRAY<STRUCT>, 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 — `. [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 diff --git a/docker-compose.yml b/docker-compose.yml index ec8a087..b8ad991 100644 --- a/docker-compose.yml +++ b/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 ',')," diff --git a/docker/databricks-mcp/Dockerfile b/docker/databricks-mcp/Dockerfile new file mode 100644 index 0000000..743a607 --- /dev/null +++ b/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"] diff --git a/docker/databricks-mcp/requirements.txt b/docker/databricks-mcp/requirements.txt new file mode 100644 index 0000000..5a98f30 --- /dev/null +++ b/docker/databricks-mcp/requirements.txt @@ -0,0 +1,3 @@ +mcp[cli]>=1.0 +databricks-sql-connector>=4.0 +sqlglot>=30.0 diff --git a/docker/databricks-mcp/server.py b/docker/databricks-mcp/server.py new file mode 100644 index 0000000..932b72b --- /dev/null +++ b/docker/databricks-mcp/server.py @@ -0,0 +1,193 @@ +"""MigrationRoom — Databricks source MCP server. + +Read-only introspection and SELECT access to ONE Databricks SQL warehouse, +exposed over MCP/SSE so the migration agent can discover the source schema +without writing Python. + +Why this exists rather than an off-the-shelf package: Databricks' own +`databricks-mcp` PyPI package is an OAuth helper for their *hosted* MCP +servers (Unity Catalog functions, vector search, Genie) — there is no +official introspect-and-SELECT MCP for a SQL warehouse. Authoring the tool +schemas ourselves also keeps them free of the JSON-Schema keywords that +Gemini's function-calling API rejects, so this server needs no shim (unlike +snowflake-source). + +Environment: + DATABRICKS_HOST required — workspace URL or bare hostname + DATABRICKS_HTTP_PATH required — e.g. /sql/1.0/warehouses/abc123 + DATABRICKS_TOKEN required — PAT for a read-only principal + DATABRICKS_NAMESPACE optional — "." default scope + MCP_PORT optional — default 8000 +""" +from __future__ import annotations + +import os +from contextlib import contextmanager +from typing import Any + +from mcp.server.fastmcp import FastMCP + +from sql_guard import SqlNotAllowed, guard + +mcp = FastMCP("databricks-source") + + +def _host() -> str: + """Bare hostname — the connector rejects a scheme or trailing slash.""" + raw = os.environ["DATABRICKS_HOST"].strip() + return raw.removeprefix("https://").removeprefix("http://").rstrip("/") + + +@contextmanager +def _cursor(): + """One short-lived connection + cursor per tool call. + + Deliberately not pooled: LibreChat holds the SSE session open for the + whole conversation, and a warehouse that auto-stops would leave a stale + connection behind. Reconnecting costs ~1 s and is far less confusing + than a silently dead handle. + """ + from databricks import sql as dbsql + + conn = dbsql.connect( + server_hostname=_host(), + http_path=os.environ["DATABRICKS_HTTP_PATH"], + access_token=os.environ["DATABRICKS_TOKEN"], + ) + try: + cur = conn.cursor() + try: + yield cur + finally: + cur.close() + finally: + conn.close() + + +def _rows(cur) -> list[dict[str, Any]]: + columns = [c[0] for c in cur.description or []] + return [dict(zip(columns, row)) for row in cur.fetchall()] + + +def _ident(name: str) -> str: + """Backtick-quote one identifier part, rejecting embedded backticks. + + Identifiers arrive as tool arguments from the model, so they are + untrusted input even though they are not user-facing. + """ + cleaned = (name or "").strip().strip("`") + if not cleaned or "`" in cleaned: + raise ValueError(f"invalid identifier: {name!r}") + return f"`{cleaned}`" + + +@mcp.tool() +def list_catalogs() -> list[dict[str, Any]]: + """List Unity Catalog catalogs visible to this principal.""" + with _cursor() as cur: + cur.execute( + "SELECT catalog_name, comment " + "FROM system.information_schema.catalogs " + "ORDER BY catalog_name" + ) + return _rows(cur) + + +@mcp.tool() +def list_schemas(catalog: str) -> list[dict[str, Any]]: + """List schemas in `catalog`.""" + with _cursor() as cur: + cur.execute( + "SELECT schema_name, comment " + "FROM system.information_schema.schemata " + "WHERE catalog_name = ? " + "ORDER BY schema_name", + [catalog], + ) + return _rows(cur) + + +@mcp.tool() +def list_tables(catalog: str, schema: str) -> list[dict[str, Any]]: + """List tables in `catalog`.`schema` with Delta size metadata. + + Row counts are NOT included: Delta metadata doesn't carry them and a + per-table COUNT(*) would make this call slow. Get them with one + UNION ALL count query via run_select_query instead. + """ + with _cursor() as cur: + cur.execute( + "SELECT table_name, table_type, comment " + "FROM system.information_schema.tables " + "WHERE table_catalog = ? " + " AND table_schema = ? " + "ORDER BY table_name", + [catalog, schema], + ) + tables = _rows(cur) + for row in tables: + row["sizeInBytes"] = None + row["numFiles"] = None + if row.get("table_type") not in (None, "MANAGED", "EXTERNAL"): + continue + fq = f"{_ident(catalog)}.{_ident(schema)}.{_ident(row['table_name'])}" + try: + cur.execute(f"DESCRIBE DETAIL {fq}") + detail = _rows(cur) + except Exception: + # Views and non-Delta tables have no DESCRIBE DETAIL. + continue + if detail: + row["sizeInBytes"] = detail[0].get("sizeInBytes") + row["numFiles"] = detail[0].get("numFiles") + return tables + + +@mcp.tool() +def describe_table(catalog: str, schema: str, table: str) -> dict[str, Any]: + """Full schema plus Delta detail for one table. + + Returns columns, plus clustering/partition columns, table features, + deletion-vector state, and recent history — the source-specific + features the migration has to make decisions about. + """ + fq = f"{_ident(catalog)}.{_ident(schema)}.{_ident(table)}" + out: dict[str, Any] = {"table": f"{catalog}.{schema}.{table}"} + with _cursor() as cur: + cur.execute(f"DESCRIBE TABLE EXTENDED {fq}") + out["describe_extended"] = _rows(cur) + try: + cur.execute(f"DESCRIBE DETAIL {fq}") + out["detail"] = _rows(cur) + except Exception as exc: + out["detail"] = {"unavailable": str(exc)} + try: + cur.execute(f"DESCRIBE HISTORY {fq} LIMIT 5") + out["history"] = _rows(cur) + except Exception as exc: + out["history"] = {"unavailable": str(exc)} + return out + + +@mcp.tool() +def run_select_query(sql: str, max_rows: int = 200) -> list[dict[str, Any]]: + """Run ONE read-only statement (SELECT / WITH / SHOW / DESCRIBE / + EXPLAIN) and return its rows. + + A LIMIT is applied when the statement has none. Mutations and + multi-statement input are refused — run DDL against the target with + the clickhousectl MCP. + """ + try: + statement = guard(sql, max_rows=max_rows) + except SqlNotAllowed as exc: + raise ValueError(str(exc)) from exc + with _cursor() as cur: + cur.execute(statement) + return _rows(cur) + + +if __name__ == "__main__": + mcp.settings.host = "0.0.0.0" + mcp.settings.port = int(os.environ.get("MCP_PORT", "8000")) + mcp.run(transport="sse") diff --git a/docker/databricks-mcp/sql_guard.py b/docker/databricks-mcp/sql_guard.py new file mode 100644 index 0000000..de45a50 --- /dev/null +++ b/docker/databricks-mcp/sql_guard.py @@ -0,0 +1,181 @@ +"""Read-only statement guard for the Databricks source MCP. + +Deliberately free of `mcp` and `databricks` imports so it can be unit-tested +on the host without the server's dependencies installed. + +A Databricks SQL warehouse has no session-level read-only switch, so the +only place we can refuse a mutation is here. Grants are the second layer +(see sources/databricks/GUIDE.md) — this is the first. + +This module used to classify statements by hand-lexing the leading keyword +and scanning for semicolons/comments/quotes with regexes. Three review +rounds each found a real bypass in that approach (a CTE prefixing DML, a +backtick-quoted identifier blinding the CTE scanner, and comment/string +lexing that diverged from Spark's grammar). Hand-matching Spark's lexical +grammar is not a small task, so this rewrite delegates it to sqlglot's +Databricks dialect tokenizer/parser instead of re-implementing it. A +CTE-prefixed INSERT then simply *is* an `exp.Insert` node, and correct +comment/string/identifier handling comes from a real tokenizer. +""" +from __future__ import annotations + +import logging + +import sqlglot +from sqlglot import exp +from sqlglot.dialects import Databricks +from sqlglot.errors import ParseError, TokenError + +# sqlglot logs a WARNING ("... contains unsupported syntax. Falling back to +# parsing as a 'Command'.") for every SHOW/EXPLAIN/OPTIMIZE/VACUUM statement, +# because those are intentionally handled via the exp.Command fallback below. +# That is expected here, not a problem to surface on every legitimate SHOW. +logging.getLogger("sqlglot").setLevel(logging.ERROR) + +# Node types that represent a read-only statement once successfully parsed. +READ_ONLY_NODES = (exp.Select, exp.Union, exp.Subquery, exp.Describe) + +# Leading verbs that are read-only even when sqlglot cannot build a full AST +# for the statement (Databricks-specific extensions like `DESC DETAIL`) or +# falls back to a generic exp.Command node (e.g. `SHOW CATALOGS`). Kept +# narrow on purpose: this is NOT "allow anything unparseable", it is "allow +# only these four verbs when parsing can't tell us more." +READ_ONLY_VERBS = frozenset({"SHOW", "EXPLAIN", "DESCRIBE", "DESC"}) + +_DATABRICKS = Databricks() +_DIALECT = "databricks" + + +class SqlNotAllowed(ValueError): + """Raised when a statement is not a single read-only statement.""" + + +def _strip_trailing_semicolon(sql: str, toks) -> str: + """Return `sql` with its single legal trailing ';' (if any) removed. + + A trailing semicolon is a statement terminator, not statement text; if + we appended a LIMIT clause after it we'd produce invalid SQL. Everything + else about the original text — including internal whitespace and + formatting — is preserved untouched. + """ + if toks and toks[-1].token_type == sqlglot.TokenType.SEMICOLON: + return sql[: toks[-1].start].rstrip() + return sql.strip() + + +def _unwrap_subquery(node: exp.Expression) -> exp.Expression: + """Follow `Subquery.this` down to the innermost wrapped query. + + `(SELECT ...)` parses to `exp.Subquery` wrapping an `exp.Select` (or + `exp.Union`), and a `LIMIT` already present in the source SQL may live + on that inner node rather than on the `Subquery` itself — e.g. + `(SELECT * FROM t LIMIT 5)` puts the limit on the inner `Select`, while + `(SELECT * FROM t) LIMIT 5` puts it on the outer `Subquery`. Checking + only one of the two slots would miss an existing limit and double it. + """ + while isinstance(node, exp.Subquery) and node.this is not None: + node = node.this + return node + + +def _existing_limit(root: exp.Expression) -> exp.Expression | None: + """Return the existing LIMIT clause for `root`, if any, checking both + the node itself and, for a parenthesized query, the query it wraps.""" + limit = root.args.get("limit") + if limit is not None: + return limit + inner = _unwrap_subquery(root) + if inner is not root: + return inner.args.get("limit") + return None + + +def guard(sql: str, max_rows: int = 200) -> str: + """Validate `sql` as one read-only statement and return it ready to run. + + A row cap is appended to SELECT/UNION statements — including ones + wrapped in parentheses, e.g. `(SELECT ...)` — that don't already have + one, so an unbounded scan can't stream a whole fact table into the + chat context. The original statement text is returned (never a + sqlglot-regenerated form) with only the trailing statement terminator + removed and, when applicable, a LIMIT clause appended. + + Raises SqlNotAllowed for anything else. + """ + if not sql or not sql.strip(): + raise SqlNotAllowed("empty statement; nothing to run") + + try: + toks = _DATABRICKS.tokenize(sql) + except TokenError as exc: + raise SqlNotAllowed( + f"could not tokenize this statement, so it cannot be verified " + f"as read-only ({exc})" + ) from exc + + # A SEMICOLON anywhere but the final token means multiple statements + # were sent. A single trailing ';' is legal and handled above/below. + if any(t.token_type == sqlglot.TokenType.SEMICOLON for t in toks[:-1]): + raise SqlNotAllowed( + "multiple statements are not allowed; send one statement at a time" + ) + + try: + # `exp.Semicolon` is a content-free node sqlglot appends to hold a + # comment that trails the statement terminator (e.g. `SELECT 1; -- + # comment`); it is truthy, so it must be filtered alongside `None` + # or a harmless trailing comment reads as a second statement. + stmts = [ + s + for s in sqlglot.parse(sql, dialect=_DIALECT) + if s and not isinstance(s, exp.Semicolon) + ] + except (ParseError, TokenError) as exc: + # Unparseable is rejected UNLESS the statement's leading token is a + # read-only verb — this covers Databricks extensions sqlglot's + # parser does not model, e.g. `DESC DETAIL`. We cannot safely modify + # text we could not parse, so it is returned as-is (minus a trailing + # terminator). + if toks and str(toks[0].text).upper() in READ_ONLY_VERBS: + return _strip_trailing_semicolon(sql, toks) + raise SqlNotAllowed( + f"this MCP is read-only and could not parse this statement to " + f"confirm that ({exc}). If this is DDL/DML for the target, use " + "the clickhousectl MCP instead." + ) from exc + + if len(stmts) != 1: + raise SqlNotAllowed( + "expected exactly one statement; send one statement at a time" + ) + + root = stmts[0] + + # exp.Command is sqlglot's "I did not model this statement" fallback. + # SHOW/EXPLAIN land here — so do OPTIMIZE and VACUUM, which is exactly + # why the verb check matters: allowing all Command nodes would be a + # new bypass. + if isinstance(root, exp.Command): + verb = str(root.this).upper() + if verb not in READ_ONLY_VERBS: + raise SqlNotAllowed( + f"this MCP is read-only; '{verb}' is not permitted. Run " + "DDL/DML against the target via the clickhousectl MCP " + "instead." + ) + return _strip_trailing_semicolon(sql, toks) + + if not isinstance(root, READ_ONLY_NODES): + raise SqlNotAllowed( + f"this MCP is read-only; '{type(root).__name__}' statements " + "are not permitted. Run DDL/DML against the target via the " + "clickhousectl MCP instead." + ) + + body = _strip_trailing_semicolon(sql, toks) + if ( + isinstance(root, (exp.Select, exp.Union, exp.Subquery)) + and _existing_limit(root) is None + ): + body = f"{body}\nLIMIT {max_rows}" + return body diff --git a/docker/migration-runner/migrationkit/__init__.py b/docker/migration-runner/migrationkit/__init__.py index 0530dcb..59e1083 100644 --- a/docker/migration-runner/migrationkit/__init__.py +++ b/docker/migration-runner/migrationkit/__init__.py @@ -29,6 +29,7 @@ PostgresSource, ClickHouseOssSource, BigQuerySource, + DatabricksSource, ) from .sources.base import UnloadResult from .targets import ClickHouseTarget @@ -48,6 +49,7 @@ "PostgresSource", "ClickHouseOssSource", "BigQuerySource", + "DatabricksSource", "ClickHouseTarget", "S3Stage", "GCSStage", diff --git a/docker/migration-runner/migrationkit/api.py b/docker/migration-runner/migrationkit/api.py index a6ed913..0bbc87a 100644 --- a/docker/migration-runner/migrationkit/api.py +++ b/docker/migration-runner/migrationkit/api.py @@ -365,6 +365,9 @@ def list_source_databases(src: str, refresh: bool = False) -> list[str]: elif src == "bigquery": from .sources.bigquery import BigQuerySource dbs = BigQuerySource.list_databases_from_env() + elif src == "databricks": + from .sources.databricks import DatabricksSource + dbs = DatabricksSource.list_databases_from_env() else: raise HTTPException( status_code=404, diff --git a/docker/migration-runner/migrationkit/migrator.py b/docker/migration-runner/migrationkit/migrator.py index db602d6..449432a 100644 --- a/docker/migration-runner/migrationkit/migrator.py +++ b/docker/migration-runner/migrationkit/migrator.py @@ -242,9 +242,9 @@ def add_table_via_s3( """Register a table for S3-staged migration: source unloads to S3, then ClickHouse Cloud loads via `INSERT FROM s3(...)`. - Only sources that override `Source.unload_to_s3()` (Snowflake - and ClickHouse OSS today) support this path — Migrator validates - at `run()` time.""" + Only sources that override `Source.unload_to_s3()` (Snowflake, + ClickHouse OSS, and Databricks today) support this path — + Migrator validates at `run()` time.""" resolved_target = target_table or name _validate_target_table_name(resolved_target) plan = _S3TablePlan( diff --git a/docker/migration-runner/migrationkit/sources/__init__.py b/docker/migration-runner/migrationkit/sources/__init__.py index 5f9d716..edb0558 100644 --- a/docker/migration-runner/migrationkit/sources/__init__.py +++ b/docker/migration-runner/migrationkit/sources/__init__.py @@ -3,6 +3,7 @@ from .postgres import PostgresSource from .clickhouse_oss import ClickHouseOssSource from .bigquery import BigQuerySource +from .databricks import DatabricksSource __all__ = [ "Source", @@ -10,4 +11,5 @@ "PostgresSource", "ClickHouseOssSource", "BigQuerySource", + "DatabricksSource", ] diff --git a/docker/migration-runner/migrationkit/sources/base.py b/docker/migration-runner/migrationkit/sources/base.py index dcd32a9..dfb308e 100644 --- a/docker/migration-runner/migrationkit/sources/base.py +++ b/docker/migration-runner/migrationkit/sources/base.py @@ -48,6 +48,9 @@ def execute_and_count(self, sql: str) -> tuple[int, float | None, float]: keyed on the session's `cursor.sfqid` - Postgres: `EXPLAIN (ANALYZE, FORMAT JSON)` → `Execution Time` - BigQuery: (planned) `jobs.get` totalSlotMs / endTime-startTime + - Databricks: SQL query-history REST API keyed on `cursor.query_id`, + falling back to `system.query.history`; `None` if + neither is available `server_ms` may be `None` only when the engine genuinely can't return a server-side timing for this query — callers fall back diff --git a/docker/migration-runner/migrationkit/sources/databricks.py b/docker/migration-runner/migrationkit/sources/databricks.py new file mode 100644 index 0000000..e8c9a81 --- /dev/null +++ b/docker/migration-runner/migrationkit/sources/databricks.py @@ -0,0 +1,304 @@ +"""Databricks SQL warehouse as a migration source. + +Mirrors SnowflakeSource: direct batch reads for small tables, plus an +S3 staging path for large facts. Unity Catalog is three-level +(catalog.schema.table) but the playground models one "source database" +per run, so DATABRICKS_NAMESPACE carries "." and +`self.database` is that dotted string. +""" +from __future__ import annotations + +import json +import os +import time +import urllib.error +import urllib.parse +import urllib.request +from typing import Any, Iterator, TYPE_CHECKING + +from .base import Source, UnloadResult + +if TYPE_CHECKING: + from ..staging.s3 import S3Stage + + +def split_namespace(namespace: str) -> tuple[str, str]: + """Split `"."` into its two parts.""" + parts = [p.strip() for p in (namespace or "").split(".") if p.strip()] + if len(parts) != 2: + raise ValueError( + "DATABRICKS_NAMESPACE must be '.' " + f"(e.g. migration_demo.tpch), got {namespace!r}" + ) + return parts[0], parts[1] + + +def normalize_host(raw: str) -> str: + """Bare hostname. The connector's `server_hostname` rejects a scheme.""" + host = (raw or "").strip() + return host.removeprefix("https://").removeprefix("http://").rstrip("/") + + +def parquet_only(objects: list) -> list: + """Keep only Parquet part-files. + + Databricks' commit protocol writes `_committed_*`, `_started_*` and + `_SUCCESS` markers alongside the data. Counting them would inflate the + file count and byte total the dashboard shows for the unload phase. + """ + return [o for o in objects if o.key.lower().endswith(".parquet")] + + +class DatabricksSource(Source): + source_type = "databricks" + + def __init__( + self, + server_hostname: str, + http_path: str, + access_token: str, + catalog: str | None = None, + schema: str | None = None, + ) -> None: + from databricks import sql as dbsql + + self.catalog = catalog + self.schema = schema + self.database = f"{catalog}.{schema}" if catalog and schema else None + self._host = normalize_host(server_hostname) + self._token = access_token + self._conn = dbsql.connect( + server_hostname=self._host, + http_path=http_path, + access_token=access_token, + catalog=catalog, + schema=schema, + ) + + @classmethod + def from_env(cls) -> "DatabricksSource": + catalog = schema = None + namespace = os.environ.get("DATABRICKS_NAMESPACE") + if namespace: + catalog, schema = split_namespace(namespace) + return cls( + server_hostname=os.environ["DATABRICKS_HOST"], + http_path=os.environ["DATABRICKS_HTTP_PATH"], + access_token=os.environ["DATABRICKS_TOKEN"], + catalog=catalog, + schema=schema, + ) + + @classmethod + def list_databases_from_env(cls) -> list[str]: + """Return `catalog.schema` pairs visible to the env credentials. + + Backs the dashboard's source-database dropdown. Every value is + also a usable SQL prefix, which is why the pair is returned as one + dotted string rather than a nested structure. + """ + from databricks import sql as dbsql + + conn = dbsql.connect( + server_hostname=normalize_host(os.environ["DATABRICKS_HOST"]), + http_path=os.environ["DATABRICKS_HTTP_PATH"], + access_token=os.environ["DATABRICKS_TOKEN"], + ) + try: + cur = conn.cursor() + try: + cur.execute( + "SELECT catalog_name, schema_name " + "FROM system.information_schema.schemata " + "WHERE schema_name <> 'information_schema' " + "ORDER BY catalog_name, schema_name" + ) + return [f"{row[0]}.{row[1]}" for row in cur.fetchall()] + finally: + cur.close() + finally: + conn.close() + + def _fq(self, table: str) -> str: + """Fully-qualify a bare table name against the run's namespace.""" + if "." in table: + return table + if not (self.catalog and self.schema): + raise ValueError( + f"table {table!r} is unqualified and no DATABRICKS_NAMESPACE " + f"is set — pass '..{table}' instead" + ) + return f"{self.catalog}.{self.schema}.{table}" + + def count_rows(self, query: str) -> int: + cur = self._conn.cursor() + try: + cur.execute(f"SELECT count(*) FROM ({query})") + (n,) = cur.fetchone() + return int(n) + finally: + cur.close() + + def iter_batches( + self, query: str, batch_size: int + ) -> Iterator[list[dict[str, Any]]]: + cur = self._conn.cursor() + try: + cur.execute(query) + columns = [c[0].lower() for c in cur.description] + while True: + rows = cur.fetchmany(batch_size) + if not rows: + return + yield [dict(zip(columns, row)) for row in rows] + finally: + cur.close() + + def execute_and_count(self, sql: str) -> tuple[int, float | None, float]: + cur = self._conn.cursor() + try: + t0 = time.monotonic() + cur.execute(sql) + rows = cur.fetchall() + wall_ms = (time.monotonic() - t0) * 1000.0 + statement_id = getattr(cur, "query_id", None) + finally: + cur.close() + server_ms = self._fetch_server_ms(statement_id) if statement_id else None + return len(rows), server_ms, wall_ms + + def _fetch_server_ms(self, statement_id: str) -> float | None: + """Server-side execution time in ms, or None. + + Three tiers, because no single surface is reliable: the SQL query + history REST API is near-instant but its response shape is not + contractually stable; `system.query.history` is stable but can lag + by minutes and may not be enabled. Returning None is explicitly + permitted by the Source ABC — Benchmarker falls back to wall_ms. + """ + for attempt in range(2): + ms = self._server_ms_from_history_api(statement_id) + if ms is not None: + return ms + if attempt == 0: + time.sleep(0.25) + return self._server_ms_from_system_table(statement_id) + + def _server_ms_from_history_api(self, statement_id: str) -> float | None: + """Returns a float or None — never raises, even for a malformed or + hostile response body. The response shape is a guess against a + live API (see module docstring / _fetch_server_ms), so both the + network call *and* the body-processing that follows it must + degrade to None rather than let an AttributeError/TypeError + escape and fail the whole benchmark row.""" + try: + filter_by = json.dumps({"statement_ids": [statement_id]}) + query = urllib.parse.urlencode({"filter_by": filter_by}) + url = f"https://{self._host}/api/2.0/sql/history/queries?{query}" + request = urllib.request.Request( + url, headers={"Authorization": f"Bearer {self._token}"} + ) + with urllib.request.urlopen(request, timeout=10) as response: + body = json.load(response) + for item in body.get("res") or []: + metrics = item.get("metrics") or {} + for key in ("execution_time_ms", "total_time_ms"): + if metrics.get(key) is not None: + return float(metrics[key]) + if item.get("duration") is not None: + return float(item["duration"]) + return None + except Exception: + return None + + def _server_ms_from_system_table(self, statement_id: str) -> float | None: + # NOTE: deviates from the brief, which interpolated statement_id + # into the SQL string with an f-string. statement_id comes from + # the connector (cursor.query_id), not user input, but the + # connector supports native positional binding (`?` placeholders, + # documented for databricks-sql-connector 3.0.0+), so we use that + # instead of hand-rolled string interpolation. + sql = ( + "SELECT execution_duration_ms, total_duration_ms " + "FROM system.query.history " + "WHERE statement_id = ? LIMIT 1" + ) + try: + cur = self._conn.cursor() + try: + cur.execute(sql, [statement_id]) + row = cur.fetchone() + finally: + cur.close() + except Exception: + return None + if not row: + return None + for value in row: + if value is not None: + return float(value) + return None + + def unload_to_s3( + self, + table: str, + stage: "S3Stage", + run_id: str, + file_format: str = "parquet", + ) -> UnloadResult: + """Bulk-export `table` to the per-run S3 prefix with + `INSERT OVERWRITE DIRECTORY ... USING PARQUET`. + + Requires a Unity Catalog external location over the staging bucket + with WRITE FILES granted to this principal — the `demo` Terraform + module provisions one. Idempotent: OVERWRITE replaces this table's + files without touching others in the run. + """ + from ..staging.s3 import list_s3_objects + + if file_format.lower() != "parquet": + raise ValueError( + f"unload_to_s3: only parquet is supported, got {file_format!r}" + ) + + # Resolve before the try so an unqualified `table` (usage error) + # raises its own ValueError instead of being caught below and + # re-wrapped as a misleading "permissions" RuntimeError. + fq_table = self._fq(table) + + target_uri = stage.s3_uri(run_id, table) + cur = self._conn.cursor() + try: + t0 = time.monotonic() + try: + cur.execute( + f"INSERT OVERWRITE DIRECTORY '{target_uri}'\n" + f"USING PARQUET\n" + f"SELECT * FROM {fq_table}" + ) + except Exception as exc: + raise RuntimeError( + f"Databricks refused to write Parquet to {target_uri}. " + "The staging bucket needs a Unity Catalog external " + "location with WRITE FILES granted to this principal — " + "sources/databricks/terraform/demo provisions one when " + "enable_s3_staging=true. Until then use " + "Migrator.add_table() for this table; the direct path " + f"needs no external location. Original error: {exc}" + ) from exc + seconds = round(time.monotonic() - t0, 3) + finally: + cur.close() + + files = parquet_only(list_s3_objects(stage, run_id, table)) + return UnloadResult( + file_count=len(files), + total_bytes=sum(f.size for f in files), + seconds=seconds, + ) + + def close(self) -> None: + try: + self._conn.close() + except Exception: + pass diff --git a/docker/migration-runner/requirements.txt b/docker/migration-runner/requirements.txt index 77633d3..9e25acb 100644 --- a/docker/migration-runner/requirements.txt +++ b/docker/migration-runner/requirements.txt @@ -10,3 +10,4 @@ google-cloud-bigquery>=3.21 google-cloud-storage>=2.16 db-dtypes>=1.2 pymongo>=4.6 +databricks-sql-connector>=4.0 diff --git a/docs/adding-a-source.md b/docs/adding-a-source.md index 2c30e13..eff8121 100644 --- a/docs/adding-a-source.md +++ b/docs/adding-a-source.md @@ -24,10 +24,23 @@ sources/ │ ├── prompts/ │ └── GUIDE.md │ -└── snowflake/ ← Snowflake → ClickHouse Cloud (cloud source — no docker/) +├── snowflake/ ← Snowflake → ClickHouse Cloud (cloud source — no docker/) +│ ├── queries/ +│ ├── scripts/ +│ ├── prompts/ +│ └── GUIDE.md +│ +├── bigquery/ ← BigQuery → ClickHouse Cloud (cloud source — no docker/) +│ ├── queries/ +│ ├── scripts/ +│ ├── prompts/ +│ └── GUIDE.md +│ +└── databricks/ ← Databricks → ClickHouse Cloud (cloud source — no docker/) ├── queries/ ├── scripts/ ├── prompts/ + ├── terraform/ ← optional demo/workspace provisioning └── GUIDE.md ``` @@ -58,7 +71,7 @@ The components fall into five groups: - **`migration-runner`** — multi-purpose service. Hosts: - **MCP server** on SSE: exposes `run_python`, `run_python_background`, `tail_python_job`, `write_workspace_file`, `read_workspace_file`, `list_workspace_files`. This is the agent's Python sandbox. - **FastAPI** on `:8001`: REST + SSE for the dashboard. Endpoints under `/api/mk/runs/*` and `/api/mk/sources/*`. Source manifests and conversation pre-creation live here too. - - **`migrationkit` Python library** (importable inside `run_python`): `Migrator`, `Validator`, `Benchmarker`, the `Source` ABC (with concrete `PostgresSource`, `SnowflakeSource`, `BigQuerySource`, `ClickHouseOssSource`), `ClickHouseTarget`, and `S3Stage` / `GCSStage` for object-storage staging paths. + - **`migrationkit` Python library** (importable inside `run_python`): `Migrator`, `Validator`, `Benchmarker`, the `Source` ABC (with concrete `PostgresSource`, `SnowflakeSource`, `BigQuerySource`, `DatabricksSource`, `ClickHouseOssSource`), `ClickHouseTarget`, and `S3Stage` / `GCSStage` for object-storage staging paths. - **SQLite WAL state store** at `/workspace/state/migrationkit.db`. Authoritative state across the stack: `runs`, `run_tables`, `events`, `batches`, `controls`, `validations`, `benchmarks`. Concurrent writers in the same container; readers via FastAPI. **Source MCP layer** — one MCP server per supported source, all exposed to LibreChat over SSE: @@ -67,6 +80,7 @@ The components fall into five groups: - **`clickhouse-oss-mcp`** (`mcp-clickhouse`) — `run_select_query`. - **`snowflake-source`** (`snowflake-labs/mcp`) + **`snowflake-source-shim`** — the shim is a Python proxy that strips JSON-Schema fields (`exclusiveMaximum`, `const`, `oneOf`, `allOf`, `$schema`) that Gemini's function-calling API rejects. LibreChat connects to the shim, not the upstream MCP. - **`bigquery-source`** (Google MCP toolbox) — `bigquery-list-dataset-ids`, `bigquery-list-table-ids`, `bigquery-get-table-info`, plus `INFORMATION_SCHEMA` access. +- **`databricks-mcp`** (`docker/databricks-mcp`, purpose-built in-repo) — `list_catalogs`, `list_schemas`, `list_tables`, `describe_table`, `run_select_query`. Databricks publishes no introspect-and-SELECT MCP for SQL warehouses (their own `databricks-mcp` package is an OAuth helper for the hosted UC-functions / vector-search / Genie servers, not this), so this one is authored in-repo, in the same spirit as `clickhousectl-mcp`. Because we author its tool schemas directly, it needs no Gemini shim — unlike `snowflake-source`. **Target MCP — write-enabled:** @@ -77,7 +91,7 @@ The components fall into five groups: - **`postgres`** — PostgreSQL 16 with `ecommerce` (~10M rows) and optional `tpch` (SF1) datasets. - **`clickhouse-oss`** — ClickHouse OSS with `analytics` (web events, ~12.2M rows) and optional `tpch` (SF1) datasets. -Snowflake and BigQuery don't have bundled databases — the source MCPs connect to partner-provided cloud accounts. +Snowflake, BigQuery, and Databricks don't have bundled databases — the source MCPs connect to partner-provided cloud accounts. --- @@ -87,7 +101,7 @@ Snowflake and BigQuery don't have bundled databases — the source MCPs connect ### MCP tools -Exposed via SSE to LibreChat. All four agents attach this MCP. +Exposed via SSE to LibreChat. All five agents attach this MCP. | Tool | Use for | |---|---| @@ -97,7 +111,7 @@ Exposed via SSE to LibreChat. All four agents attach this MCP. | `write_workspace_file(path, content)` | Writes a file under `/workspace/`. The agent uses this to persist the migration script before dispatching it. | | `read_workspace_file(path)` / `list_workspace_files()` | Read-only file ops on `/workspace/`. | -**Why dispatch + ONE tail instead of polling:** every `tail_python_job` call replays the full conversation context to the LLM — repeated polling burns tokens quadratically with no UX benefit. The dashboard's Migration tab already streams per-table progress via SSE; the agent's job is to dispatch + stop, and the partner watches the dashboard. This is enforced in every source's `sources//prompts/02-migrate-data.md` and in the system prompts at `librechat/sources/{snowflake,bigquery}-instructions.md`. +**Why dispatch + ONE tail instead of polling:** every `tail_python_job` call replays the full conversation context to the LLM — repeated polling burns tokens quadratically with no UX benefit. The dashboard's Migration tab already streams per-table progress via SSE; the agent's job is to dispatch + stop, and the partner watches the dashboard. This is enforced in every source's `sources//prompts/02-migrate-data.md` and in the system prompts at `librechat/sources/{snowflake,bigquery,databricks}-instructions.md`. ### FastAPI HTTP surface @@ -135,7 +149,7 @@ The Python library inside the runner that handles data movement. Critical pieces ```python from migrationkit import ( Migrator, Validator, Benchmarker, - PostgresSource, SnowflakeSource, BigQuerySource, ClickHouseOssSource, + PostgresSource, SnowflakeSource, BigQuerySource, DatabricksSource, ClickHouseOssSource, ClickHouseTarget, S3Stage, GCSStage, ) @@ -361,7 +375,12 @@ volumes: -data: ``` -> **Port allocation:** Postgres MCP uses host port `8001`, ClickHouse OSS MCP uses `8002`. Use the next available port (e.g. `8003`) for a third source to avoid conflicts. +> **Port allocation:** currently in use — `8001` Postgres MCP, `8002` +> ClickHouse OSS MCP, `8003` clickhousectl MCP, `8004` Snowflake source, +> `8005`/`8006` migration-runner (MCP / HTTP API), `8007` BigQuery +> source, `8008` Databricks MCP, `8014` Snowflake's Gemini-compat shim. +> Use the next available port (e.g. `8009`) for a new source to avoid +> conflicts. Also add the new MCP service to LibreChat's `depends_on`: @@ -428,6 +447,7 @@ Do **not** add the cloud MCP to LibreChat's `depends_on` with `service_healthy` | SQLite | `@modelcontextprotocol/server-sqlite` | Wraps with supergateway | | Snowflake | `@datawizardinc/mcp-snowflake-server` | Wraps with supergateway | | BigQuery | `@ergut/mcp-bigquery-server` | Wraps with supergateway | +| Databricks | *(none suitable)* | Databricks publishes no introspect-and-SELECT MCP for SQL warehouses; see `docker/databricks-mcp` for a purpose-built one | | Redshift | use Postgres MCP with Redshift endpoint | Standard psycopg2 connection | > MCP package availability and names change frequently. Check [npmjs.com](https://www.npmjs.com) and [glama.ai/mcp/servers](https://glama.ai/mcp/servers) for the latest options before wiring up a new source. @@ -623,11 +643,12 @@ The agent's behaviour is controlled by a set of modular instruction files. Edit | File | Purpose | |---|---| -| [librechat/clickhouse-cloud-instructions.md](../librechat/clickhouse-cloud-instructions.md) | Base ClickHouse Cloud rules — injected into `mcpServers.clickhousectl.serverInstructions` (shared by all four agents). | +| [librechat/clickhouse-cloud-instructions.md](../librechat/clickhouse-cloud-instructions.md) | Base ClickHouse Cloud rules — injected into `mcpServers.clickhousectl.serverInstructions` (shared by all five agents). | | `agent-skills/.../AGENTS.md` | ClickHouse best-practices skill (cloned from [ClickHouse/agent-skills](https://github.com/ClickHouse/agent-skills) as a submodule). Appended to the same `clickhousectl` instructions. | | [librechat/sources/postgres-instructions.md](../librechat/sources/postgres-instructions.md) | Postgres-specific rules. | | [librechat/sources/snowflake-instructions.md](../librechat/sources/snowflake-instructions.md) | Snowflake-specific rules. | | [librechat/sources/bigquery-instructions.md](../librechat/sources/bigquery-instructions.md) | BigQuery-specific rules. | +| [librechat/sources/databricks-instructions.md](../librechat/sources/databricks-instructions.md) | Databricks-specific rules. | | [librechat/sources/clickhouse-oss-instructions.md](../librechat/sources/clickhouse-oss-instructions.md) | ClickHouse OSS-specific rules. | Each pre-built agent attaches exactly the MCPs for its source, so it transparently receives only the relevant `serverInstructions` — no per-agent `instructions` field is set. `make setup` rebuilds the injected blocks idempotently. **Don't edit anything below the `--- Migration Rules (auto-injected …) ---` marker in `librechat.yaml`**; the MCP-purpose blurb above the marker IS hand-editable. @@ -658,15 +679,19 @@ 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 + Gemini shim make up-bigquery # also start the BigQuery source MCP +make up-databricks # also start the Databricks source MCP (databricks-mcp) make snowflake-setup # set up MIGRATION_DEMO.RETAIL workload in existing Snowflake (Path A) make snowflake-provision # provision a fresh Snowflake demo env with Terraform (Path B) +make databricks-setup # seed migration_demo.tpch in an existing Databricks workspace (manual entry path) +make databricks-provision # provision the demo catalog/schema/warehouse into an existing workspace with Terraform +make databricks-provision-workspace # provision a fresh serverless workspace, then chain into databricks-provision make tpch-data # generate TPC-H SF1 .tbl files in workloads/tpch/data/sf1/ make tpch-load-bigquery # load TPC-H + augmentations into BigQuery make tpch-load-postgres # load TPC-H into Postgres make tpch-load-clickhouse-oss # load TPC-H into the bundled ClickHouse OSS make down # stop without removing data make reset # destroy volumes and start fresh -make reset-agent # delete + recreate the four pre-built agents (after AGENT_PROVIDER/AGENT_MODEL changes in .env) +make reset-agent # delete + recreate the five pre-built agents (after AGENT_PROVIDER/AGENT_MODEL changes in .env) make health # check all services are healthy make migration-status # check progress of a running migration script (target row counts) make logs # tail all service logs diff --git a/docs/architecture.mmd b/docs/architecture.mmd index 84beb86..f8e86cd 100644 --- a/docs/architecture.mmd +++ b/docs/architecture.mmd @@ -42,6 +42,7 @@ flowchart TB sf_shim["snowflake-source-shim\n(Gemini schema-strip proxy)"] sf_mcp["snowflake-source\n(Snowflake-Labs MCP)"] bq_mcp["bigquery-source\n(Google MCP toolbox)"] + dbx_mcp["databricks-mcp\n(purpose-built, in-repo)"] chctl["clickhousectl-mcp\nrun_query · run_command\n(write-enabled CH Cloud MCP)"] sf_shim --> sf_mcp end @@ -61,6 +62,7 @@ flowchart TB direction LR sf_acct[("Snowflake account")] bq_proj[("BigQuery project")] + dbx_ws[("Databricks workspace")] end subgraph staging["Optional object-storage staging"] direction LR @@ -86,6 +88,7 @@ flowchart TB librechat <-->|"SSE"| ch_oss_mcp librechat <-->|"SSE"| sf_shim librechat <-->|"SSE"| bq_mcp + librechat <-->|"SSE"| dbx_mcp librechat <-->|"SSE"| chctl %% migrationkit data plane (read source, write target, write state) @@ -95,6 +98,7 @@ flowchart TB migrationkit -->|"read"| ch_oss migrationkit -->|"read"| sf_acct migrationkit -->|"read"| bq_proj + migrationkit -->|"read"| dbx_ws migrationkit -->|"HTTPS INSERT"| ch_cloud %% Source-MCP backing-DB connections @@ -102,6 +106,7 @@ flowchart TB ch_oss_mcp <--> ch_oss sf_mcp <-->|"snowflake-connector"| sf_acct bq_mcp <-->|"google-cloud-bigquery"| bq_proj + dbx_mcp <-->|"databricks-sql-connector"| dbx_ws chctl <-->|"HTTPS"| ch_cloud %% Optional staging paths (dashed) @@ -119,10 +124,10 @@ flowchart TB classDef dbBox fill:#f3e8ff,stroke:#9333ea,stroke-width:1.5px,color:#3b0764 classDef stagingBox fill:#e0f2fe,stroke:#0284c7,stroke-width:1.5px,color:#0c4a6e - class llm,ch_cloud,sf_acct,bq_proj cloudBox + class llm,ch_cloud,sf_acct,bq_proj,dbx_ws cloudBox class s3,gcs stagingBox class partner hostNode class nginx,dashboard,librechat,mongo uiBox class runner_mcp,runner_api,runner_state,migrationkit runnerBox - class pg_mcp,ch_oss_mcp,sf_mcp,sf_shim,bq_mcp,chctl mcpBox + class pg_mcp,ch_oss_mcp,sf_mcp,sf_shim,bq_mcp,dbx_mcp,chctl mcpBox class pg,ch_oss dbBox diff --git a/docs/architecture.png b/docs/architecture.png index a39515d..f1682a7 100644 Binary files a/docs/architecture.png and b/docs/architecture.png differ diff --git a/docs/migration-checklist.md b/docs/migration-checklist.md index f29bcce..84b1de6 100644 --- a/docs/migration-checklist.md +++ b/docs/migration-checklist.md @@ -5,8 +5,8 @@ to track progress on a source → ClickHouse Cloud migration. The phases map 1:1 to the dashboard's six steps and to each source's `sources//prompts/01..06-*.md` files. -This checklist is shared across all four supported sources (Snowflake, -BigQuery, Postgres, ClickHouse OSS). The type-mapping table in Phase 2 +This checklist is shared across all five supported sources (Snowflake, +BigQuery, Databricks, Postgres, ClickHouse OSS). The type-mapping table in Phase 2 is written with **Postgres-specific** types as a concrete example; other source engines have analogous mappings — see the per-source prompts and `librechat/sources/-instructions.md` for the @@ -20,6 +20,7 @@ authoritative lists. - [ ] Query patterns analysed — WHERE / GROUP BY / JOIN columns identified - [ ] Source-engine-specific features listed for translation (VARIANT / Streams / Dynamic Tables for Snowflake; STRUCT / Materialized Views for BigQuery; + VARIANT / liquid clustering / deletion vectors for Databricks; JSONB / arrays / ENUMs for Postgres; AggregatingMergeTree / MVs for CH OSS) ## Phase 2 — ClickHouse Schema Design diff --git a/librechat/librechat.yaml b/librechat/librechat.yaml index 57ae809..a8d9cc6 100644 --- a/librechat/librechat.yaml +++ b/librechat/librechat.yaml @@ -22,6 +22,7 @@ mcpSettings: - snowflake-source - snowflake-source-shim - bigquery-source + - databricks-mcp - migration-runner mcpServers: postgres-source: @@ -1519,18 +1520,256 @@ mcpServers: Bytes scanned during a migration session shows up on the partner's GCP bill — flag if the workload is large enough that this matters. + databricks-source: + type: sse + url: http://databricks-mcp:8000/sse + timeout: 60000 + serverInstructions: |- + SOURCE Databricks SQL warehouse. Read-only — statements are guarded + so only SELECT / WITH / SHOW / DESCRIBE / EXPLAIN reach the + warehouse. Unity Catalog is three-level (catalog.schema.table); the + schema is unknown ahead of time, so discover it dynamically with + list_catalogs, list_schemas, list_tables, describe_table, and + run_select_query. Use this MCP — not migration-runner — for all + source introspection. + + --- Migration Rules (auto-injected, do not edit below) --- + + ## Databricks Source — Migration Instructions + + This section applies when the SOURCE database is Databricks. + + --- + + ## Unity Catalog is Three-Level — Always Fully Qualify + + Databricks names objects `catalog.schema.table`. ClickHouse has two levels + (`database.table`). + + - **Always fully qualify** source tables in generated SQL and in + `migrationkit` `source_query=` values. A bare `orders` resolves against + whatever the connection's default namespace happens to be, which may not + be the namespace the partner selected. + - The playground's "source database" is the `catalog.schema` pair as one + dotted string (`DATABRICKS_NAMESPACE`, e.g. `migration_demo.tpch`). When + you see a single dotted value in a prompt, that is what it is. + - Identifiers are case-insensitive and stored lower-case. Quote with + backticks when a name needs it — never double quotes. + + --- + + ## Schema Discovery — Don't Assume Anything + + Never assume column names, table names, types, or that any particular + Databricks feature is or is not in use. Discover the real schema at + runtime via the `databricks-source` MCP. + + **Use `databricks-source` (not `migration-runner`) for all schema + discovery and read-only inspection.** `migration-runner`'s `run_python` is + **only** for data movement once the schema is understood. Reasons: + + - `databricks-source` is workspace-scoped: `list_catalogs` returns every + catalog the principal can see, not just one default namespace. + - `migration-runner` inherits env vars from the playground's `.env` via + `env_file:`. A stale `DATABRICKS_NAMESPACE` there silently scopes a + connector session to the wrong catalog/schema and your inventory + quietly misses everything else. + - The MCP is guarded read-only, so introspection cannot mutate the source + even by accident. + + **The discovery checklist for every Databricks migration:** + + 1. **List catalogs and schemas:** + `list_catalogs()` then `list_schemas(catalog)`. + + 2. **List tables with sizes:** + `list_tables(catalog, schema)` — returns `table_type`, `comment`, + `sizeInBytes`, `numFiles`. It deliberately does **not** return row + counts: Delta metadata doesn't carry them and a per-table `COUNT(*)` + would make the call slow. + + 3. **Get row counts in ONE query**, not one call per table: + ```sql + SELECT 'orders' AS t, count(*) AS n FROM ..orders + UNION ALL SELECT 'lineitem', count(*) FROM ..lineitem + ORDER BY n DESC + ``` + + 4. **Full schema plus Delta detail per table:** + `describe_table(catalog, schema, table)`. Read all three sections of + the response — `describe_extended` for columns, `detail` for + clustering/partition columns and table features, `history` for whether + the table is actively mutated. Clustering columns and deletion-vector + state live only in `detail`. + + 5. **Inventory the source-specific features you actually found:** + ```sql + -- VARIANT / nested columns, generated columns, identity columns + SELECT column_name, full_data_type, is_nullable, generation_expression + FROM system.information_schema.columns + WHERE table_catalog = '' AND table_schema = '' + ORDER BY table_name, ordinal_position + ``` + ```sql + -- Materialized views and streaming tables are separate object types; + -- a table listing alone misses them. + SELECT table_name, table_type + FROM system.information_schema.tables + WHERE table_catalog = '' AND table_schema = '' + ``` + + 6. **Sample data and check nullability + cardinality** for every column + before choosing types. This is what tells you whether a column is + `Nullable(T)`, whether a string is a `LowCardinality(String)` + candidate, and whether a `DECIMAL` needs full precision: + ```sql + SELECT * FROM .. LIMIT 5 + SELECT count(*), count(), count(DISTINCT ) FROM ..
+ ``` + + 7. **Read the partner's queries.** `ORDER BY` key selection on the + ClickHouse side comes from the columns in WHERE / JOIN / GROUP BY of + the actual workload — **not** from the source's `CLUSTER BY` columns. + + Produce a migration inventory before generating any target schema. + + --- + + ## Databricks → ClickHouse Type Mapping + + | Databricks | ClickHouse | Notes | + |---|---|---| + | `BOOLEAN` | `Bool` | | + | `TINYINT` / `SMALLINT` / `INT` / `BIGINT` | `Int8` / `Int16` / `Int32` / `Int64` | Databricks integers are always signed | + | `FLOAT` / `DOUBLE` | `Float32` / `Float64` | | + | `DECIMAL(p, s)` | `Decimal(p, s)` | **Never** `Float64` — TPC-H money columns need exact arithmetic | + | `STRING` | `String`, or `LowCardinality(String)` | Use `LowCardinality` below ~10k distinct values | + | `BINARY` | `String` | | + | `DATE` | `Date32` | `Date` only reaches 2149 and starts at 1970 | + | `TIMESTAMP` | `DateTime64(6, 'UTC')` | Databricks `TIMESTAMP` is instant-with-timezone; normalise to UTC at the source | + | `TIMESTAMP_NTZ` | `DateTime64(6)` | No timezone — do **not** attach one | + | `INTERVAL` | `Int64` seconds, or `String` | No native equivalent; state the unit in a comment | + | `ARRAY` | `Array(T)` | | + | `MAP` | `Map(K, V)` | Missing-key lookup returns the type default, not NULL | + | `STRUCT` | named `Tuple(a T, …)` | Positional in ClickHouse — field order matters | + | `ARRAY>` | `Nested(...)` or `Array(Tuple(...))` | `Nested` is easier to query; `Array(Tuple)` is easier to insert | + | `VARIANT` | `JSON`, or hot keys extracted to typed columns | Extracting the 2–3 keys the workload filters on usually beats a whole `JSON` column | + | generated column | `MATERIALIZED` or `ALIAS` | `MATERIALIZED` stores it, `ALIAS` recomputes on read | + | identity column | `Int64` + no auto-generation | ClickHouse has no identity; the migrated values are what you keep | + + **Nullability:** `system.information_schema.columns.is_nullable` is + authoritative. Either declare `Nullable()` on the target, or declare + non-Nullable with an explicit `DEFAULT` **and** map `None` to that default + in a step-2 `transform=`. A non-Nullable target column with neither fails + mid-batch on the first NULL. + + --- + + ## Delta / Unity Catalog Objects — Migration Patterns + + ### Liquid clustering (`CLUSTER BY`) + Not an index and not a sort order you can copy. Treat it as a hint about + which columns the workload filters on, then choose the ClickHouse + `ORDER BY` from the actual queries. Say in chat when your choice differs + from the source clustering and why. + + ### Deletion vectors + Deletes are recorded as vectors rather than rewritten files, so a table can + report rows that a `SELECT` won't return. Two consequences: source + `COUNT(*)` is the only trustworthy count (never sum file statistics), and a + table being actively deleted from can legitimately change count mid-run. + + ### Time travel (`VERSION AS OF` / `TIMESTAMP AS OF`) + No ClickHouse equivalent. Use it to pin a consistent read for the + migration itself (`SELECT * FROM t VERSION AS OF `) so a concurrent + write doesn't skew validation. Do not try to reproduce the history. + + ### Materialized views and streaming tables + Recreate as a ClickHouse Materialized View over an `AggregatingMergeTree`, + and **backfill** it from the base table after loading — an MV only sees + rows inserted after it exists. Get the defining query from + `describe_table`'s extended output. + + ### Change Data Feed (`delta.enableChangeDataFeed`) + Ongoing CDC is out of scope for a one-shot migration. If the partner needs + it, the answer is ClickPipes or a `ReplacingMergeTree` with a version + column — say so and move on rather than half-building it. + + ### Unity Catalog volumes and external locations + Only relevant to the S3 staging path. A `/Volumes/...` path is not + readable by ClickHouse; the staged unload writes to `s3://` directly. + + --- + + ## Migration Script Rules + + - Connect with `DatabricksSource.from_env()`. Never construct + `databricks.sql.connect(...)` by hand in a generated script — the + helper owns namespace splitting and host normalisation. + - Fully qualify every `source_query=`. + - Row dict keys are **lower-case**; `iter_batches` lowercases them. + - `batch_size` ~100k narrow, ~25–50k for rows carrying VARIANT or + `ARRAY`. Never above 500k. + - **S3 staging requires a Unity Catalog external location** over the + staging bucket with `WRITE FILES` granted. `unload_to_s3` raises a + message saying exactly that; fall back to `add_table()` per table. + - **Target column order must match the source** for staged tables — that + path is `INSERT INTO … SELECT * FROM s3(...)`, which is positional. + - The staged path supports neither `batch_size` nor `transform`. Any + table needing per-row transformation uses the direct path. + + --- + + ## Query Rewriting Notes + + | Databricks | ClickHouse | + |---|---| + | `QUALIFY ` | subquery + `WHERE` over the window | + | `LATERAL VIEW explode(arr) AS x` | `ARRAY JOIN` / `arrayJoin(arr)` | + | `transform(arr, x -> f(x))` | `arrayMap(x -> f(x), arr)` (order flips) | + | `filter(arr, x -> p(x))` | `arrayFilter(x -> p(x), arr)` (order flips) | + | `aggregate(arr, 0, (a, x) -> a + x)` | `arraySum(arr)` / `arrayReduce('sum', arr)` | + | `v:a.b` / `variant_get(v, '$.a')` | `JSONExtract*(v, 'a')`, or `v.a` on a `JSON` column | + | `named_struct('a', 1)` | `tuple(1)` — positional | + | `try_divide(a, b)` | `if(b = 0, NULL, a / b)` — `/` returns `inf`, not NULL | + | `try_cast(x AS t)` | `accurateCastOrNull(x, 't')` | + | `datediff(a, b)` | `dateDiff('day', b, a)` — order flips | + | `date_trunc('month', d)` | `toStartOfMonth(d)` | + | `SEMI JOIN` / `ANTI JOIN` | `LEFT SEMI JOIN` / `LEFT ANTI JOIN` | + | `catalog.schema.table` | `database.table` | + + State `NULLS FIRST` / `NULLS LAST` explicitly — the engines' defaults differ. + + --- + + ## Known Gotchas + + - **Cold warehouse.** A stopped SQL warehouse takes 30s+ to serve its + first statement. Never present that as query latency; run a throwaway + `SELECT 1` before benchmarking. + - **`samples` catalog is read-only.** Copy out of it, never into it. + - **`VARIANT` needs DBSQL 2024.35+ / DBR 15.3+.** On older runtimes the + column type simply doesn't exist. + - **`system.query.history` can lag** by minutes and may be disabled, so + `server_ms` is sometimes `None`. Report wall-clock as wall-clock. + - **Auto-stopped warehouse mid-migration.** A long direct-path migration + keeps the connection busy, but a paused run can let the warehouse stop; + resuming pays cold-start again. Expected, not a bug. migration-runner: type: sse url: http://migration-runner:8000/sse timeout: 900000 serverInstructions: |- Execute Python migration scripts inside the playground. Pre-installed: - snowflake-connector-python, psycopg2, clickhouse-connect, pyarrow. + snowflake-connector-python, psycopg2, clickhouse-connect, pyarrow, + databricks-sql-connector. Reachable hosts: postgres:5432, clickhouse-oss:8123, Snowflake (via - SNOWFLAKE_* env vars), ClickHouse Cloud (via CLICKHOUSE_CLOUD_*). + SNOWFLAKE_* env vars), Databricks (via DATABRICKS_* env vars), + ClickHouse Cloud (via CLICKHOUSE_CLOUD_*). Tools: - run_python(code, timeout_seconds=600) -> {stdout, stderr, exit_code, duration_seconds} - write_workspace_file(path, content) / read_workspace_file(path) / list_workspace_files() - Use this when the source is Snowflake (see snowflake-instructions.md). + Use this for any source's migration scripts (see that source's + *-instructions.md). Stream long outputs to the user. On errors, surface the full traceback and propose a fix before retrying. diff --git a/librechat/sources/databricks-instructions.md b/librechat/sources/databricks-instructions.md new file mode 100644 index 0000000..e4b76d3 --- /dev/null +++ b/librechat/sources/databricks-instructions.md @@ -0,0 +1,220 @@ +## Databricks Source — Migration Instructions + +This section applies when the SOURCE database is Databricks. + +--- + +## Unity Catalog is Three-Level — Always Fully Qualify + +Databricks names objects `catalog.schema.table`. ClickHouse has two levels +(`database.table`). + +- **Always fully qualify** source tables in generated SQL and in + `migrationkit` `source_query=` values. A bare `orders` resolves against + whatever the connection's default namespace happens to be, which may not + be the namespace the partner selected. +- The playground's "source database" is the `catalog.schema` pair as one + dotted string (`DATABRICKS_NAMESPACE`, e.g. `migration_demo.tpch`). When + you see a single dotted value in a prompt, that is what it is. +- Identifiers are case-insensitive and stored lower-case. Quote with + backticks when a name needs it — never double quotes. + +--- + +## Schema Discovery — Don't Assume Anything + +Never assume column names, table names, types, or that any particular +Databricks feature is or is not in use. Discover the real schema at +runtime via the `databricks-source` MCP. + +**Use `databricks-source` (not `migration-runner`) for all schema +discovery and read-only inspection.** `migration-runner`'s `run_python` is +**only** for data movement once the schema is understood. Reasons: + +- `databricks-source` is workspace-scoped: `list_catalogs` returns every + catalog the principal can see, not just one default namespace. +- `migration-runner` inherits env vars from the playground's `.env` via + `env_file:`. A stale `DATABRICKS_NAMESPACE` there silently scopes a + connector session to the wrong catalog/schema and your inventory + quietly misses everything else. +- The MCP is guarded read-only, so introspection cannot mutate the source + even by accident. + +**The discovery checklist for every Databricks migration:** + +1. **List catalogs and schemas:** + `list_catalogs()` then `list_schemas(catalog)`. + +2. **List tables with sizes:** + `list_tables(catalog, schema)` — returns `table_type`, `comment`, + `sizeInBytes`, `numFiles`. It deliberately does **not** return row + counts: Delta metadata doesn't carry them and a per-table `COUNT(*)` + would make the call slow. + +3. **Get row counts in ONE query**, not one call per table: + ```sql + SELECT 'orders' AS t, count(*) AS n FROM ..orders + UNION ALL SELECT 'lineitem', count(*) FROM ..lineitem + ORDER BY n DESC + ``` + +4. **Full schema plus Delta detail per table:** + `describe_table(catalog, schema, table)`. Read all three sections of + the response — `describe_extended` for columns, `detail` for + clustering/partition columns and table features, `history` for whether + the table is actively mutated. Clustering columns and deletion-vector + state live only in `detail`. + +5. **Inventory the source-specific features you actually found:** + ```sql + -- VARIANT / nested columns, generated columns, identity columns + SELECT column_name, full_data_type, is_nullable, generation_expression + FROM system.information_schema.columns + WHERE table_catalog = '' AND table_schema = '' + ORDER BY table_name, ordinal_position + ``` + ```sql + -- Materialized views and streaming tables are separate object types; + -- a table listing alone misses them. + SELECT table_name, table_type + FROM system.information_schema.tables + WHERE table_catalog = '' AND table_schema = '' + ``` + +6. **Sample data and check nullability + cardinality** for every column + before choosing types. This is what tells you whether a column is + `Nullable(T)`, whether a string is a `LowCardinality(String)` + candidate, and whether a `DECIMAL` needs full precision: + ```sql + SELECT * FROM ..
LIMIT 5 + SELECT count(*), count(), count(DISTINCT ) FROM ..
+ ``` + +7. **Read the partner's queries.** `ORDER BY` key selection on the + ClickHouse side comes from the columns in WHERE / JOIN / GROUP BY of + the actual workload — **not** from the source's `CLUSTER BY` columns. + +Produce a migration inventory before generating any target schema. + +--- + +## Databricks → ClickHouse Type Mapping + +| Databricks | ClickHouse | Notes | +|---|---|---| +| `BOOLEAN` | `Bool` | | +| `TINYINT` / `SMALLINT` / `INT` / `BIGINT` | `Int8` / `Int16` / `Int32` / `Int64` | Databricks integers are always signed | +| `FLOAT` / `DOUBLE` | `Float32` / `Float64` | | +| `DECIMAL(p, s)` | `Decimal(p, s)` | **Never** `Float64` — TPC-H money columns need exact arithmetic | +| `STRING` | `String`, or `LowCardinality(String)` | Use `LowCardinality` below ~10k distinct values | +| `BINARY` | `String` | | +| `DATE` | `Date32` | `Date` only reaches 2149 and starts at 1970 | +| `TIMESTAMP` | `DateTime64(6, 'UTC')` | Databricks `TIMESTAMP` is instant-with-timezone; normalise to UTC at the source | +| `TIMESTAMP_NTZ` | `DateTime64(6)` | No timezone — do **not** attach one | +| `INTERVAL` | `Int64` seconds, or `String` | No native equivalent; state the unit in a comment | +| `ARRAY` | `Array(T)` | | +| `MAP` | `Map(K, V)` | Missing-key lookup returns the type default, not NULL | +| `STRUCT` | named `Tuple(a T, …)` | Positional in ClickHouse — field order matters | +| `ARRAY>` | `Nested(...)` or `Array(Tuple(...))` | `Nested` is easier to query; `Array(Tuple)` is easier to insert | +| `VARIANT` | `JSON`, or hot keys extracted to typed columns | Extracting the 2–3 keys the workload filters on usually beats a whole `JSON` column | +| generated column | `MATERIALIZED` or `ALIAS` | `MATERIALIZED` stores it, `ALIAS` recomputes on read | +| identity column | `Int64` + no auto-generation | ClickHouse has no identity; the migrated values are what you keep | + +**Nullability:** `system.information_schema.columns.is_nullable` is +authoritative. Either declare `Nullable()` on the target, or declare +non-Nullable with an explicit `DEFAULT` **and** map `None` to that default +in a step-2 `transform=`. A non-Nullable target column with neither fails +mid-batch on the first NULL. + +--- + +## Delta / Unity Catalog Objects — Migration Patterns + +### Liquid clustering (`CLUSTER BY`) +Not an index and not a sort order you can copy. Treat it as a hint about +which columns the workload filters on, then choose the ClickHouse +`ORDER BY` from the actual queries. Say in chat when your choice differs +from the source clustering and why. + +### Deletion vectors +Deletes are recorded as vectors rather than rewritten files, so a table can +report rows that a `SELECT` won't return. Two consequences: source +`COUNT(*)` is the only trustworthy count (never sum file statistics), and a +table being actively deleted from can legitimately change count mid-run. + +### Time travel (`VERSION AS OF` / `TIMESTAMP AS OF`) +No ClickHouse equivalent. Use it to pin a consistent read for the +migration itself (`SELECT * FROM t VERSION AS OF `) so a concurrent +write doesn't skew validation. Do not try to reproduce the history. + +### Materialized views and streaming tables +Recreate as a ClickHouse Materialized View over an `AggregatingMergeTree`, +and **backfill** it from the base table after loading — an MV only sees +rows inserted after it exists. Get the defining query from +`describe_table`'s extended output. + +### Change Data Feed (`delta.enableChangeDataFeed`) +Ongoing CDC is out of scope for a one-shot migration. If the partner needs +it, the answer is ClickPipes or a `ReplacingMergeTree` with a version +column — say so and move on rather than half-building it. + +### Unity Catalog volumes and external locations +Only relevant to the S3 staging path. A `/Volumes/...` path is not +readable by ClickHouse; the staged unload writes to `s3://` directly. + +--- + +## Migration Script Rules + +- Connect with `DatabricksSource.from_env()`. Never construct + `databricks.sql.connect(...)` by hand in a generated script — the + helper owns namespace splitting and host normalisation. +- Fully qualify every `source_query=`. +- Row dict keys are **lower-case**; `iter_batches` lowercases them. +- `batch_size` ~100k narrow, ~25–50k for rows carrying VARIANT or + `ARRAY`. Never above 500k. +- **S3 staging requires a Unity Catalog external location** over the + staging bucket with `WRITE FILES` granted. `unload_to_s3` raises a + message saying exactly that; fall back to `add_table()` per table. +- **Target column order must match the source** for staged tables — that + path is `INSERT INTO … SELECT * FROM s3(...)`, which is positional. +- The staged path supports neither `batch_size` nor `transform`. Any + table needing per-row transformation uses the direct path. + +--- + +## Query Rewriting Notes + +| Databricks | ClickHouse | +|---|---| +| `QUALIFY ` | subquery + `WHERE` over the window | +| `LATERAL VIEW explode(arr) AS x` | `ARRAY JOIN` / `arrayJoin(arr)` | +| `transform(arr, x -> f(x))` | `arrayMap(x -> f(x), arr)` (order flips) | +| `filter(arr, x -> p(x))` | `arrayFilter(x -> p(x), arr)` (order flips) | +| `aggregate(arr, 0, (a, x) -> a + x)` | `arraySum(arr)` / `arrayReduce('sum', arr)` | +| `v:a.b` / `variant_get(v, '$.a')` | `JSONExtract*(v, 'a')`, or `v.a` on a `JSON` column | +| `named_struct('a', 1)` | `tuple(1)` — positional | +| `try_divide(a, b)` | `if(b = 0, NULL, a / b)` — `/` returns `inf`, not NULL | +| `try_cast(x AS t)` | `accurateCastOrNull(x, 't')` | +| `datediff(a, b)` | `dateDiff('day', b, a)` — order flips | +| `date_trunc('month', d)` | `toStartOfMonth(d)` | +| `SEMI JOIN` / `ANTI JOIN` | `LEFT SEMI JOIN` / `LEFT ANTI JOIN` | +| `catalog.schema.table` | `database.table` | + +State `NULLS FIRST` / `NULLS LAST` explicitly — the engines' defaults differ. + +--- + +## Known Gotchas + +- **Cold warehouse.** A stopped SQL warehouse takes 30s+ to serve its + first statement. Never present that as query latency; run a throwaway + `SELECT 1` before benchmarking. +- **`samples` catalog is read-only.** Copy out of it, never into it. +- **`VARIANT` needs DBSQL 2024.35+ / DBR 15.3+.** On older runtimes the + column type simply doesn't exist. +- **`system.query.history` can lag** by minutes and may be disabled, so + `server_ms` is sometimes `None`. Report wall-clock as wall-clock. +- **Auto-stopped warehouse mid-migration.** A long direct-path migration + keeps the connection busy, but a paused run can let the warehouse stop; + resuming pays cold-start again. Expected, not a bug. diff --git a/scripts/build-librechat-runtime.sh b/scripts/build-librechat-runtime.sh index d2ead10..cf274e4 100755 --- a/scripts/build-librechat-runtime.sh +++ b/scripts/build-librechat-runtime.sh @@ -15,6 +15,7 @@ # Optional sources (gated by Compose profile -> only included when profile is active): # snowflake-source -> profile "snowflake" # bigquery-source -> profile "bigquery" +# databricks-source -> profile "databricks" # All other MCPs are unconditional. set -euo pipefail @@ -74,6 +75,14 @@ else removed+=("bigquery-source") fi +# databricks-source: MCP key "databricks-source", host is databricks-mcp +if is_active databricks; then + kept+=("databricks-source") +else + drop_mcp "databricks-source" "databricks-mcp" + removed+=("databricks-source") +fi + echo "Runtime librechat config: $RUNTIME_FILE" echo " Active profiles: ${profiles_csv:-}" if [ ${#kept[@]} -gt 0 ]; then diff --git a/scripts/reset-agent.sh b/scripts/reset-agent.sh index a801d17..3d27229 100755 --- a/scripts/reset-agent.sh +++ b/scripts/reset-agent.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# Delete the four pre-built migration agents from MongoDB and re-run +# Delete the five pre-built migration agents from MongoDB and re-run # librechat-init to recreate them. Use after changing AGENT_PROVIDER / # AGENT_MODEL (or any AGENT_PROVIDER_) in .env, since the init # container's idempotency check skips agents that already match. @@ -12,6 +12,7 @@ AGENTS=( "Snowflake → ClickHouse Cloud" "BigQuery → ClickHouse Cloud" "ClickHouse OSS → ClickHouse Cloud" + "Databricks → ClickHouse Cloud" ) echo "This will delete these agents from MongoDB:" @@ -27,7 +28,8 @@ docker compose exec -T mongodb mongosh "mongodb://mongodb:27017/LibreChat" --qui "Postgres → ClickHouse Cloud", "Snowflake → ClickHouse Cloud", "BigQuery → ClickHouse Cloud", - "ClickHouse OSS → ClickHouse Cloud" + "ClickHouse OSS → ClickHouse Cloud", + "Databricks → ClickHouse Cloud" ]} }) ' @@ -47,6 +49,9 @@ if [ -z "${COMPOSE_PROFILES:-}" ]; then if grep -q '^bigquery-source$' <<<"$running"; then detected="${detected:+$detected,}bigquery" fi + if grep -q '^databricks-mcp$' <<<"$running"; then + detected="${detected:+$detected,}databricks" + fi export COMPOSE_PROFILES="$detected" echo "Detected active profiles: ${COMPOSE_PROFILES:-}" fi diff --git a/sources/databricks/GUIDE.md b/sources/databricks/GUIDE.md new file mode 100644 index 0000000..27e64f7 --- /dev/null +++ b/sources/databricks/GUIDE.md @@ -0,0 +1,432 @@ +# Migration Guide — Databricks → ClickHouse Cloud + +This guide walks you through a complete Databricks → ClickHouse Cloud +migration using the **MigrationRoom** dashboard. The dashboard +orchestrates the work: you pick a source, click six step buttons in +order, and watch the AI agent do each step live. The agent has MCP +connections to your Databricks source (`databricks-mcp`, host port +`8008`, Compose profile `databricks`), your ClickHouse Cloud target, +an in-chat Python runtime, and the `migrationkit` Python library that +handles data movement. + +**Demo workload:** `migration_demo.tpch` — TPC-H tables downsampled +from the `samples.tpch` catalog every Databricks workspace ships with. +`samples.tpch` is scale factor 1000 (~1 TB, ~6 billion `lineitem` +rows); `setup_workload.sql`'s predicates cut it to SF1-equivalent +cardinality (~6M `lineitem` rows), so this runs in minutes and the +eventual ClickHouse Cloud migration moves megabytes, not terabytes. +Eight Databricks-specific augmentations are layered on top so the agent +has to make real decisions instead of a mechanical type-for-type copy: + +| Source object | ClickHouse decision | +|---|---| +| `orders.o_metadata` (VARIANT) | `JSON` column, or extract hot keys to typed columns | +| `lineitem.l_shipping_events` (`ARRAY`) | `Nested(...)`, or `Array(Tuple(...))` | +| `lineitem.l_attributes` (`MAP`) | `Map(String, String)` | +| `lineitem.l_committed_at` / `l_committed_at_ntz` (`TIMESTAMP` / `TIMESTAMP_NTZ`) | `DateTime64(6, 'UTC')` / `DateTime64(6)` | +| `lineitem` liquid clustering (`CLUSTER BY`) | `ORDER BY (...)` chosen from the actual queries | +| `lineitem` deletion vectors | No equivalent; `ReplacingMergeTree`, ClickPipes, or defer | +| `orders.o_orderyear` (`GENERATED ALWAYS AS`) | `MATERIALIZED` (stored) or `ALIAS` (computed) column | +| `daily_order_summary` (materialized view, serverless-only) | ClickHouse Materialized View on `AggregatingMergeTree` | + +**Total time:** ~60 minutes including setup (estimate — see the +Honesty section below; this flow has not been timed end-to-end +against a live workspace). +**Workflow:** the dashboard's six step buttons drive the migration. The +prompt files in [prompts/](prompts/) are what each button fires — you +don't need to paste them by hand. + +--- + +## Phase 0 — Databricks setup (~5–20 min depending on path, estimate) + +Pick the entry point that matches what you already have. All three end +with the same `migration_demo.tpch` workload sitting in a Databricks +workspace and the four `DATABRICKS_*` variables in `.env`. + +### New Databricks environment + +You have no Databricks workspace at all. Terraform provisions a +**serverless** workspace, then chains straight into the demo-object +module below — one command, no copy-pasting a workspace URL between +two applies. + +The one manual step: create a Databricks **account** (if you don't +have one) and, in the account console +(`accounts.cloud.databricks.com` → Settings → Identity and access → +Service principals), an **account-admin** service principal. Terraform +authenticates *as* this service principal, so it cannot also create +it. Generate an OAuth secret and note the client ID and account UUID. + +```bash +cd sources/databricks/terraform/workspace +cp terraform.tfvars.example terraform.tfvars +# Edit: databricks_account_id, databricks_client_id, databricks_client_secret + +cd ../../../.. +make databricks-provision-workspace +``` + +`make databricks-provision-workspace` runs the `workspace` module, +then the `demo` module against the workspace it just created, then +prints where to capture the `.env` block: + +```bash +cd sources/databricks/terraform/demo && terraform output -raw env_block >> ../../../../.env +``` + +If the region has no Unity Catalog metastore yet (rare, but possible +for a brand-new account), the chained `demo` apply fails outright +because there's nothing for the new workspace to attach to. Set +`create_metastore = true` in `terraform/workspace/terraform.tfvars` +and re-run — see +[terraform/workspace/README.md](terraform/workspace/README.md#create_metastore) +for how to tell in advance whether you need it. This is the most +likely first failure of this "one command" path. + +See [terraform/workspace/README.md](terraform/workspace/README.md) +and [terraform/demo/README.md](terraform/demo/README.md) for what each +module creates and why account-admin (not workspace-admin) is required. + +### Existing workspace, provision the demo objects + +You already have a Databricks workspace with Unity Catalog. Terraform +creates the catalog/schema, a dedicated serverless SQL warehouse, a +demo service principal + token, and (optionally) the S3 staging path, +then seeds the workload for you. + +```bash +cd sources/databricks/terraform/demo +cp terraform.tfvars.example terraform.tfvars +# Edit: workspace_url, databricks_token +# (this module creates a catalog, and CREATE CATALOG is a metastore-level +# privilege, not a workspace-level one — see the README's Prerequisites +# for what the PAT's principal needs) + +cd ../../../.. +make databricks-provision +``` + +```bash +cd sources/databricks/terraform/demo && terraform output -raw env_block >> ../../../../.env +``` + +See [terraform/demo/README.md](terraform/demo/README.md) — in +particular the **re-apply caveat** if you set `enable_s3_staging = +true`: IAM propagation can make the first `apply` fail on +`databricks_external_location.staging`; running `terraform apply` +again succeeds once IAM catches up. + +### Existing workspace, workload only (fully manual) + +You want to do everything by hand instead of running Terraform. This +enumerates every step the two modules above automate. + +This path needs **two identities**, mirroring the Terraform `demo` +module's design (`terraform/demo/main.tf`'s provisioner deliberately +runs the seeding script as the admin PAT, not as the demo principal's +read-only token): a **setup** principal with write access, used only +to provision and seed, and a **runtime** principal with read-only +access, whose token is the one that ends up in `.env`. + +1. Create a Unity Catalog **catalog** and **schema**. The namespace + must be **exactly** `migration_demo.tpch` — `setup_workload.sql` + hard-codes that name in all 25 of its statements, so any other + catalog/schema name fails at the first one. +2. Create a **serverless** SQL warehouse. Serverless is required for + the `daily_order_summary` materialized-view augmentation — a + classic warehouse works for everything else, but skips that one + object (see the Troubleshooting entry below). +3. Create a **personal access token for the setup principal** — the + identity you'll run step 7 as. It needs write access: `CREATE + CATALOG` at the **metastore** (not workspace) level, plus `CREATE + SCHEMA`, `CREATE TABLE`, and `MODIFY` (covers `ALTER TABLE`, + `UPDATE`, `DELETE`) on the catalog from step 1. An existing + workspace-admin/account-admin PAT works fine here too. +4. Grant `USE CATALOG`, `USE SCHEMA`, and `SELECT` on the demo catalog, + **and** on the built-in `samples` catalog — the workload + `CREATE TABLE ... AS SELECT` copies out of `samples.tpch` — to a + **separate, read-only principal**. This is the identity that ends + up in `.env` for the playground to actually run as. +5. *(Optional, only for the S3-staged migration path)* create a Unity + Catalog **external location** over your staging bucket with + `WRITE FILES` granted. +6. Put the **setup** principal's token from step 3 in `.env` as + `DATABRICKS_TOKEN` — temporarily, just for the next step: + ```bash + 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 + ``` +7. Seed the workload: + ```bash + make databricks-setup + ``` + `databricks-setup` installs `databricks-sql-connector`, then runs + `sources/databricks/scripts/setup_workload.py`, which executes + `setup_workload.sql` against your warehouse — the 8 TPC-H tables + plus the eight augmentations above. **Once this succeeds, replace + `DATABRICKS_TOKEN` in `.env` with the read-only principal's token + from step 4** before starting the playground in Phase 1 — the + runtime identity does not need, and should not have, write access. + +--- + +## Phase 1 — Launch the playground (~5 min) + +```bash +make up-databricks +``` + +`make up-databricks` regenerates `librechat.runtime.yaml` for the +`databricks` profile, pulls/builds images, and starts every service +including `databricks-mcp`. Default `make up` skips it — like +Snowflake and BigQuery, the Databricks MCP is profile-gated because it +needs account credentials in `.env`. + +Check `docker compose ps` — `databricks-mcp` should show `healthy` +alongside the rest of the stack. Its healthcheck only probes that the +SSE endpoint responds; the actual Databricks connection is opened +lazily per tool call, so a healthy container does **not** by itself +prove your credentials are valid (see Troubleshooting). + +Open **** (accept the self-signed cert) +and sign in (`admin@playground.local` / `playground`). You'll land on +the **MigrationRoom** dashboard with the chat panel on the right. + +In the **SETUP** card at the top: + +- **Source**: pick `Databricks`. The chat panel auto-switches to the + `Databricks → ClickHouse Cloud` agent. +- **Source database**: the `catalog.schema` pair, e.g. + `migration_demo.tpch`. +- **Queries**: open **Edit · N OLAP** and confirm the OLAP queries are + loaded (steps 1 and 4 use them, for schema design and rewrite). + +--- + +## Phase 2 — Run the migration (~45 min, estimate) + +The dashboard has **six step buttons** at the top of the **STEPS** +panel. Click each in order. All six are clickable at any time, so you +can re-fire a step (e.g. re-run validation after fixing the schema). + +### Step 1 — Discover & Design Schema + +Agent introspects the source via the `databricks-source` MCP +(`list_catalogs`, `list_schemas`, `list_tables`, `describe_table`, +`run_select_query` — never `run_python` for discovery), reads the OLAP +queries to drive `ORDER BY` choices, proposes the ClickHouse target +schema, and runs the DDL via `clickhousectl`. + +**Watch in chat** for the agent's decisions: `Decimal(15, 2)` for +money, `Date32` for dates, `DateTime64(6, 'UTC')` vs `DateTime64(6)` +for the `TIMESTAMP` / `TIMESTAMP_NTZ` pair, `JSON` (not `String`) for +`o_metadata`, `Nested(...)` vs `Array(Tuple(...))` for +`l_shipping_events`, and an explicit call on `MATERIALIZED` vs `ALIAS` +for `o_orderyear`. **Confirm the target database name** when the agent +asks (default suggestion `migration_demo`). + +### Step 2 — Migrate Data + +Agent writes a short Python script using the `migrationkit` library, +dispatches it as a background job via `migration-runner`, issues +ONE `tail_python_job` to confirm `status=running`, then stops. The +dashboard's **Migration** tab streams live progress: rows/sec, ETA, +per-table progress bars, milestone events. + +Large tables (over ~1M rows) can also take the S3-stage path if +`STAGING_S3_*` is set — the agent picks direct vs staged per table. + +**This step is meant to look quiet in chat.** The agent dispatches and +stops by design (see [`docs/adding-a-source.md`](../../docs/adding-a-source.md) +on why polling isn't used); silence in the conversation while the +dashboard streams progress bars is expected, not stuck. + +### Step 3 — Validate + +Agent runs `Validator(...).validate()` — row count parity per table, +source vs target. Results land on the dashboard's **Validation** tab. +If anything mismatches the agent **stops and reports** in chat — fix +the schema and re-fire step 2, don't ask the agent to patch the target +by hand. + +### Step 4 — Rewrite Queries + +Agent translates each OLAP query from Databricks SQL to ClickHouse SQL +**in chat**. No script — this is a reasoning step. Walk through each +rewrite, push back on unfamiliar substitutions (`QUALIFY` → subquery, +`LATERAL VIEW explode` → `ARRAY JOIN`, `aggregate`/`filter` → +`arrayReduce`/`arrayFilter`, etc.). + +### Step 5 — Benchmark + +Agent runs `Benchmarker(...).benchmark(queries=[...])` — each query on +source and target, server-side timing on both. Results land on the +**Benchmark** tab as `source_ms / target_ms / speedup` per query. + +Databricks timing comes from the SQL query-history REST API, with +`system.query.history` as a fallback; if neither is available +`server_ms` is `None` and wall-clock is shown instead — the agent +should say so rather than presenting wall time as server time. The +first query in a session also pays Databricks warehouse cold-start; +expect a throwaway `SELECT 1` before the real numbers. + +### Step 6 — Optimize + +Agent proposes ClickHouse-Cloud-specific optimizations for the +slowest queries: Materialized Views on `AggregatingMergeTree`, +Projections, codec adjustments. Iterate in chat — once you apply an +optimization, re-fire step 5 to confirm the speedup. + +--- + +## Validation + +Compare the agent's final state against: + +- **Schema:** [queries/expected_ch_schema.sql](queries/expected_ch_schema.sql) +- **Queries:** [queries/expected_ch_queries.sql](queries/expected_ch_queries.sql) +- **Checklist:** [../../docs/migration-checklist.md](../../docs/migration-checklist.md) + +Bit-for-bit identity isn't expected — what matters is that the agent +made defensible choices: `Decimal` (not `Float`) for money, `Date32` +for dates, `JSON` (not a bare `String`) for `o_metadata`, a +`Nested`/`Map` decision for the nested lineitem columns, and a clear +call on the generated column and the materialized view. + +**Row-count parity has one legitimate exception.** `lineitem` has +deletion vectors enabled, and `setup_workload.sql` deletes 500 rows +from it as part of seeding a real history to time-travel over. If a +partner (or another process) deletes more rows from the source table +**while step 2 is running**, the source count can legitimately drop +mid-migration — that's not a migration bug, and step 3's `Validator` +output should be read with that in mind rather than assumed to be +stale. + +--- + +## Teardown and cost + +Everything provisioned in Phase 0 bills against your own +Databricks/AWS account, not against MigrationRoom. + +```bash +cd sources/databricks/terraform/demo && terraform destroy +cd ../workspace && terraform destroy # only if you used the workspace path +``` + +Destroy `demo` before `workspace` — the workspace module has no +`demo` in scope and destroying it first takes everything `demo` +created (catalog, warehouse, tokens) down with it regardless of order, +but destroying in `demo` → `workspace` order lets each module report +what it removed cleanly. + +- The serverless SQL warehouse's `auto_stop_minutes` defaults to 10 + minutes of idle time before it suspends — an idle warehouse still + bills, so don't raise this casually. +- If you enabled S3 staging, the staging bucket has a 7-day lifecycle + expiry on staged objects, and both the catalog/schema and the bucket + are created with `force_destroy = true` so `terraform destroy` won't + wedge on leftover demo tables or objects. + +--- + +## Troubleshooting + +**`databricks-mcp` shows healthy but every tool call errors:** +Credentials are checked lazily, per call — not at container startup — +so a healthy container proves the SSE endpoint responds, not that +`DATABRICKS_HOST` / `DATABRICKS_HTTP_PATH` / `DATABRICKS_TOKEN` are +valid. Check `docker compose logs databricks-mcp`, fix `.env`, then +`docker compose --profile databricks restart databricks-mcp`. + +**Chained `demo` apply fails because the workspace has no metastore:** +Some regions/accounts have no Unity Catalog metastore auto-provisioned +yet. Set `create_metastore = true` in +`terraform/workspace/terraform.tfvars` and re-apply — see +[terraform/workspace/README.md](terraform/workspace/README.md#create_metastore) +for how to check in advance whether you need it. This is the most +likely first failure of the "one command" new-environment path. + +**`databricks_grants.samples` fails on the first `terraform apply`:** +The built-in `samples` catalog is often owned by a different metastore +admin than the identity running Terraform. If the grant fails, either +apply again after getting the metastore admin to grant ownership (or +`SELECT`/`USE CATALOG`/`USE SCHEMA` on `samples.tpch` directly) to your +provisioning principal, or run the grant by hand once and re-apply — +this is a per-account permission quirk, not a bug in the module. + +**`setup_workload.sql` fails on the hand-written `orders` table (`CREATE +OR REPLACE TABLE migration_demo.tpch.orders`):** +That DDL's column types were written without a live workspace to test +against and may not exactly match `samples.tpch.orders` on yours. Run +`DESCRIBE TABLE samples.tpch.orders`, compare against the column list +in `setup_workload.sql`, adjust the mismatched types, then re-run +`make databricks-setup`. + +**First query of the session is slow:** +A stopped SQL warehouse takes 30s+ to serve its first statement. This +is warehouse cold-start, not a ClickHouse-vs-Databricks comparison — +run a throwaway `SELECT 1` before benchmarking, or note that query 1's +number includes cold-start. + +**`VARIANT` column rejected / type doesn't exist:** +Needs DBSQL 2024.35+ or DBR 15.3+. Upgrade the SQL warehouse's channel +to Current, or use a newer runtime, then re-run `make databricks-setup`. + +**Staged unload refused:** +The S3 path needs a Unity Catalog **external location** over the +staging bucket with `WRITE FILES` granted. `unload_to_s3` raises an +error saying exactly that; either provision the external location +(`enable_s3_staging=true` in the `demo` module) or let the agent fall +back to the direct path for that table. + +**`terraform apply` on the `demo` module fails once on +`databricks_external_location.staging`, then succeeds on retry:** +Expected, not a bug. The external location validates itself by +assuming the freshly created IAM role, and IAM is eventually +consistent — a role that was just created can fail to assume for a +few seconds even after Terraform reports the attachment done. Just +run `terraform apply` again; don't delete the role or the storage +credential. + +**`daily_order_summary` (and sample query 7) missing:** +The materialized-view augmentation is `@optional` in +`setup_workload.sql` — it needs a serverless SQL warehouse, and is +silently skipped on a classic one, which is a perfectly reasonable +demo environment otherwise. If it's absent, drop sample query 7 (and +its rewrite in `expected_ch_queries.sql`) rather than treating the +setup as broken. + +**Serverless egress control blocking the staging bucket:** +If the workspace has a network connectivity configuration (NCC) +restricting serverless egress, the staging bucket's endpoint must be +added to the allowed list, or the serverless warehouse can't read or +write staged Parquet through the external location. This is a +workspace-level network setting outside Terraform's control here — +check with whoever manages the workspace's NCC. + +--- + +## Honesty about what's been verified + +No live migration has been run end-to-end against a real Databricks +workspace as part of building this source — that needs a workspace, a +SQL warehouse, and a token, none of which exist in the environment +this guide was written in. `terraform apply` has **not** been run for +either Terraform module; `terraform validate` proves the configuration +is well-formed, not that an apply succeeds. See each module's README +(["workspace"](terraform/workspace/README.md), +["demo"](terraform/demo/README.md)) for the specific attributes most +likely to need adjustment on first contact with a real account. + +Likewise, `make up-databricks` — Phase 1's very first command — has +not actually been run in this environment: the `databricks-mcp` image +has not been built, and none of the Compose services have been booted +under the `databricks` profile. What's been verified is static: the +Dockerfile, compose service definition, and profile gating are +consistent with the other sources' equivalents, and `docker compose +config` accepts them. Whether the image builds cleanly and the +container reaches `healthy` has not been exercised. diff --git a/sources/databricks/manifest.json b/sources/databricks/manifest.json new file mode 100644 index 0000000..7c8116e --- /dev/null +++ b/sources/databricks/manifest.json @@ -0,0 +1,6 @@ +{ + "label": "Databricks", + "default_database_env": "DATABRICKS_NAMESPACE", + "default_database_fallback": "migration_demo.tpch", + "agent_name": "Databricks → ClickHouse Cloud" +} diff --git a/sources/databricks/prompts/01-discover-and-design.md b/sources/databricks/prompts/01-discover-and-design.md new file mode 100644 index 0000000..30a4765 --- /dev/null +++ b/sources/databricks/prompts/01-discover-and-design.md @@ -0,0 +1,88 @@ +# Step 1 — Discover the source and design the ClickHouse Cloud target schema + +You are migrating from `{source}` to ClickHouse Cloud. + +- **Source namespace** (where the data lives today): `{database}` — a + Unity Catalog `catalog.schema` pair, selected by the partner in the + dashboard. Use it as-is. +- **Target database** (where the data will land in ClickHouse Cloud): not + chosen yet. Propose a name in this step and confirm with the partner. + +If the partner has told you in this conversation to use a different +namespace, follow their chat instruction instead. + +## Source + +Use the `databricks-source` MCP — **not** `run_python`. Its five tools are +`list_catalogs`, `list_schemas`, `list_tables`, `describe_table`, and +`run_select_query`. + +1. `list_tables(catalog, schema)` for the namespace above. It returns + `sizeInBytes` and `numFiles` per table but **no row counts** — Delta + metadata doesn't carry them. +2. Get row counts in ONE query rather than one per table: + ```sql + SELECT 'orders' AS t, count(*) AS n FROM migration_demo.tpch.orders + UNION ALL SELECT 'lineitem', count(*) FROM migration_demo.tpch.lineitem + -- … one line per table + ORDER BY n DESC + ``` +3. `describe_table(catalog, schema, table)` for every table. Read the + `detail` and `history` sections too, not just the columns — that is + where clustering columns, partition columns, table features, and + deletion-vector state live. +4. Inventory the Databricks-specific features you actually find. Do not + assume any are present: VARIANT columns, `ARRAY` / `MAP` / + `STRUCT`, generated columns, liquid clustering (`CLUSTER BY`), + deletion vectors, `TIMESTAMP` vs `TIMESTAMP_NTZ`, materialized views, + streaming tables. +5. Sample rows and check cardinality before designing types: + ```sql + SELECT * FROM ..
LIMIT 5 + SELECT count(*), count(), count(DISTINCT ) FROM ..
+ ``` +6. Identify fact vs dimension tables and the join graph. + +## Analytical workload + +The partner will run these against the migrated data. Use them to choose +`ORDER BY` keys, partitioning, and projections — the ordering should come +from the columns in WHERE / JOIN / GROUP BY here, **not** from the source's +clustering columns: + +```sql +{olap_queries} +``` + +## Target + +Use the `clickhousectl` MCP to: + +1. Create the target database (suggested default `migration_demo`; confirm + first). +2. `CREATE TABLE` for every source table, following the ClickHouse Cloud + best-practice rules attached to **clickhousectl**. Justify each engine, + `ORDER BY`, `PARTITION BY`, and codec choice in chat. +3. Map Databricks types — the full table is in your Databricks source + instructions. The decisions worth surfacing to the partner: + - `VARIANT` → `JSON`, or extract hot keys into typed columns + - `ARRAY>` → `Nested(...)` or `Array(Tuple(...))` + - `MAP` → `Map(String, String)` + - `DECIMAL(p, s)` → `Decimal(p, s)`, never `Float64` + - `TIMESTAMP` → `DateTime64(6, 'UTC')`; `TIMESTAMP_NTZ` → `DateTime64(6)` + - generated column → `MATERIALIZED` (stored) or `ALIAS` (computed) +4. **Column order must match the source table's column order** for every + table you plan to migrate through S3 staging in step 2 — that path does + `INSERT INTO … SELECT * FROM s3(...)`, which is positional. +5. **Handle nullable columns deliberately.** `describe_table` reports + nullability. For each nullable column either declare + `Nullable()`, or declare it non-Nullable with an explicit `DEFAULT` + AND add a `transform=` lambda in step 2 mapping `None` to that default. + A non-Nullable column with neither will fail mid-batch on the first NULL. +6. Verify with `SHOW TABLES`. + +## When you're done + +Summarise the source namespace, target database name, and the key schema +decisions in chat — later steps refer back to them. Do **not** insert any +data; that is step 2. diff --git a/sources/databricks/prompts/02-migrate-data.md b/sources/databricks/prompts/02-migrate-data.md new file mode 100644 index 0000000..3a69d19 --- /dev/null +++ b/sources/databricks/prompts/02-migrate-data.md @@ -0,0 +1,115 @@ +# Step 2 — Migrate data using `migrationkit` + +The target schema from step 1 is in place. Copy the data from Databricks +into ClickHouse Cloud with the `migrationkit` library — it handles +batching, per-batch checkpointing, pause/resume/cancel, and the live +progress events the dashboard renders. + +## Pick a path per table + +| Path | Use when | API | +|---|---|---| +| **Direct** | `total_rows ≤ 1_000_000` | `m.add_table(...)` | +| **S3 staging** | `total_rows > 1_000_000` AND `STAGING_S3_BUCKET` is set | `m.add_table_via_s3(name=..., stage=S3Stage.from_env())` | + +```python +import os +USE_S3 = bool(os.environ.get("STAGING_S3_BUCKET")) +``` + +If `USE_S3` is False, use the direct path for every table and note in chat +that the partner hasn't configured S3 staging. Don't fail the migration — +direct works at any size, it's just slower. + +**The S3 path additionally needs a Unity Catalog external location** over +the staging bucket with `WRITE FILES` granted. If `unload_to_s3` raises, +the error says so; fall back to `add_table()` for that table and tell the +partner. + +## What to write + +One Python script (~25 lines), dispatched with `run_python_background`, +confirmed with ONE `tail_python_job` call. + +```python +import os +import time +from migrationkit import Migrator, DatabricksSource, ClickHouseTarget, S3Stage + +USE_S3 = bool(os.environ.get("STAGING_S3_BUCKET")) +stage = S3Stage.from_env() if USE_S3 else None + +m = Migrator( + run_id=f"migrate-databricks-{int(time.time())}", + source=DatabricksSource.from_env(), + target=ClickHouseTarget.from_env(), + # REQUIRED: the ClickHouse Cloud database from step 1. + target_database="", +) + +# Direct path: dimensions and small facts. `target_table` is a BARE name — +# never `db.table`; the Migrator owns the database via target_database=. +m.add_table( + name="", + source_query="SELECT * FROM ..", + target_table="", + batch_size=100_000, +) + +# S3-staged path: large facts, only when stage is set. +if stage is not None: + m.add_table_via_s3(name="", target_table="", stage=stage) +else: + m.add_table( + name="", + source_query="SELECT * FROM ..", + target_table="", + batch_size=50_000, + ) + +# … one m.add_table(...) or m.add_table_via_s3(...) per source table. + +m.run() +``` + +Chat-side flow: + +```text +1. call: write_workspace_file(path="migrate.py", content=