Add Databricks migration source - #8
Open
sishuo-yang wants to merge 21 commits into
Open
Conversation
Add keyword_after_ctes() helper to parse CTE definitions and extract the actual statement keyword. This prevents mutations like: WITH x AS (SELECT 1) INSERT INTO orders SELECT * FROM x The guard now verifies that statements starting with WITH have SELECT (not INSERT/UPDATE/DELETE/MERGE) as the actual command after the CTE list. Add 23 comprehensive tests covering comments, string literals, nested parens, RECURSIVE modifier, multiple CTEs, column lists, and VALUES form. All 52 tests pass (29 original + 23 new CTE tests).
Replace three separate quote/comment-handling regexes with a single
_neutralize() function that makes one pass through the input, tracking
mutually-exclusive quoting/commenting states:
- line comments: -- to end of line
- block comments: /* ... */
- single-quoted strings: '...' ('''' escapes)
- double-quoted strings: "..." ("" escapes)
- backtick identifiers: `...` (`` escapes)
This closes the bypass where backtick identifiers containing ) or keywords
could hide mutations. Example:
WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) INSERT INTO t...
The backtick closes the CTE early; with layered regexes this was invisible.
With unified state tracking, the paren and SELECT are neutralized; the
scanner correctly sees INSERT as the statement keyword and rejects it.
Benefits of unified normalization:
1. All quoting/commenting interactions are correct (states are mutually exclusive)
2. Single source of truth: _neutralize() is the ONLY place that understands quoting
3. Callers are simpler (no regex substitution chains)
4. No false rejections: 'a;b' no longer triggers multi-statement error
5. No double-LIMIT: LIMIT inside backtick doesn't confuse the check
Neutralization returns identical-length strings with quoted/commented
regions replaced by spaces, enabling offset-preserving analysis.
Strip_comments, is_multi_statement, and keyword_after_ctes all now use
_neutralize() instead of their own regexes. Deleted old regex patterns.
Add 18 new tests covering backtick security cases:
- Backtick hiding DML after CTE (INSERT, UPDATE, DELETE, MERGE)
- Backtick containing closing paren
- Backtick with semicolon (false rejection fixed)
- Backtick with -- and /* markers (not treated as comments)
- Escaped backtick (doubled backtick)
- String/comment marker precedence tests
- LIMIT inside backtick (real LIMIT still injected)
All 70 tests pass (52 previous + 18 new).
Three review rounds each found a real mutation bypass in the hand-rolled lexer (CTE-prefixed DML, a backtick identifier blinding the CTE scanner, and comment/string lexing that diverged from Spark's grammar). Delegate classification to sqlglot's Databricks dialect tokenizer/parser instead of re-implementing Spark's lexical grammar by hand: tokenize first to catch multi-statement input correctly (comments/strings/backticks/escapes all handled by a real tokenizer), then require the parsed statement to be a Select/Union/Subquery/Describe, or a Command/unparseable statement whose leading verb is one of SHOW/EXPLAIN/DESCRIBE/DESC. LIMIT injection now reads the AST's own limit clause instead of a regex, so a LIMIT-shaped alias/backtick text can no longer suppress the cap. The original statement text is returned untouched (only a trailing semicolon stripped and a LIMIT appended when needed), never a sqlglot-regenerated form. Removes READ_ONLY_LEADING_KEYWORDS and all the hand-rolled lexing helpers per the amended spec. Adds sqlglot>=30.0 to both requirements files.
Two regressions from the sqlglot rewrite, found by re-review: - A top-level `(SELECT ...)` parses to exp.Subquery, not exp.Select/ exp.Union, so it was exempt from LIMIT injection entirely -- silently defeating the row cap for parenthesized queries (and their UNION variants). Fix: include exp.Subquery in the LIMIT branch and check both the Subquery's own limit slot and the inner wrapped query's, since sqlglot puts the limit on whichever level `LIMIT` textually followed in the source (`(SELECT ... LIMIT 5)` vs `(SELECT ...) LIMIT 5`). Verified `(SELECT ...) LIMIT n` is valid Databricks SQL by round-tripping it back through sqlglot's own parser and generator. - sqlglot.parse() emits a synthetic exp.Semicolon node to hold a comment trailing the statement terminator (`SELECT 1; -- comment`); it survived the `if s` None-filter and made an ordinary trailing comment look like a second statement. Fix: filter exp.Semicolon alongside None. The multi-statement gate itself is untouched -- a real second statement after the comment is still rejected (interior semicolon is caught before parsing is even attempted). Left the Command-fallback trailing-text behavior (e.g. `SHOW TABLES\nDROP TABLE t`) as-is per review: it's a single malformed statement Spark's own parser would reject, and hardening it would mean re-introducing keyword scanning on Command argument text -- the exact pattern that caused three prior bypasses.
…alues _literal() escaped ' but not \, so a value ending in a backslash could break out of the string literal and inject additional SQL. Delete it and bind catalog/schema values with the connector's native ? parameters in list_schemas and list_tables instead. Identifiers are unaffected: _ident() was already sound and continues to build the fully-qualified names passed to DESCRIBE DETAIL/HISTORY/TABLE EXTENDED.
Adds DatabricksSource to migrationkit, implementing the Source ABC for Databricks SQL warehouses: direct batch reads via databricks-sql-connector plus an INSERT OVERWRITE DIRECTORY ... USING PARQUET path to S3 staging, mirroring SnowflakeSource's shape. server_ms degrades through a query history REST API tier and a system.query.history fallback tier, returning None (never raising) when both fail. Deviates from the task brief in one place: _server_ms_from_system_table uses the connector's native positional parameter binding (`?` placeholder) for statement_id instead of f-string interpolation, per Task 2's guidance to prefer binding over hand-rolled quoting wherever the API allows it. Wires DatabricksSource into migrationkit/sources/__init__.py and migrationkit/__init__.py exports, and into api.py's list_source_databases dispatch. Adds databricks-sql-connector>=4.0 to requirements.txt.
…3 error framing Review found three issues in DatabricksSource, all fixed: - _server_ms_from_history_api's try/except only guarded the network call and json.load; the body-processing loop after it ran unguarded and raised AttributeError/TypeError on any malformed response shape (the shape is an admitted guess against a live API). Widened the try to wrap the whole function and broadened to `except Exception`, so it now always returns float or None, matching the docstring's contract that Benchmarker can fall back to wall_ms. - unload_to_s3 evaluated self._fq(table) inside the try around cur.execute, so an unqualified table name (a caller usage error) got caught and re-wrapped as a misleading "Databricks refused... needs a Unity Catalog external location" RuntimeError. Moved _fq(table) before the try so its ValueError propagates unwrapped. - Added tests for _fq (bare-name qualification, already-qualified passthrough, ValueError with no namespace) and a 5-case parametrized regression test for the history-API fix, covering null/list/malformed response bodies via a fake urlopen response. Test count: 22 (up from 14), full suite 83 passed.
Adds setup_workload.sql (TPC-H copy from samples.tpch plus VARIANT, generated column, nested ARRAY<STRUCT>/MAP, TIMESTAMP_NTZ, liquid clustering, deletion vectors, and an optional materialized view), setup_workload.py (directive-aware runner with @requires/@optional exit-code semantics), and the parser's unit tests. Deviates from the task brief in one spot: the lineitem UPDATE now assigns CANCELLED to ~1/10 of rows (pmod(l_orderkey, 10) = 0) in the second shipping event, so sample query 6's `status <> 'CANCELLED'` filter is not a no-op against real data.
…00 source Task 5 review (ruled on by the human) found two defects in the original brief that this script inherited: 1. Databricks' ALTER TABLE ADD COLUMN grammar has no GENERATED clause, so the o_orderyear generated column can only be declared at CREATE TABLE time. Replaced the plain orders CTAS with an explicit CREATE TABLE (naming all columns plus the generated column) and a separate INSERT with an explicit column list, since the source has no o_orderyear to select. The @requires directive moved to the CREATE TABLE, with its hint rewritten to flag the real risk: hard-coded column types may not match samples.tpch.orders on a given workspace. 2. samples.tpch is scale factor 1000 (~1 TB), not SF1 (~6M rows) as the brief assumed. Added key-range/derived-key predicates to all 8 table population statements to downsample to SF1-equivalent cardinality: orders/lineitem share the same o_orderkey/l_orderkey <= 6000000 range so the fact-to-fact join holds by construction; customer is derived from the orders actually kept; region/nation stay whole copies; part/partsupp/supplier get cheap independent key-range slices. customer's CREATE TABLE moved to after orders since its predicate subqueries migration_demo.tpch.orders. Statement and directive counts are unchanged (25 statements; 5 required / 1 optional) — the deleted ALTER TABLE and the new INSERT net to zero, and the @requires directive relocated rather than disappearing.
Add the complete set of six markdown prompts fired by dashboard step buttons:
01-discover-and-design, 02-migrate-data, 03-validate, 04-rewrite-queries,
05-benchmark, 06-optimize. Each embeds the exact instruction flow for that
migration stage, with placeholders {source}, {database}, {olap_queries}
substituted client-side. No polling loops per repo rule; single tail_python_job
calls per background task.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two content defects from the brief: 1. Prompt 05 (benchmark): Fix benchmark() call shape. Was passing tuples (source_sql, target_sql) but benchmarker.py:93 expects list[dict] with keys "name", "source_sql", "target_sql". Update template and comments. 2. Prompt 04 (rewrite-queries): Fix ClickHouse SEMI/ANTI JOIN syntax in dialect table. Correct order is LEFT SEMI JOIN / LEFT ANTI JOIN, not SEMI LEFT JOIN / ANTI LEFT JOIN. Verified: no polling loops remain, placeholders intact, six filenames unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add the Databricks source migration-instructions markdown and register the databricks-source MCP server (and databricks-mcp allowedDomains entry) in librechat.yaml so build-instructions.sh can inject the rules into serverInstructions.
Adds the databricks-mcp Compose service (profile "databricks", host port 8008), wires it into librechat's depends_on and librechat-init's agent bootstrap loop (agent tuple + OPTIONAL_PROFILE), extends the runtime librechat.yaml profile-stripping script with a third gate, adds make up-databricks / databricks-setup targets (and reserves databricks-provision / databricks-provision-workspace for later tasks), and documents the four DATABRICKS_* runtime env vars in .env.example.
Provisions Unity Catalog namespace, serverless SQL warehouse, demo service principal + token, grants, optional S3 staging path, and the TPC-H workload inside a workspace the partner already has. Adds the databricks-provision Makefile recipe.
Adds the `workspace` module, which creates a serverless Databricks workspace from nothing (no IAM role, root bucket, or VPC needed) and grants the provisioning service principal workspace admin, then chains into the existing `demo` module. demo/ gains OAuth (client_id/client_secret) as an alternative to databricks_token, since a fresh workspace has no PAT yet, plus a short-lived provisioner token for the workload-setup script so it always runs with write privileges regardless of auth mode. The Makefile's merge step between the two `terraform apply` calls is a standalone script (sources/databricks/scripts/merge_workspace_tfvars.py) rather than an inline heredoc: a heredoc spanning multiple Makefile recipe lines only works under GNU Make's .ONESHELL (3.82+), and the `make` on this machine is 3.81, which runs each recipe line in its own shell and silently mis-splits a multi-line heredoc.
… CATALOG assumption merge_workspace_tfvars.py silently wrote an empty client_secret when a partner supplied it via TF_VAR_databricks_client_secret instead of the tfvars file (the script only reads the file); demo/'s null-coalescing provider block then turned that into an ambiguous OAuth failure far from the cause. It now exits non-zero naming the missing variable, states it must be in workspace/terraform.tfvars specifically, and explains the script doesn't see TF_VAR_* env vars — plus purpose-written messages for a missing input JSON or missing tfvars file instead of bare tracebacks. Also documents, in both READMEs, that the provisioning service principal needs account admin (not just workspace admin) because demo/'s CREATE CATALOG is a metastore-level privilege that only account-admin's implicit metastore-admin capability satisfies — no metastore grant was added, since guessing the metastore ID would be worse than stating the assumption.
… updates Ship the reference target schema and rewritten OLAP queries for the Databricks source (verified consistent with each other and with setup_workload.sql by script), a partner-facing GUIDE.md covering the three Phase 0 entry points, and sweep README.md/docs/adding-a-source.md/ docs/architecture.mmd for the fifth source and its MCP node.
Consolidated fix wave from the whole-branch review before proposing the Databricks source for merge: - reset-agent.sh: add Databricks to both agent lists and auto-detect the `databricks` Compose profile via the `databricks-mcp` service, so `make reset-agent` recreates the agent instead of only deleting it. - terraform/demo/variables.tf: constrain catalog_name/schema_name to their defaults with actionable validation errors, since setup_workload.sql hard-codes the migration_demo.tpch namespace. - GUIDE.md: fix the manual path's two-identity story (write-capable setup principal vs. read-only runtime principal, mirroring the Terraform module), pin the namespace to exactly migration_demo.tpch, add create_metastore / samples-grant / orders-DDL troubleshooting entries, correct several stale facts (augmentation count, PAT wording, auto_stop_minutes, olap_queries steps, restart command, what's actually been run), and demo/README.md: note the samples-grant caveat. - librechat.yaml: generalize migration-runner's serverInstructions beyond Snowflake and list databricks-sql-connector / DATABRICKS_* env vars. - Repo-wide staleness: port allocations, sources/ tree, instructions file globs, "four sources" / "three pre-built agents" counts, TPC-H augmentation tables, migrator.py/base.py docstrings, and stale Makefile/docker-compose.yml comments now all account for Databricks. - expected_ch_schema.sql: Decimal(18,2) to match the source precision, correct the augmented-column-order prose, and replace the unreachable "explicit column list" prescription with the actual fact that ClickHouse's s3()/file() match Parquet columns by name. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
databricksas a fifth migration source, at the same level of completeness as Snowflake and BigQuery: a partner points the playground at their Databricks workspace, picks Databricks in the dashboard's Source dropdown, and runs the six migration steps against a TPC-H workload decorated with Delta/Unity-Catalog features.Design doc:
docs/superpowers/specs/2026-08-06-databricks-source-design.mdWhat's here
docker/databricks-mcp/list_catalogs,list_schemas,list_tables,describe_table,run_select_query). Purpose-built because Databricks publishes no introspect-and-SELECT MCP for SQL warehouses — theirdatabricks-mcppackage is an OAuth helper for the hosted UC/vector-search/Genie servers.migrationkit/sources/databricks.pyDatabricksSource(Source)— direct batch reads plus an S3-staged path viaINSERT OVERWRITE DIRECTORY. MirrorsSnowflakeSource.sources/databricks/sources/databricks/terraform/workspace/creates a serverless workspace from nothing;demo/provisions the demo objects into an existing one.librechat/sources/databricks-instructions.mdbuild-instructions.sh.Three entry paths, so both audiences are served:
make databricks-provision-workspace. One account-console visit (create an account-admin service principal — the credential Terraform authenticates with, so it can't provision itself), then one command. Serverless workspaces need no cross-account IAM role, root bucket, or VPC.make databricks-provision.make databricks-setup, everything else configured by hand.Notable design decisions
The read-only guard is AST-based, not keyword-based.
docker/databricks-mcp/sql_guard.pyparses withsqlglot's Databricks dialect and judges the root node type. Three consecutive review rounds found real mutation bypasses in the original leading-keyword approach —WITH x AS (SELECT 1) INSERT INTO orders SELECT * FROM x(Databricks permits a CTE before DML), a backtick-quoted identifier smuggling a)past the CTE scanner, and a\r-terminated line comment hiding a; DROP(Spark's rule is~[\r\n], not~[\n]). Each fix was correct for the case reported and wrong elsewhere, because hand-matching Spark's lexical grammar is the task. With a real parser, a CTE-prefixedINSERTsimply is anexp.Insertnode. The rewrite was a net −301 lines.samples.tpchis scale factor 1000, not 1. Databricks documents it as "approximately 1 TB", besidetpcds_sf1at "approximately 1 GB". A plainSELECT *CTAS would have copied ~8.7 billion rows. The workload slicesordersandlineitemon the sameo_orderkey/l_orderkeyrange so their FK relationship holds by construction, and derivescustomerfrom the orders actually kept so theorders → customer → nationjoin stays exact. The cutoff is exact rather than approximate: TPC-H assigns order keys by the scale-independent formula(seq/8)*32 + (seq%8) + 1, which tops out at 5,999,976 for the last SF1 order.The generated column is declared at
CREATE TABLE, never added byALTER TABLE. Databricks'ALTER TABLE … ADD COLUMNgrammar has no generation-expression clause. Downgrading it to a plain computed column was rejected:generation_expressionwould beNULL, the agent would see an ordinaryINT, and theMATERIALIZED-vs-ALIASdecision the augmentation exists to force would silently vanish.Values are bound, identifiers are validated. The MCP passes
catalog/schemaas native parameters (Databricks documents connector 3.0.0+ parameters as injection-safe, the inline style as not). Identifiers can't be bound, soDESCRIBE DETAILuses a quoting helper that rejects any embedded backtick.Two Terraform modules, not one with a flag.
demo/'s provider needshost = <workspace URL>, which doesn't exist untilworkspace/is applied, and Terraform can't plan resources whose provider configuration depends on a resource created in the same apply.Verified
pytest tests/— 91 passing (new suite; the repo had none before).fmt -check,init -backend=false,validateagainst provider 1.124.0.librechat.yamlparses;build-instructions.shinjects idempotently; the Compose service resolves only under its profile; profile-strip works in both directions.JSONExtractStringover the nativeJSONtype,ARRAY JOINoverNested, theQUALIFY→subquery rewrite,countMerge/sumMergeover theAggregatingMergeTree.NOT verified — needs a reviewer with a workspace
Please treat these as open until someone exercises them:
make up-databricks, the first command a partner runs.terraform applyfor either module.validateproves the configuration is well-formed, not that an apply succeeds. The storage-credential ↔ IAM ordering (skip_validation, constructed ARN,time_sleep) and the metastore lookup can only be confirmed against a real account.DatabricksSource, each annotated in the code and each degrading to wall-clock timing rather than failing:cursor.query_id, the SQL query-history REST response shape, andsystem.query.historycolumn names.ordersDDL types insetup_workload.sql—o_totalprice'sDECIMALprecision is a guess. The statement is@requires-labelled with a hint telling the operator to runDESCRIBE TABLE samples.tpch.ordersand adjust.system.information_schemais metastore-wide. Three MCP tools and the source-database dropdown assume it spans the metastore rather than thesystemcatalog alone. One live query settles it; if it's wrong, discovery returns nothing.Each of these is also stated in the shipped docs —
GUIDE.mdand both module READMEs carry explicit "what has and hasn't been verified" sections.Known follow-ups (deliberately not in this PR)
migrator.py:22's module-level docstring still says "(Snowflake + ClickHouse OSS today)" — the method docstring three lines away was corrected.GUIDE.md:223saysDecimal(15, 2)for money where the reference schema now saysDecimal(18, 2).demo/'scatalog_name/schema_nameare pinned to their defaults byvalidationblocks, becausesetup_workload.sqlhard-codesmigration_demo.tpch. Templating the SQL is the real fix; constraining was the right call at a merge gate.SetupRail.tsx's analytical-query count is wrong for every source (BigQuery shows 1 of 6, Snowflake 1 of 5) — asplit(";")filter that drops any query preceded by a comment. Pre-existing; dashboard changes were a spec non-goal.Review notes
The commit history is deliberately legible: each
feat(...)is a task, eachfix(...)is a review round. Every defect found during implementation originated in the spec or plan rather than the implementation, and the spec records each amendment with its rationale and rejected alternatives.🤖 Generated with Claude Code