From 5a8c836ce0d4838a87bffa00f979a79cbd526baa Mon Sep 17 00:00:00 2001 From: sishuo Date: Thu, 6 Aug 2026 17:19:01 +0700 Subject: [PATCH 01/21] feat(databricks): read-only SQL guard for the source MCP --- docker/databricks-mcp/sql_guard.py | 78 ++++++++++++++++++++++ tests/requirements.txt | 1 + tests/test_sql_guard.py | 100 +++++++++++++++++++++++++++++ 3 files changed, 179 insertions(+) create mode 100644 docker/databricks-mcp/sql_guard.py create mode 100644 tests/requirements.txt create mode 100644 tests/test_sql_guard.py diff --git a/docker/databricks-mcp/sql_guard.py b/docker/databricks-mcp/sql_guard.py new file mode 100644 index 0000000..09f5cac --- /dev/null +++ b/docker/databricks-mcp/sql_guard.py @@ -0,0 +1,78 @@ +"""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. +""" +from __future__ import annotations + +import re + +READ_ONLY_LEADING_KEYWORDS = frozenset( + {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"} +) + +_LINE_COMMENT = re.compile(r"--[^\n]*") +_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.S) +_SINGLE_QUOTED = re.compile(r"'(?:''|[^'])*'") +_DOUBLE_QUOTED = re.compile(r'"(?:""|[^"])*"') +_LEADING_WORD = re.compile(r"[(\s]*([A-Za-z_]+)") +_TRAILING_LIMIT = re.compile(r"\blimit\b\s+\d+\s*$", re.I) + + +class SqlNotAllowed(ValueError): + """Raised when a statement is not a single read-only statement.""" + + +def strip_comments(sql: str) -> str: + """Remove block and line comments so they can't hide a leading keyword.""" + return _LINE_COMMENT.sub(" ", _BLOCK_COMMENT.sub(" ", sql)) + + +def leading_keyword(sql: str) -> str: + """First bare word of the statement, upper-cased. Leading parens and + comments are skipped so `(SELECT 1)` and `/* c */ SELECT 1` both read + as SELECT.""" + match = _LEADING_WORD.match(strip_comments(sql).strip()) + return match.group(1).upper() if match else "" + + +def is_multi_statement(sql: str) -> bool: + """True when a ';' separates statements. Semicolons inside string + literals and comments don't count, and a single trailing ';' is fine.""" + body = strip_comments(sql) + body = _DOUBLE_QUOTED.sub('""', _SINGLE_QUOTED.sub("''", body)) + return ";" in body.strip().rstrip(";") + + +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/WITH statements that don't already end + in one, so an unbounded scan can't stream a whole fact table into the + chat context. + + Raises SqlNotAllowed for anything else.""" + if not sql or not sql.strip(): + raise SqlNotAllowed("empty statement") + if is_multi_statement(sql): + raise SqlNotAllowed( + "multiple statements are not allowed; send one statement" + ) + keyword = leading_keyword(sql) + if keyword not in READ_ONLY_LEADING_KEYWORDS: + allowed = ", ".join(sorted(READ_ONLY_LEADING_KEYWORDS)) + raise SqlNotAllowed( + f"this MCP is read-only; only {allowed} are permitted, " + f"got {keyword or '?'}. Run DDL against the target via " + f"the clickhousectl MCP instead." + ) + body = sql.strip().rstrip(";").rstrip() + if keyword in {"SELECT", "WITH"} and not _TRAILING_LIMIT.search( + strip_comments(body) + ): + body = f"{body}\nLIMIT {max_rows}" + return body diff --git a/tests/requirements.txt b/tests/requirements.txt new file mode 100644 index 0000000..039d26e --- /dev/null +++ b/tests/requirements.txt @@ -0,0 +1 @@ +pytest>=8.0 diff --git a/tests/test_sql_guard.py b/tests/test_sql_guard.py new file mode 100644 index 0000000..1ecbf92 --- /dev/null +++ b/tests/test_sql_guard.py @@ -0,0 +1,100 @@ +"""Unit tests for the Databricks MCP read-only statement guard. + +Run from the repo root: + python3 -m pip install -r tests/requirements.txt + python3 -m pytest tests/test_sql_guard.py -v +""" +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "docker" / "databricks-mcp")) + +from sql_guard import READ_ONLY_LEADING_KEYWORDS, SqlNotAllowed, guard # noqa: E402 + + +@pytest.mark.parametrize( + "sql", + [ + "SELECT 1", + "select c_custkey from customer", + "WITH x AS (SELECT 1) SELECT * FROM x", + "SHOW CATALOGS", + "DESCRIBE TABLE migration_demo.tpch.orders", + "DESC DETAIL migration_demo.tpch.lineitem", + "EXPLAIN SELECT 1", + " \n SELECT 1 \n ", + "(SELECT 1)", + ], +) +def test_read_only_statements_are_allowed(sql): + assert guard(sql) + + +@pytest.mark.parametrize( + "sql", + [ + "INSERT INTO t VALUES (1)", + "UPDATE t SET a = 1", + "DELETE FROM t", + "DROP TABLE t", + "CREATE TABLE t (a INT)", + "ALTER TABLE t ADD COLUMN b INT", + "MERGE INTO t USING s ON t.a = s.a", + "TRUNCATE TABLE t", + "COPY INTO t FROM 's3://b/k'", + "GRANT SELECT ON TABLE t TO `u`", + ], +) +def test_mutating_statements_are_rejected(sql): + with pytest.raises(SqlNotAllowed): + guard(sql) + + +def test_multi_statement_is_rejected(): + with pytest.raises(SqlNotAllowed) as excinfo: + guard("SELECT 1; DROP TABLE t") + assert "one statement" in str(excinfo.value) + + +def test_trailing_semicolon_is_not_multi_statement(): + assert guard("SELECT 1;").startswith("SELECT 1") + + +def test_semicolon_inside_string_literal_is_not_multi_statement(): + assert guard("SELECT 'a;b' AS s") + + +def test_comment_hidden_mutation_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("-- harmless\nDROP TABLE t") + + +def test_block_comment_hidden_mutation_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("/* SELECT */ DROP TABLE t") + + +def test_limit_is_injected_when_absent(): + assert guard("SELECT * FROM orders", max_rows=50).endswith("LIMIT 50") + + +def test_existing_limit_is_not_doubled(): + out = guard("SELECT * FROM orders LIMIT 5", max_rows=50) + assert out.lower().count("limit") == 1 + + +def test_show_does_not_get_a_limit(): + assert "LIMIT" not in guard("SHOW CATALOGS") + + +def test_empty_statement_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard(" ") + + +def test_keyword_set_is_read_only(): + assert READ_ONLY_LEADING_KEYWORDS == frozenset( + {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"} + ) From 3925a9738d9245d97dcc00b6ac4234d1d2148908 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 14:23:59 +0800 Subject: [PATCH 02/21] fix(databricks): detect and reject CTE-prefixed DML mutations 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). --- docker/databricks-mcp/sql_guard.py | 136 +++++++++++++++++++++++++++++ tests/test_sql_guard.py | 112 ++++++++++++++++++++++++ 2 files changed, 248 insertions(+) diff --git a/docker/databricks-mcp/sql_guard.py b/docker/databricks-mcp/sql_guard.py index 09f5cac..1388604 100644 --- a/docker/databricks-mcp/sql_guard.py +++ b/docker/databricks-mcp/sql_guard.py @@ -48,6 +48,131 @@ def is_multi_statement(sql: str) -> bool: return ";" in body.strip().rstrip(";") +def keyword_after_ctes(sql: str) -> str: + """Extract the top-level SQL keyword following CTE definitions. + + When a statement starts with WITH, this returns the keyword of the actual + statement following the CTE list. For example: + - "WITH x AS (SELECT 1) INSERT INTO ..." returns "INSERT" + - "WITH x AS (SELECT 1) SELECT * FROM x" returns "SELECT" + + Handles comments, string literals, RECURSIVE modifier, multiple CTEs, + and column lists. Uses a tokenizer to safely parse paren-balanced CTEs. + """ + # Strip comments + body = strip_comments(sql).strip() + + # Replace string literals with placeholders to avoid parsing their contents + strings = {} + counter = [0] + + def replace_string(match): + placeholder = f"__STR_{counter[0]}__" + strings[placeholder] = match.group(0) + counter[0] += 1 + return placeholder + + body = _SINGLE_QUOTED.sub(replace_string, body) + body = _DOUBLE_QUOTED.sub(replace_string, body) + + # Tokenize: extract words and structural tokens (parens, commas) + tokens = [] + i = 0 + while i < len(body): + # Skip whitespace + while i < len(body) and body[i] in " \t\n\r": + i += 1 + if i >= len(body): + break + + # Extract words: letters, underscores, and digits + if body[i].isalpha() or body[i] == "_": + start = i + while i < len(body) and (body[i].isalnum() or body[i] == "_"): + i += 1 + tokens.append(body[start:i].upper()) + # Extract structural tokens + elif body[i] in "(),": + tokens.append(body[i]) + i += 1 + else: + # Skip other characters (operators, etc.) + i += 1 + + # Verify it starts with WITH + if not tokens or tokens[0] != "WITH": + return "" + + pos = 1 + + # Skip optional RECURSIVE keyword + if pos < len(tokens) and tokens[pos] == "RECURSIVE": + pos += 1 + + # Parse the CTE list: each CTE is name [(cols)] AS (body) + # Multiple CTEs are separated by commas + while pos < len(tokens): + # Pattern: NAME [ ( ... ) ] AS ( ... ) [, more CTEs] + + # Expect a name + if pos >= len(tokens) or not tokens[pos][0].isalpha(): + break + pos += 1 + + # Skip optional column list (parens before AS) + if pos < len(tokens) and tokens[pos] == "(": + # Look ahead: if next closing ) is followed by AS, this is a column list + save_pos = pos + 1 + depth = 1 + while save_pos < len(tokens) and depth > 0: + if tokens[save_pos] == "(": + depth += 1 + elif tokens[save_pos] == ")": + depth -= 1 + save_pos += 1 + + # Check if AS follows the closing paren + if save_pos < len(tokens) and tokens[save_pos] == "AS": + # Yes, it's a column list - skip it + pos = save_pos + # If not AS, this paren must be the CTE body, so don't skip + + # Expect AS + if pos >= len(tokens) or tokens[pos] != "AS": + # Malformed; return what we have + if pos < len(tokens): + return tokens[pos] + return "" + pos += 1 + + # Expect and skip balanced parens (CTE body) + if pos >= len(tokens) or tokens[pos] != "(": + if pos < len(tokens): + return tokens[pos] + return "" + + depth = 1 + pos += 1 + while pos < len(tokens) and depth > 0: + if tokens[pos] == "(": + depth += 1 + elif tokens[pos] == ")": + depth -= 1 + pos += 1 + + # After closing paren of CTE body, check for comma + if pos < len(tokens) and tokens[pos] == ",": + pos += 1 + # Continue to parse next CTE + else: + # No comma - CTE list is done, next token is the statement keyword + if pos < len(tokens): + return tokens[pos] + return "" + + return "" + + def guard(sql: str, max_rows: int = 200) -> str: """Validate `sql` as one read-only statement and return it ready to run. @@ -70,6 +195,17 @@ def guard(sql: str, max_rows: int = 200) -> str: f"got {keyword or '?'}. Run DDL against the target via " f"the clickhousectl MCP instead." ) + + # CTE-prefixed DML detection: WITH keyword is allowed, but only if + # the actual statement following the CTE definitions is SELECT + if keyword == "WITH": + actual_keyword = keyword_after_ctes(sql) + if actual_keyword != "SELECT": + raise SqlNotAllowed( + f"CTE-prefixed DML is not permitted; got {actual_keyword or '?'}, " + f"only SELECT is allowed after CTE definitions" + ) + body = sql.strip().rstrip(";").rstrip() if keyword in {"SELECT", "WITH"} and not _TRAILING_LIMIT.search( strip_comments(body) diff --git a/tests/test_sql_guard.py b/tests/test_sql_guard.py index 1ecbf92..d2dc098 100644 --- a/tests/test_sql_guard.py +++ b/tests/test_sql_guard.py @@ -98,3 +98,115 @@ def test_keyword_set_is_read_only(): assert READ_ONLY_LEADING_KEYWORDS == frozenset( {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"} ) + + +# CTE-prefixed DML rejection tests +@pytest.mark.parametrize( + "sql", + [ + "WITH x AS (SELECT 1) INSERT INTO t VALUES (1)", + "WITH x AS (SELECT 1) UPDATE t SET a = 1", + "WITH x AS (SELECT 1) DELETE FROM t", + "WITH x AS (SELECT 1) MERGE INTO t USING s ON t.a = s.a", + ], +) +def test_cte_with_dml_is_rejected(sql): + with pytest.raises(SqlNotAllowed) as excinfo: + guard(sql) + assert "CTE-prefixed DML" in str(excinfo.value) + + +def test_cte_with_recursive_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("WITH RECURSIVE x AS (SELECT 1) INSERT INTO t VALUES (1)") + + +def test_cte_with_column_list_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("WITH t (a, b) AS (SELECT 1, 2) INSERT INTO t VALUES (1, 2)") + + +def test_cte_with_multiple_ctes_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("WITH a AS (SELECT 1), b AS (SELECT 2) INSERT INTO t SELECT * FROM a") + + +def test_cte_with_nested_parens_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard( + "WITH x AS (SELECT 1 FROM (SELECT 2) y) INSERT INTO t SELECT * FROM x" + ) + + +def test_cte_with_string_literal_paren_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("WITH x AS (SELECT '(' AS s) INSERT INTO t SELECT * FROM x") + + +def test_cte_with_comment_hiding_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("WITH x AS (SELECT 1) /* c */ INSERT INTO t VALUES (1)") + + +def test_cte_with_values_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("WITH t (n) AS (VALUES (1)) INSERT INTO t SELECT * FROM t") + + +def test_cte_with_nested_cte_and_insert_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard( + "WITH x AS (WITH y AS (SELECT 1) SELECT * FROM y) " + "INSERT INTO t SELECT * FROM x" + ) + + +def test_cte_with_select_is_allowed(): + assert guard("WITH x AS (SELECT 1) SELECT * FROM x") + + +def test_cte_with_multiple_ctes_select_is_allowed(): + assert guard("WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b") + + +def test_cte_with_recursive_select_is_allowed(): + assert guard("WITH RECURSIVE x AS (SELECT 1) SELECT * FROM x") + + +def test_cte_with_column_list_select_is_allowed(): + assert guard("WITH t (a, b) AS (SELECT 1, 2) SELECT * FROM t") + + +def test_cte_with_nested_parens_select_is_allowed(): + assert guard( + "WITH x AS (SELECT 1 FROM (SELECT 2) y) SELECT * FROM x" + ) + + +def test_cte_with_string_literal_paren_select_is_allowed(): + assert guard("WITH x AS (SELECT '(' AS s) SELECT * FROM x") + + +def test_cte_with_values_select_is_allowed(): + assert guard("WITH t (n) AS (VALUES (1)) SELECT * FROM t") + + +def test_cte_with_nested_cte_select_is_allowed(): + assert guard("WITH x AS (WITH y AS (SELECT 1) SELECT * FROM y) SELECT * FROM x") + + +def test_cte_with_select_gets_limit(): + out = guard("WITH x AS (SELECT 1) SELECT * FROM x", max_rows=50) + assert out.endswith("LIMIT 50") + + +def test_cte_with_select_existing_limit_not_doubled(): + out = guard("WITH x AS (SELECT 1) SELECT * FROM x LIMIT 5", max_rows=50) + assert out.lower().count("limit") == 1 + + +def test_cte_with_multiple_ctes_select_gets_limit(): + out = guard( + "WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b", max_rows=50 + ) + assert out.endswith("LIMIT 50") From b0a1e9b5f33acda23034a76ce86064fff978f640 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 14:31:45 +0800 Subject: [PATCH 03/21] fix(databricks): close backtick bypass with unified normalizer 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). --- docker/databricks-mcp/sql_guard.py | 160 ++++++++++++++++++++++++----- tests/test_sql_guard.py | 99 ++++++++++++++++++ 2 files changed, 233 insertions(+), 26 deletions(-) diff --git a/docker/databricks-mcp/sql_guard.py b/docker/databricks-mcp/sql_guard.py index 1388604..79ca24c 100644 --- a/docker/databricks-mcp/sql_guard.py +++ b/docker/databricks-mcp/sql_guard.py @@ -15,10 +15,6 @@ {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"} ) -_LINE_COMMENT = re.compile(r"--[^\n]*") -_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.S) -_SINGLE_QUOTED = re.compile(r"'(?:''|[^'])*'") -_DOUBLE_QUOTED = re.compile(r'"(?:""|[^"])*"') _LEADING_WORD = re.compile(r"[(\s]*([A-Za-z_]+)") _TRAILING_LIMIT = re.compile(r"\blimit\b\s+\d+\s*$", re.I) @@ -27,25 +23,149 @@ class SqlNotAllowed(ValueError): """Raised when a statement is not a single read-only statement.""" +def _neutralize(sql: str) -> str: + """Neutralize all quoted/commented regions in one pass, preserving length. + + Returns a string of identical length where all characters inside any of + these five mutually-exclusive regions are replaced by spaces: + - line comments: -- to end of line (but not the newline itself) + - block comments: /* ... */ + - single-quoted strings: '...' ('' escapes) + - double-quoted strings: "..." ("" escapes) + - backtick identifiers: `...` (`` escapes) + + Whichever region opens first (in reading order) wins; once inside one, + the others cannot begin. This single-pass approach handles all interactions + correctly: '/*' is a string, not a comment; `--` is an identifier, not a + comment; -- 'unclosed remains a comment and eats the rest of the line. + """ + result = [] + i = 0 + + while i < len(sql): + # Line comment: -- to end of line + # Must come before backtick/quote checks + if i + 1 < len(sql) and sql[i : i + 2] == "--": + result.append(" ") + result.append(" ") + i += 2 + # Consume until newline but preserve the newline + while i < len(sql) and sql[i] not in "\n": + result.append(" ") + i += 1 + if i < len(sql) and sql[i] == "\n": + result.append("\n") + i += 1 + continue + + # Block comment: /* ... */ + if i + 1 < len(sql) and sql[i : i + 2] == "/*": + result.append(" ") + result.append(" ") + i += 2 + # Consume until */ + while i + 1 < len(sql): + if sql[i : i + 2] == "*/": + result.append(" ") + result.append(" ") + i += 2 + break + result.append(" ") + i += 1 + continue + + # Single-quoted string: '...' with '' as escape + if sql[i] == "'": + result.append(" ") + i += 1 + while i < len(sql): + if sql[i] == "'": + if i + 1 < len(sql) and sql[i + 1] == "'": + # Escaped quote + result.append(" ") + result.append(" ") + i += 2 + else: + # End of string + result.append(" ") + i += 1 + break + else: + result.append(" ") + i += 1 + continue + + # Double-quoted string: "..." with "" as escape + if sql[i] == '"': + result.append(" ") + i += 1 + while i < len(sql): + if sql[i] == '"': + if i + 1 < len(sql) and sql[i + 1] == '"': + # Escaped quote + result.append(" ") + result.append(" ") + i += 2 + else: + # End of string + result.append(" ") + i += 1 + break + else: + result.append(" ") + i += 1 + continue + + # Backtick-quoted identifier: `...` with `` as escape + if sql[i] == "`": + result.append(" ") + i += 1 + while i < len(sql): + if sql[i] == "`": + if i + 1 < len(sql) and sql[i + 1] == "`": + # Escaped backtick + result.append(" ") + result.append(" ") + i += 2 + else: + # End of identifier + result.append(" ") + i += 1 + break + else: + result.append(" ") + i += 1 + continue + + # Regular character + result.append(sql[i]) + i += 1 + + return "".join(result) + + def strip_comments(sql: str) -> str: """Remove block and line comments so they can't hide a leading keyword.""" - return _LINE_COMMENT.sub(" ", _BLOCK_COMMENT.sub(" ", sql)) + neutralized = _neutralize(sql) + # Replace spaces that came from neutralization with empty space + # (but keep real spaces from the original) + return neutralized def leading_keyword(sql: str) -> str: """First bare word of the statement, upper-cased. Leading parens and comments are skipped so `(SELECT 1)` and `/* c */ SELECT 1` both read as SELECT.""" - match = _LEADING_WORD.match(strip_comments(sql).strip()) + neutralized = _neutralize(sql).strip() + match = _LEADING_WORD.match(neutralized) return match.group(1).upper() if match else "" def is_multi_statement(sql: str) -> bool: """True when a ';' separates statements. Semicolons inside string literals and comments don't count, and a single trailing ';' is fine.""" - body = strip_comments(sql) - body = _DOUBLE_QUOTED.sub('""', _SINGLE_QUOTED.sub("''", body)) - return ";" in body.strip().rstrip(";") + neutralized = _neutralize(sql) + return ";" in neutralized.strip().rstrip(";") def keyword_after_ctes(sql: str) -> str: @@ -56,24 +176,12 @@ def keyword_after_ctes(sql: str) -> str: - "WITH x AS (SELECT 1) INSERT INTO ..." returns "INSERT" - "WITH x AS (SELECT 1) SELECT * FROM x" returns "SELECT" - Handles comments, string literals, RECURSIVE modifier, multiple CTEs, - and column lists. Uses a tokenizer to safely parse paren-balanced CTEs. + Handles comments, string literals (all types including backticks), + RECURSIVE modifier, multiple CTEs, and column lists. Uses a tokenizer + to safely parse paren-balanced CTEs against the neutralized version. """ - # Strip comments - body = strip_comments(sql).strip() - - # Replace string literals with placeholders to avoid parsing their contents - strings = {} - counter = [0] - - def replace_string(match): - placeholder = f"__STR_{counter[0]}__" - strings[placeholder] = match.group(0) - counter[0] += 1 - return placeholder - - body = _SINGLE_QUOTED.sub(replace_string, body) - body = _DOUBLE_QUOTED.sub(replace_string, body) + # Neutralize quotes and comments so they can't affect paren counting + body = _neutralize(sql).strip() # Tokenize: extract words and structural tokens (parens, commas) tokens = [] diff --git a/tests/test_sql_guard.py b/tests/test_sql_guard.py index d2dc098..d4b75e0 100644 --- a/tests/test_sql_guard.py +++ b/tests/test_sql_guard.py @@ -210,3 +210,102 @@ def test_cte_with_multiple_ctes_select_gets_limit(): "WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b", max_rows=50 ) assert out.endswith("LIMIT 50") + + +# Backtick identifier tests (security: backticks can hide parens and keywords) +@pytest.mark.parametrize( + "sql", + [ + # Backtick containing ) closes CTE early, revealing the hidden DML + "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) INSERT INTO orders SELECT * FROM x", + "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) UPDATE t SET a = 1", + "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) DELETE FROM t", + "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) MERGE INTO t USING s ON t.a = s.a", + ], +) +def test_backtick_with_hidden_dml_is_rejected(sql): + with pytest.raises(SqlNotAllowed) as excinfo: + guard(sql) + assert "CTE-prefixed DML" in str(excinfo.value) + + +def test_backtick_with_closing_paren_alone_and_insert(): + with pytest.raises(SqlNotAllowed): + guard( + "WITH x AS (SELECT 1 AS `a)` FROM t) INSERT INTO orders SELECT * FROM x" + ) + + +def test_backtick_with_semicolon_is_allowed(): + # This is a legitimate query: backtick contains semicolon + # Should NOT be treated as multi-statement + assert guard("SELECT 1 FROM `a;b`") + + +def test_backtick_with_line_comment_marker_is_allowed(): + # Backtick containing -- should not be treated as start of line comment + assert guard("SELECT 1 AS `a--b` FROM t") + + +def test_backtick_with_block_comment_marker_is_allowed(): + # Backtick containing /* should not start a block comment + assert guard("SELECT 1 AS `a/*b*/c` FROM t") + + +def test_backtick_with_doubled_backtick_escape(): + # Backtick with escaped backtick inside + assert guard("SELECT 1 AS `a``b` FROM t") + + +def test_backtick_with_doubled_backtick_and_paren(): + # Complex case: doubled backtick followed by paren + assert guard("SELECT 1 AS `a``)`b` FROM t") + + +def test_string_with_line_comment_marker_not_a_comment(): + # '--' inside single quotes is not a comment + assert guard("SELECT '--' AS s") + + +def test_backtick_identifier_with_block_comment_marker_not_comment(): + # '/*' inside backticks is not a comment + assert guard("SELECT 1 AS `/*` FROM t") + + +def test_line_comment_with_unclosed_string_eats_rest(): + # After --, everything to end of line is a comment, including 'unclosed string + # This is a legitimate statement with a line comment at the end + assert guard("SELECT 1 FROM t -- 'this is a comment") + + +def test_limit_inside_backtick_still_gets_real_limit(): + # If LIMIT appears inside a backtick identifier, a real LIMIT is still added + out = guard("SELECT * FROM t WHERE col = `LIMIT 5`", max_rows=50) + assert out.endswith("LIMIT 50") + # Verify there's only one real LIMIT (the injected one) + assert out.lower().count("limit") == 2 # 1 in backtick, 1 real + + +def test_select_with_backtick_column_and_limit_not_doubled(): + # Existing real LIMIT is not doubled even with backtick + out = guard("SELECT * FROM t WHERE col = `LIMIT` LIMIT 10", max_rows=50) + assert out.lower().count("limit") == 2 # 1 in backtick, 1 real + + +def test_backtick_in_cte_with_closing_paren_and_select(): + # Backtick with ) in CTE body, followed by legitimate SELECT + assert guard( + "WITH x AS (SELECT 1 AS `a)b` FROM t) SELECT * FROM x" + ) + + +def test_backtick_in_cte_with_closing_paren_gets_limit(): + out = guard( + "WITH x AS (SELECT 1 AS `a)b` FROM t) SELECT * FROM x", max_rows=50 + ) + assert out.endswith("LIMIT 50") + + +def test_string_literal_with_paren_not_confused_with_column_list(): + # String containing ( should not be confused with column list + assert guard("SELECT '(' AS s") From 7aa555ee48e9ce1ac6edb8c90497ba99638dcba6 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 15:48:40 +0800 Subject: [PATCH 04/21] fix(databricks): rewrite SQL guard on sqlglot's Databricks parser 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. --- docker/databricks-mcp/requirements.txt | 3 + docker/databricks-mcp/sql_guard.py | 386 +++++++------------------ tests/requirements.txt | 1 + tests/test_sql_guard.py | 359 ++++++++--------------- 4 files changed, 224 insertions(+), 525 deletions(-) create mode 100644 docker/databricks-mcp/requirements.txt 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/sql_guard.py b/docker/databricks-mcp/sql_guard.py index 79ca24c..9555d36 100644 --- a/docker/databricks-mcp/sql_guard.py +++ b/docker/databricks-mcp/sql_guard.py @@ -6,317 +6,137 @@ 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 re - -READ_ONLY_LEADING_KEYWORDS = frozenset( - {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"} -) - -_LEADING_WORD = re.compile(r"[(\s]*([A-Za-z_]+)") -_TRAILING_LIMIT = re.compile(r"\blimit\b\s+\d+\s*$", re.I) - - -class SqlNotAllowed(ValueError): - """Raised when a statement is not a single read-only statement.""" - - -def _neutralize(sql: str) -> str: - """Neutralize all quoted/commented regions in one pass, preserving length. - - Returns a string of identical length where all characters inside any of - these five mutually-exclusive regions are replaced by spaces: - - line comments: -- to end of line (but not the newline itself) - - block comments: /* ... */ - - single-quoted strings: '...' ('' escapes) - - double-quoted strings: "..." ("" escapes) - - backtick identifiers: `...` (`` escapes) - - Whichever region opens first (in reading order) wins; once inside one, - the others cannot begin. This single-pass approach handles all interactions - correctly: '/*' is a string, not a comment; `--` is an identifier, not a - comment; -- 'unclosed remains a comment and eats the rest of the line. - """ - result = [] - i = 0 +import logging - while i < len(sql): - # Line comment: -- to end of line - # Must come before backtick/quote checks - if i + 1 < len(sql) and sql[i : i + 2] == "--": - result.append(" ") - result.append(" ") - i += 2 - # Consume until newline but preserve the newline - while i < len(sql) and sql[i] not in "\n": - result.append(" ") - i += 1 - if i < len(sql) and sql[i] == "\n": - result.append("\n") - i += 1 - continue +import sqlglot +from sqlglot import exp +from sqlglot.dialects import Databricks +from sqlglot.errors import ParseError, TokenError - # Block comment: /* ... */ - if i + 1 < len(sql) and sql[i : i + 2] == "/*": - result.append(" ") - result.append(" ") - i += 2 - # Consume until */ - while i + 1 < len(sql): - if sql[i : i + 2] == "*/": - result.append(" ") - result.append(" ") - i += 2 - break - result.append(" ") - i += 1 - continue +# 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) - # Single-quoted string: '...' with '' as escape - if sql[i] == "'": - result.append(" ") - i += 1 - while i < len(sql): - if sql[i] == "'": - if i + 1 < len(sql) and sql[i + 1] == "'": - # Escaped quote - result.append(" ") - result.append(" ") - i += 2 - else: - # End of string - result.append(" ") - i += 1 - break - else: - result.append(" ") - i += 1 - continue +# Node types that represent a read-only statement once successfully parsed. +READ_ONLY_NODES = (exp.Select, exp.Union, exp.Subquery, exp.Describe) - # Double-quoted string: "..." with "" as escape - if sql[i] == '"': - result.append(" ") - i += 1 - while i < len(sql): - if sql[i] == '"': - if i + 1 < len(sql) and sql[i + 1] == '"': - # Escaped quote - result.append(" ") - result.append(" ") - i += 2 - else: - # End of string - result.append(" ") - i += 1 - break - else: - result.append(" ") - i += 1 - continue +# 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"}) - # Backtick-quoted identifier: `...` with `` as escape - if sql[i] == "`": - result.append(" ") - i += 1 - while i < len(sql): - if sql[i] == "`": - if i + 1 < len(sql) and sql[i + 1] == "`": - # Escaped backtick - result.append(" ") - result.append(" ") - i += 2 - else: - # End of identifier - result.append(" ") - i += 1 - break - else: - result.append(" ") - i += 1 - continue +_DATABRICKS = Databricks() +_DIALECT = "databricks" - # Regular character - result.append(sql[i]) - i += 1 - - return "".join(result) - - -def strip_comments(sql: str) -> str: - """Remove block and line comments so they can't hide a leading keyword.""" - neutralized = _neutralize(sql) - # Replace spaces that came from neutralization with empty space - # (but keep real spaces from the original) - return neutralized - - -def leading_keyword(sql: str) -> str: - """First bare word of the statement, upper-cased. Leading parens and - comments are skipped so `(SELECT 1)` and `/* c */ SELECT 1` both read - as SELECT.""" - neutralized = _neutralize(sql).strip() - match = _LEADING_WORD.match(neutralized) - return match.group(1).upper() if match else "" - - -def is_multi_statement(sql: str) -> bool: - """True when a ';' separates statements. Semicolons inside string - literals and comments don't count, and a single trailing ';' is fine.""" - neutralized = _neutralize(sql) - return ";" in neutralized.strip().rstrip(";") +class SqlNotAllowed(ValueError): + """Raised when a statement is not a single read-only statement.""" -def keyword_after_ctes(sql: str) -> str: - """Extract the top-level SQL keyword following CTE definitions. - When a statement starts with WITH, this returns the keyword of the actual - statement following the CTE list. For example: - - "WITH x AS (SELECT 1) INSERT INTO ..." returns "INSERT" - - "WITH x AS (SELECT 1) SELECT * FROM x" returns "SELECT" +def _strip_trailing_semicolon(sql: str, toks) -> str: + """Return `sql` with its single legal trailing ';' (if any) removed. - Handles comments, string literals (all types including backticks), - RECURSIVE modifier, multiple CTEs, and column lists. Uses a tokenizer - to safely parse paren-balanced CTEs against the neutralized version. + 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. """ - # Neutralize quotes and comments so they can't affect paren counting - body = _neutralize(sql).strip() - - # Tokenize: extract words and structural tokens (parens, commas) - tokens = [] - i = 0 - while i < len(body): - # Skip whitespace - while i < len(body) and body[i] in " \t\n\r": - i += 1 - if i >= len(body): - break - - # Extract words: letters, underscores, and digits - if body[i].isalpha() or body[i] == "_": - start = i - while i < len(body) and (body[i].isalnum() or body[i] == "_"): - i += 1 - tokens.append(body[start:i].upper()) - # Extract structural tokens - elif body[i] in "(),": - tokens.append(body[i]) - i += 1 - else: - # Skip other characters (operators, etc.) - i += 1 - - # Verify it starts with WITH - if not tokens or tokens[0] != "WITH": - return "" - - pos = 1 - - # Skip optional RECURSIVE keyword - if pos < len(tokens) and tokens[pos] == "RECURSIVE": - pos += 1 - - # Parse the CTE list: each CTE is name [(cols)] AS (body) - # Multiple CTEs are separated by commas - while pos < len(tokens): - # Pattern: NAME [ ( ... ) ] AS ( ... ) [, more CTEs] - - # Expect a name - if pos >= len(tokens) or not tokens[pos][0].isalpha(): - break - pos += 1 - - # Skip optional column list (parens before AS) - if pos < len(tokens) and tokens[pos] == "(": - # Look ahead: if next closing ) is followed by AS, this is a column list - save_pos = pos + 1 - depth = 1 - while save_pos < len(tokens) and depth > 0: - if tokens[save_pos] == "(": - depth += 1 - elif tokens[save_pos] == ")": - depth -= 1 - save_pos += 1 - - # Check if AS follows the closing paren - if save_pos < len(tokens) and tokens[save_pos] == "AS": - # Yes, it's a column list - skip it - pos = save_pos - # If not AS, this paren must be the CTE body, so don't skip - - # Expect AS - if pos >= len(tokens) or tokens[pos] != "AS": - # Malformed; return what we have - if pos < len(tokens): - return tokens[pos] - return "" - pos += 1 - - # Expect and skip balanced parens (CTE body) - if pos >= len(tokens) or tokens[pos] != "(": - if pos < len(tokens): - return tokens[pos] - return "" - - depth = 1 - pos += 1 - while pos < len(tokens) and depth > 0: - if tokens[pos] == "(": - depth += 1 - elif tokens[pos] == ")": - depth -= 1 - pos += 1 - - # After closing paren of CTE body, check for comma - if pos < len(tokens) and tokens[pos] == ",": - pos += 1 - # Continue to parse next CTE - else: - # No comma - CTE list is done, next token is the statement keyword - if pos < len(tokens): - return tokens[pos] - return "" - - return "" + if toks and toks[-1].token_type == sqlglot.TokenType.SEMICOLON: + return sql[: toks[-1].start].rstrip() + return sql.strip() 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/WITH statements that don't already end - in one, so an unbounded scan can't stream a whole fact table into the - chat context. + A row cap is appended to SELECT/UNION statements 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.""" + Raises SqlNotAllowed for anything else. + """ if not sql or not sql.strip(): - raise SqlNotAllowed("empty statement") - if is_multi_statement(sql): + 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" + "multiple statements are not allowed; send one statement at a time" ) - keyword = leading_keyword(sql) - if keyword not in READ_ONLY_LEADING_KEYWORDS: - allowed = ", ".join(sorted(READ_ONLY_LEADING_KEYWORDS)) + + try: + stmts = [s for s in sqlglot.parse(sql, dialect=_DIALECT) if s] + 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; only {allowed} are permitted, " - f"got {keyword or '?'}. Run DDL against the target via " - f"the clickhousectl MCP instead." + 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" ) - # CTE-prefixed DML detection: WITH keyword is allowed, but only if - # the actual statement following the CTE definitions is SELECT - if keyword == "WITH": - actual_keyword = keyword_after_ctes(sql) - if actual_keyword != "SELECT": + 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"CTE-prefixed DML is not permitted; got {actual_keyword or '?'}, " - f"only SELECT is allowed after CTE definitions" + 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 = sql.strip().rstrip(";").rstrip() - if keyword in {"SELECT", "WITH"} and not _TRAILING_LIMIT.search( - strip_comments(body) - ): + body = _strip_trailing_semicolon(sql, toks) + if isinstance(root, (exp.Select, exp.Union)) and root.args.get("limit") is None: body = f"{body}\nLIMIT {max_rows}" return body diff --git a/tests/requirements.txt b/tests/requirements.txt index 039d26e..5a15428 100644 --- a/tests/requirements.txt +++ b/tests/requirements.txt @@ -1 +1,2 @@ pytest>=8.0 +sqlglot>=30.0 diff --git a/tests/test_sql_guard.py b/tests/test_sql_guard.py index d4b75e0..abd53e1 100644 --- a/tests/test_sql_guard.py +++ b/tests/test_sql_guard.py @@ -3,6 +3,14 @@ Run from the repo root: python3 -m pip install -r tests/requirements.txt python3 -m pytest tests/test_sql_guard.py -v + +sql_guard.py classifies statements by tokenizing/parsing them with sqlglot's +Databricks dialect rather than hand-lexing, because three earlier review +rounds each found a real bypass in the hand-lexed version (CTE-prefixed DML, +a backtick identifier that blinded the CTE scanner, and comment/string +lexing that diverged from Spark's grammar). The must-allow/must-reject +lists below are the cases empirically verified against sqlglot 30.15.0 +during the design of this rewrite, including all three of those bypasses. """ import sys from pathlib import Path @@ -11,301 +19,168 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "docker" / "databricks-mcp")) -from sql_guard import READ_ONLY_LEADING_KEYWORDS, SqlNotAllowed, guard # noqa: E402 +from sql_guard import SqlNotAllowed, guard # noqa: E402 + +# --- Verified cases ------------------------------------------------------- +# +# The two entries containing a literal carriage return are written with the +# "\r" escape below (not a raw CR byte) so they survive editors/git intact; +# Python turns the escape into an actual CR character at parse time. + +MUST_ALLOW = [ + "SELECT 1", + "SELECT 1;", + "select c from customer", + "(SELECT 1)", + "SELECT 1 UNION ALL SELECT 2", + "WITH x AS (SELECT 1) SELECT * FROM x", + "WITH t (n) AS (VALUES (1)) SELECT * FROM t", + "WITH RECURSIVE r (n) AS (VALUES (1) UNION ALL SELECT n+1 FROM r WHERE n<6) SELECT * FROM r", + "SELECT 1 FROM `a;b`", + "SELECT 'a;b' AS s", + "SELECT 1 AS `a--b` FROM t", + "SHOW CATALOGS", + "DESCRIBE TABLE m.t.o", + "DESCRIBE DETAIL m.t.l", + "DESCRIBE HISTORY m.t.l LIMIT 5", + "DESC DETAIL m.t.l", + "EXPLAIN SELECT 1", +] + +MUST_REJECT = [ + # The first five are the previously-verified bypasses. + "WITH x AS (SELECT 1) INSERT INTO orders SELECT * FROM x", + "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) INSERT INTO orders SELECT * FROM x", + "SHOW TABLES -- x\r; DROP TABLE t", + "SELECT 1 -- c\r; INSERT INTO orders VALUES (1)", + "SELECT 'a\\' AS x, ' ; DROP TABLE t", + "SELECT 1; DROP TABLE t", + "INSERT INTO t VALUES (1)", + "UPDATE t SET a=1", + "DELETE FROM t", + "DROP TABLE t", + "CREATE TABLE t (a INT)", + "ALTER TABLE t ADD COLUMN b INT", + "TRUNCATE TABLE t", + "COPY INTO t FROM 's3://b/k'", + "GRANT SELECT ON TABLE t TO `u`", + "MERGE INTO t USING s ON t.a=s.a WHEN MATCHED THEN UPDATE SET t.b=s.b", + "WITH x AS (SELECT 1) MERGE INTO t USING x ON t.a=x.a WHEN MATCHED THEN DELETE", + "OPTIMIZE m.t.l ZORDER BY (a)", + "VACUUM m.t.l", + " ", # empty / whitespace-only +] @pytest.mark.parametrize( - "sql", - [ - "SELECT 1", - "select c_custkey from customer", - "WITH x AS (SELECT 1) SELECT * FROM x", - "SHOW CATALOGS", - "DESCRIBE TABLE migration_demo.tpch.orders", - "DESC DETAIL migration_demo.tpch.lineitem", - "EXPLAIN SELECT 1", - " \n SELECT 1 \n ", - "(SELECT 1)", - ], + "sql", MUST_ALLOW, ids=[f"allow-{i:02d}" for i in range(len(MUST_ALLOW))] ) -def test_read_only_statements_are_allowed(sql): +def test_verified_read_only_statements_are_allowed(sql): assert guard(sql) @pytest.mark.parametrize( - "sql", - [ - "INSERT INTO t VALUES (1)", - "UPDATE t SET a = 1", - "DELETE FROM t", - "DROP TABLE t", - "CREATE TABLE t (a INT)", - "ALTER TABLE t ADD COLUMN b INT", - "MERGE INTO t USING s ON t.a = s.a", - "TRUNCATE TABLE t", - "COPY INTO t FROM 's3://b/k'", - "GRANT SELECT ON TABLE t TO `u`", - ], + "sql", MUST_REJECT, ids=[f"reject-{i:02d}" for i in range(len(MUST_REJECT))] ) -def test_mutating_statements_are_rejected(sql): +def test_verified_mutations_and_malformed_input_are_rejected(sql): with pytest.raises(SqlNotAllowed): guard(sql) -def test_multi_statement_is_rejected(): +def test_multi_statement_error_message_mentions_one_statement(): with pytest.raises(SqlNotAllowed) as excinfo: guard("SELECT 1; DROP TABLE t") assert "one statement" in str(excinfo.value) -def test_trailing_semicolon_is_not_multi_statement(): - assert guard("SELECT 1;").startswith("SELECT 1") - - -def test_semicolon_inside_string_literal_is_not_multi_statement(): - assert guard("SELECT 'a;b' AS s") - - -def test_comment_hidden_mutation_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("-- harmless\nDROP TABLE t") - - -def test_block_comment_hidden_mutation_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("/* SELECT */ DROP TABLE t") - - -def test_limit_is_injected_when_absent(): - assert guard("SELECT * FROM orders", max_rows=50).endswith("LIMIT 50") - - -def test_existing_limit_is_not_doubled(): - out = guard("SELECT * FROM orders LIMIT 5", max_rows=50) - assert out.lower().count("limit") == 1 - - -def test_show_does_not_get_a_limit(): - assert "LIMIT" not in guard("SHOW CATALOGS") - - -def test_empty_statement_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard(" ") - - -def test_keyword_set_is_read_only(): - assert READ_ONLY_LEADING_KEYWORDS == frozenset( - {"SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN"} - ) - - -# CTE-prefixed DML rejection tests -@pytest.mark.parametrize( - "sql", - [ - "WITH x AS (SELECT 1) INSERT INTO t VALUES (1)", - "WITH x AS (SELECT 1) UPDATE t SET a = 1", - "WITH x AS (SELECT 1) DELETE FROM t", - "WITH x AS (SELECT 1) MERGE INTO t USING s ON t.a = s.a", - ], -) -def test_cte_with_dml_is_rejected(sql): +def test_mutation_error_message_points_to_clickhousectl(): with pytest.raises(SqlNotAllowed) as excinfo: - guard(sql) - assert "CTE-prefixed DML" in str(excinfo.value) - - -def test_cte_with_recursive_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("WITH RECURSIVE x AS (SELECT 1) INSERT INTO t VALUES (1)") + guard("DROP TABLE t") + assert "clickhousectl" in str(excinfo.value) -def test_cte_with_column_list_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("WITH t (a, b) AS (SELECT 1, 2) INSERT INTO t VALUES (1, 2)") - - -def test_cte_with_multiple_ctes_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("WITH a AS (SELECT 1), b AS (SELECT 2) INSERT INTO t SELECT * FROM a") +# --- LIMIT injection ------------------------------------------------------- -def test_cte_with_nested_parens_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard( - "WITH x AS (SELECT 1 FROM (SELECT 2) y) INSERT INTO t SELECT * FROM x" - ) +def test_limit_word_inside_backtick_alias_gets_exactly_one_real_limit(): + out = guard("SELECT 1 AS `x LIMIT 5` FROM t") + # The backtick alias contains the word LIMIT verbatim; only the + # appended clause should introduce a *new* line starting with LIMIT. + assert out.count("\nLIMIT") == 1 + assert out.endswith("LIMIT 200") -def test_cte_with_string_literal_paren_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("WITH x AS (SELECT '(' AS s) INSERT INTO t SELECT * FROM x") - - -def test_cte_with_comment_hiding_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("WITH x AS (SELECT 1) /* c */ INSERT INTO t VALUES (1)") - - -def test_cte_with_values_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard("WITH t (n) AS (VALUES (1)) INSERT INTO t SELECT * FROM t") - - -def test_cte_with_nested_cte_and_insert_is_rejected(): - with pytest.raises(SqlNotAllowed): - guard( - "WITH x AS (WITH y AS (SELECT 1) SELECT * FROM y) " - "INSERT INTO t SELECT * FROM x" - ) - - -def test_cte_with_select_is_allowed(): - assert guard("WITH x AS (SELECT 1) SELECT * FROM x") - - -def test_cte_with_multiple_ctes_select_is_allowed(): - assert guard("WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b") - - -def test_cte_with_recursive_select_is_allowed(): - assert guard("WITH RECURSIVE x AS (SELECT 1) SELECT * FROM x") - - -def test_cte_with_column_list_select_is_allowed(): - assert guard("WITH t (a, b) AS (SELECT 1, 2) SELECT * FROM t") - - -def test_cte_with_nested_parens_select_is_allowed(): - assert guard( - "WITH x AS (SELECT 1 FROM (SELECT 2) y) SELECT * FROM x" - ) - - -def test_cte_with_string_literal_paren_select_is_allowed(): - assert guard("WITH x AS (SELECT '(' AS s) SELECT * FROM x") - - -def test_cte_with_values_select_is_allowed(): - assert guard("WITH t (n) AS (VALUES (1)) SELECT * FROM t") - - -def test_cte_with_nested_cte_select_is_allowed(): - assert guard("WITH x AS (WITH y AS (SELECT 1) SELECT * FROM y) SELECT * FROM x") - - -def test_cte_with_select_gets_limit(): - out = guard("WITH x AS (SELECT 1) SELECT * FROM x", max_rows=50) - assert out.endswith("LIMIT 50") - - -def test_cte_with_select_existing_limit_not_doubled(): - out = guard("WITH x AS (SELECT 1) SELECT * FROM x LIMIT 5", max_rows=50) - assert out.lower().count("limit") == 1 - - -def test_cte_with_multiple_ctes_select_gets_limit(): - out = guard( - "WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b", max_rows=50 - ) - assert out.endswith("LIMIT 50") - - -# Backtick identifier tests (security: backticks can hide parens and keywords) @pytest.mark.parametrize( "sql", [ - # Backtick containing ) closes CTE early, revealing the hidden DML - "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) INSERT INTO orders SELECT * FROM x", - "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) UPDATE t SET a = 1", - "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) DELETE FROM t", - "WITH x AS (SELECT 1 AS `a) SELECT 2 AS b` FROM t) MERGE INTO t USING s ON t.a = s.a", + "SHOW CATALOGS", + "DESCRIBE TABLE m.t.o", + "DESC DETAIL m.t.l", + "EXPLAIN SELECT 1", ], ) -def test_backtick_with_hidden_dml_is_rejected(sql): - with pytest.raises(SqlNotAllowed) as excinfo: - guard(sql) - assert "CTE-prefixed DML" in str(excinfo.value) - - -def test_backtick_with_closing_paren_alone_and_insert(): - with pytest.raises(SqlNotAllowed): - guard( - "WITH x AS (SELECT 1 AS `a)` FROM t) INSERT INTO orders SELECT * FROM x" - ) +def test_show_describe_explain_receive_no_limit(sql): + assert "LIMIT" not in guard(sql) -def test_backtick_with_semicolon_is_allowed(): - # This is a legitimate query: backtick contains semicolon - # Should NOT be treated as multi-statement - assert guard("SELECT 1 FROM `a;b`") +def test_limit_is_injected_using_max_rows(): + assert guard("SELECT * FROM orders", max_rows=50).endswith("LIMIT 50") -def test_backtick_with_line_comment_marker_is_allowed(): - # Backtick containing -- should not be treated as start of line comment - assert guard("SELECT 1 AS `a--b` FROM t") +def test_existing_limit_is_not_doubled(): + out = guard("SELECT * FROM orders LIMIT 5", max_rows=50) + assert out.lower().count("limit") == 1 -def test_backtick_with_block_comment_marker_is_allowed(): - # Backtick containing /* should not start a block comment - assert guard("SELECT 1 AS `a/*b*/c` FROM t") +def test_union_without_limit_gets_one_appended(): + out = guard("SELECT 1 UNION ALL SELECT 2", max_rows=10) + assert out.count("\nLIMIT") == 1 + assert out.endswith("LIMIT 10") -def test_backtick_with_doubled_backtick_escape(): - # Backtick with escaped backtick inside - assert guard("SELECT 1 AS `a``b` FROM t") +def test_union_with_existing_limit_is_not_doubled(): + out = guard("SELECT 1 UNION ALL SELECT 2 LIMIT 1", max_rows=10) + assert out.lower().count("limit") == 1 -def test_backtick_with_doubled_backtick_and_paren(): - # Complex case: doubled backtick followed by paren - assert guard("SELECT 1 AS `a``)`b` FROM t") +# --- guard returns the original text, never a sqlglot regeneration -------- -def test_string_with_line_comment_marker_not_a_comment(): - # '--' inside single quotes is not a comment - assert guard("SELECT '--' AS s") +def test_guard_returns_original_text_with_only_limit_appended(): + sql = "SELECT 1, `weird Name` FROM t" + out = guard(sql, max_rows=77) + assert out == sql + "\nLIMIT 77" -def test_backtick_identifier_with_block_comment_marker_not_comment(): - # '/*' inside backticks is not a comment - assert guard("SELECT 1 AS `/*` FROM t") +def test_guard_preserves_original_formatting_when_no_limit_is_appended(): + sql = "SHOW CATALOGS" + assert guard(sql) == sql -def test_line_comment_with_unclosed_string_eats_rest(): - # After --, everything to end of line is a comment, including 'unclosed string - # This is a legitimate statement with a line comment at the end - assert guard("SELECT 1 FROM t -- 'this is a comment") +def test_guard_strips_the_trailing_semicolon_before_appending_limit(): + out = guard("SELECT 1;", max_rows=5) + assert out == "SELECT 1\nLIMIT 5" -def test_limit_inside_backtick_still_gets_real_limit(): - # If LIMIT appears inside a backtick identifier, a real LIMIT is still added - out = guard("SELECT * FROM t WHERE col = `LIMIT 5`", max_rows=50) - assert out.endswith("LIMIT 50") - # Verify there's only one real LIMIT (the injected one) - assert out.lower().count("limit") == 2 # 1 in backtick, 1 real +def test_guard_strips_trailing_semicolon_when_no_limit_is_needed(): + out = guard("SHOW CATALOGS;") + assert out == "SHOW CATALOGS" -def test_select_with_backtick_column_and_limit_not_doubled(): - # Existing real LIMIT is not doubled even with backtick - out = guard("SELECT * FROM t WHERE col = `LIMIT` LIMIT 10", max_rows=50) - assert out.lower().count("limit") == 2 # 1 in backtick, 1 real +def test_guard_does_not_double_strip_semicolon_inside_backtick(): + sql = "SELECT 1 FROM `a;b`" + out = guard(sql, max_rows=5) + assert out == sql + "\nLIMIT 5" -def test_backtick_in_cte_with_closing_paren_and_select(): - # Backtick with ) in CTE body, followed by legitimate SELECT - assert guard( - "WITH x AS (SELECT 1 AS `a)b` FROM t) SELECT * FROM x" - ) +# --- misc ------------------------------------------------------------- -def test_backtick_in_cte_with_closing_paren_gets_limit(): - out = guard( - "WITH x AS (SELECT 1 AS `a)b` FROM t) SELECT * FROM x", max_rows=50 - ) - assert out.endswith("LIMIT 50") +def test_empty_statement_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard("") -def test_string_literal_with_paren_not_confused_with_column_list(): - # String containing ( should not be confused with column list - assert guard("SELECT '(' AS s") +def test_whitespace_only_statement_is_rejected(): + with pytest.raises(SqlNotAllowed): + guard(" \n\t ") From efdd08780849a6daf7087be5cb1edbac4d285461 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 16:03:12 +0800 Subject: [PATCH 05/21] fix(databricks): cap parenthesized queries; allow trailing comments 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. --- docker/databricks-mcp/sql_guard.py | 47 +++++++++++++++++++++++++++--- tests/test_sql_guard.py | 46 +++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) diff --git a/docker/databricks-mcp/sql_guard.py b/docker/databricks-mcp/sql_guard.py index 9555d36..de45a50 100644 --- a/docker/databricks-mcp/sql_guard.py +++ b/docker/databricks-mcp/sql_guard.py @@ -63,11 +63,39 @@ def _strip_trailing_semicolon(sql: str, toks) -> str: 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 that don't already - have one, so an unbounded scan can't stream a whole fact table into the + 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. @@ -93,7 +121,15 @@ def guard(sql: str, max_rows: int = 200) -> str: ) try: - stmts = [s for s in sqlglot.parse(sql, dialect=_DIALECT) if s] + # `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 @@ -137,6 +173,9 @@ def guard(sql: str, max_rows: int = 200) -> str: ) body = _strip_trailing_semicolon(sql, toks) - if isinstance(root, (exp.Select, exp.Union)) and root.args.get("limit") is None: + 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/tests/test_sql_guard.py b/tests/test_sql_guard.py index abd53e1..af980cd 100644 --- a/tests/test_sql_guard.py +++ b/tests/test_sql_guard.py @@ -99,6 +99,23 @@ def test_mutation_error_message_points_to_clickhousectl(): assert "clickhousectl" in str(excinfo.value) +def test_trailing_comment_after_semicolon_is_not_a_second_statement(): + # Regression: sqlglot.parse() emits a synthetic exp.Semicolon node to + # carry a comment that trails the terminator. That node is truthy, so + # it survived the `if s` None-filter and made this look like two + # statements. It must not, since it fails closed either way but a + # trailing comment on a query is ordinary input. + assert guard("SELECT 1; -- comment") == "SELECT 1\nLIMIT 200" + + +def test_trailing_comment_fix_does_not_widen_the_multi_statement_gate(): + # Paired negative case: a *real* second statement after the comment + # must still be rejected. The exp.Semicolon filter must not swallow + # an actual statement, only the content-free placeholder node. + with pytest.raises(SqlNotAllowed): + guard("SELECT 1; -- comment\nDROP TABLE t") + + # --- LIMIT injection ------------------------------------------------------- @@ -143,6 +160,35 @@ def test_union_with_existing_limit_is_not_doubled(): assert out.lower().count("limit") == 1 +def test_parenthesized_select_without_limit_gets_one_appended(): + # Regression: a top-level `(SELECT ...)` parses to exp.Subquery, not + # exp.Select/exp.Union, and was previously exempted from LIMIT + # injection entirely — silently defeating the row cap. + out = guard("(SELECT * FROM orders)", max_rows=50) + assert out == "(SELECT * FROM orders)\nLIMIT 50" + + +def test_parenthesized_select_with_existing_limit_is_not_doubled(): + # The existing LIMIT lives on the inner Select, not the outer + # Subquery; both slots must be checked or this gets a second LIMIT. + out = guard("(SELECT * FROM orders LIMIT 5)", max_rows=50) + assert out == "(SELECT * FROM orders LIMIT 5)" + assert out.lower().count("limit") == 1 + + +def test_parenthesized_union_without_limit_gets_one_appended(): + out = guard("(SELECT * FROM a UNION ALL SELECT * FROM b)", max_rows=50) + assert out == "(SELECT * FROM a UNION ALL SELECT * FROM b)\nLIMIT 50" + + +def test_parenthesized_union_with_existing_limit_is_not_doubled(): + out = guard( + "(SELECT * FROM a UNION ALL SELECT * FROM b LIMIT 5)", max_rows=50 + ) + assert out == "(SELECT * FROM a UNION ALL SELECT * FROM b LIMIT 5)" + assert out.lower().count("limit") == 1 + + # --- guard returns the original text, never a sqlglot regeneration -------- From 88e74a953eb21704fd0bc501d651401566820988 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 16:13:34 +0800 Subject: [PATCH 06/21] feat(databricks): read-only source MCP server (5 tools, SSE) --- docker/databricks-mcp/Dockerfile | 10 ++ docker/databricks-mcp/server.py | 196 +++++++++++++++++++++++++++++++ 2 files changed, 206 insertions(+) create mode 100644 docker/databricks-mcp/Dockerfile create mode 100644 docker/databricks-mcp/server.py 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/server.py b/docker/databricks-mcp/server.py new file mode 100644 index 0000000..4ab1d8b --- /dev/null +++ b/docker/databricks-mcp/server.py @@ -0,0 +1,196 @@ +"""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}`" + + +def _literal(value: str) -> str: + """Single-quote a string literal for interpolation into SQL.""" + return "'" + str(value).replace("'", "''") + "'" + + +@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 " + f"WHERE catalog_name = {_literal(catalog)} " + "ORDER BY schema_name" + ) + 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 " + f"WHERE table_catalog = {_literal(catalog)} " + f" AND table_schema = {_literal(schema)} " + "ORDER BY table_name" + ) + 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") From 87a181f9fc3e627cc1de4f43b4a26bbd1e3b6169 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 16:28:09 +0800 Subject: [PATCH 07/21] fix(databricks): use bound parameters instead of _literal() for SQL values _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. --- docker/databricks-mcp/server.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/docker/databricks-mcp/server.py b/docker/databricks-mcp/server.py index 4ab1d8b..932b72b 100644 --- a/docker/databricks-mcp/server.py +++ b/docker/databricks-mcp/server.py @@ -81,11 +81,6 @@ def _ident(name: str) -> str: return f"`{cleaned}`" -def _literal(value: str) -> str: - """Single-quote a string literal for interpolation into SQL.""" - return "'" + str(value).replace("'", "''") + "'" - - @mcp.tool() def list_catalogs() -> list[dict[str, Any]]: """List Unity Catalog catalogs visible to this principal.""" @@ -105,8 +100,9 @@ def list_schemas(catalog: str) -> list[dict[str, Any]]: cur.execute( "SELECT schema_name, comment " "FROM system.information_schema.schemata " - f"WHERE catalog_name = {_literal(catalog)} " - "ORDER BY schema_name" + "WHERE catalog_name = ? " + "ORDER BY schema_name", + [catalog], ) return _rows(cur) @@ -123,9 +119,10 @@ def list_tables(catalog: str, schema: str) -> list[dict[str, Any]]: cur.execute( "SELECT table_name, table_type, comment " "FROM system.information_schema.tables " - f"WHERE table_catalog = {_literal(catalog)} " - f" AND table_schema = {_literal(schema)} " - "ORDER BY table_name" + "WHERE table_catalog = ? " + " AND table_schema = ? " + "ORDER BY table_name", + [catalog, schema], ) tables = _rows(cur) for row in tables: From 77260c890939d50476170e57ca03a0ddc3570edf Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 16:34:16 +0800 Subject: [PATCH 08/21] feat(databricks): DatabricksSource with direct + S3-staged paths 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. --- .../migration-runner/migrationkit/__init__.py | 2 + docker/migration-runner/migrationkit/api.py | 3 + .../migrationkit/sources/__init__.py | 2 + .../migrationkit/sources/databricks.py | 293 ++++++++++++++++++ docker/migration-runner/requirements.txt | 1 + tests/test_databricks_source.py | 73 +++++ 6 files changed, 374 insertions(+) create mode 100644 docker/migration-runner/migrationkit/sources/databricks.py create mode 100644 tests/test_databricks_source.py 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/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/databricks.py b/docker/migration-runner/migrationkit/sources/databricks.py new file mode 100644 index 0000000..cf8f90f --- /dev/null +++ b/docker/migration-runner/migrationkit/sources/databricks.py @@ -0,0 +1,293 @@ +"""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: + 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}"} + ) + try: + with urllib.request.urlopen(request, timeout=10) as response: + body = json.load(response) + except (urllib.error.URLError, OSError, ValueError): + return None + 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 + + 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}" + ) + + 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 {self._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/tests/test_databricks_source.py b/tests/test_databricks_source.py new file mode 100644 index 0000000..726b229 --- /dev/null +++ b/tests/test_databricks_source.py @@ -0,0 +1,73 @@ +"""Unit tests for DatabricksSource's pure helpers. + +These are the parts testable without a Databricks workspace. The class +imports `databricks.sql` lazily inside __init__, so importing the module +needs no connector installed. + +Run from the repo root: + python3 -m pytest tests/test_databricks_source.py -v +""" +import sys +from dataclasses import dataclass +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "docker" / "migration-runner")) + +from migrationkit.sources.databricks import ( # noqa: E402 + normalize_host, + parquet_only, + split_namespace, +) + + +@dataclass +class FakeObject: + key: str + size: int + + +def test_split_namespace_returns_catalog_and_schema(): + assert split_namespace("migration_demo.tpch") == ("migration_demo", "tpch") + + +def test_split_namespace_tolerates_whitespace(): + assert split_namespace(" migration_demo . tpch ") == ("migration_demo", "tpch") + + +@pytest.mark.parametrize("bad", ["", " ", "tpch", "a.b.c", ".", "a."]) +def test_split_namespace_rejects_anything_but_two_parts(bad): + with pytest.raises(ValueError) as excinfo: + split_namespace(bad) + assert "DATABRICKS_NAMESPACE" in str(excinfo.value) + + +@pytest.mark.parametrize( + "raw,expected", + [ + ("https://dbc-abc.cloud.databricks.com", "dbc-abc.cloud.databricks.com"), + ("http://dbc-abc.cloud.databricks.com/", "dbc-abc.cloud.databricks.com"), + ("dbc-abc.cloud.databricks.com", "dbc-abc.cloud.databricks.com"), + (" https://dbc-abc.cloud.databricks.com/ ", "dbc-abc.cloud.databricks.com"), + ], +) +def test_normalize_host_strips_scheme_and_trailing_slash(raw, expected): + assert normalize_host(raw) == expected + + +def test_parquet_only_drops_commit_protocol_markers(): + objects = [ + FakeObject("p/run/lineitem/part-00000-abc.snappy.parquet", 1000), + FakeObject("p/run/lineitem/part-00001-def.snappy.parquet", 2000), + FakeObject("p/run/lineitem/_SUCCESS", 0), + FakeObject("p/run/lineitem/_committed_12345", 120), + FakeObject("p/run/lineitem/_started_12345", 80), + ] + kept = parquet_only(objects) + assert len(kept) == 2 + assert sum(o.size for o in kept) == 3000 + + +def test_parquet_only_is_case_insensitive(): + assert len(parquet_only([FakeObject("a/B.PARQUET", 1)])) == 1 From 467078179fdf789b38c176ed2628a31d00c36a09 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 16:44:56 +0800 Subject: [PATCH 09/21] fix(databricks): degrade history-API parsing to None, fix unload_to_s3 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. --- .../migrationkit/sources/databricks.py | 43 +++++++---- tests/test_databricks_source.py | 74 +++++++++++++++++++ 2 files changed, 101 insertions(+), 16 deletions(-) diff --git a/docker/migration-runner/migrationkit/sources/databricks.py b/docker/migration-runner/migrationkit/sources/databricks.py index cf8f90f..e8c9a81 100644 --- a/docker/migration-runner/migrationkit/sources/databricks.py +++ b/docker/migration-runner/migrationkit/sources/databricks.py @@ -185,25 +185,31 @@ def _fetch_server_ms(self, statement_id: str) -> float | None: return self._server_ms_from_system_table(statement_id) def _server_ms_from_history_api(self, statement_id: str) -> float | None: - 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}"} - ) + """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) - except (urllib.error.URLError, OSError, ValueError): + 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 - 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 def _server_ms_from_system_table(self, statement_id: str) -> float | None: # NOTE: deviates from the brief, which interpolated statement_id @@ -255,6 +261,11 @@ def unload_to_s3( 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: @@ -263,7 +274,7 @@ def unload_to_s3( cur.execute( f"INSERT OVERWRITE DIRECTORY '{target_uri}'\n" f"USING PARQUET\n" - f"SELECT * FROM {self._fq(table)}" + f"SELECT * FROM {fq_table}" ) except Exception as exc: raise RuntimeError( diff --git a/tests/test_databricks_source.py b/tests/test_databricks_source.py index 726b229..e840c6b 100644 --- a/tests/test_databricks_source.py +++ b/tests/test_databricks_source.py @@ -7,7 +7,9 @@ Run from the repo root: python3 -m pytest tests/test_databricks_source.py -v """ +import json import sys +import urllib.request from dataclasses import dataclass from pathlib import Path @@ -16,6 +18,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "docker" / "migration-runner")) from migrationkit.sources.databricks import ( # noqa: E402 + DatabricksSource, normalize_host, parquet_only, split_namespace, @@ -71,3 +74,74 @@ def test_parquet_only_drops_commit_protocol_markers(): def test_parquet_only_is_case_insensitive(): assert len(parquet_only([FakeObject("a/B.PARQUET", 1)])) == 1 + + +def _unconnected_source(catalog, schema): + """Build a DatabricksSource without touching __init__ (no network, + no connector import) — enough to exercise the pure `_fq` logic.""" + src = DatabricksSource.__new__(DatabricksSource) + src.catalog = catalog + src.schema = schema + return src + + +def test_fq_qualifies_a_bare_table_name(): + src = _unconnected_source("migration_demo", "tpch") + assert src._fq("orders") == "migration_demo.tpch.orders" + + +def test_fq_passes_through_an_already_qualified_name_unchanged(): + src = _unconnected_source("migration_demo", "tpch") + assert src._fq("other_catalog.other_schema.orders") == "other_catalog.other_schema.orders" + + +def test_fq_raises_for_bare_name_without_a_configured_namespace(): + src = _unconnected_source(None, None) + with pytest.raises(ValueError): + src._fq("orders") + + +class _FakeHTTPResponse: + """Minimal stand-in for what `urllib.request.urlopen` returns — supports + the context-manager protocol and the `.read()` that `json.load` uses, + without touching the network.""" + + def __init__(self, body): + self._raw = json.dumps(body).encode() + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def read(self): + return self._raw + + +@pytest.mark.parametrize( + "malformed_body", + [ + None, # null body + [], # bare list body + {"res": "not-a-list"}, # res is a string + {"res": ["not-a-dict"]}, # res item not a dict + {"res": [{"metrics": "not-a-dict"}]}, # metrics not a dict + ], +) +def test_server_ms_from_history_api_returns_none_on_malformed_response( + monkeypatch, malformed_body +): + """Regression test: the response shape is an admitted guess against a + live API. If the guess is wrong, this must degrade to None, never + raise — an uncaught exception here fails the whole benchmark row + instead of falling back to wall_ms.""" + src = DatabricksSource.__new__(DatabricksSource) + src._host = "dbc-abc.cloud.databricks.com" + src._token = "fake-token" + + monkeypatch.setattr( + urllib.request, "urlopen", lambda *a, **kw: _FakeHTTPResponse(malformed_body) + ) + + assert src._server_ms_from_history_api("stmt-id-123") is None From bdc0369b2da183fcf409e81e899d05a3b8b76b02 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 16:50:43 +0800 Subject: [PATCH 10/21] feat(databricks): source manifest and sample OLAP queries --- sources/databricks/manifest.json | 6 + .../queries/sample_olap_queries.sql | 105 ++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 sources/databricks/manifest.json create mode 100644 sources/databricks/queries/sample_olap_queries.sql 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/queries/sample_olap_queries.sql b/sources/databricks/queries/sample_olap_queries.sql new file mode 100644 index 0000000..297e8a4 --- /dev/null +++ b/sources/databricks/queries/sample_olap_queries.sql @@ -0,0 +1,105 @@ +-- Sample OLAP queries for the migration_demo.tpch workload. +-- Classic TPC-H analytical patterns plus queries that hit the +-- Databricks-specific augmentations (VARIANT, STRUCT/ARRAY/MAP, +-- liquid clustering, generated column, materialized view). +-- +-- The dashboard substitutes this file into the step-1 prompt, where the +-- agent uses it to choose ORDER BY keys, partitioning, and codecs. + +-- 1. Revenue by ship date and priority — the bread-and-butter rollup. +-- Drives the ORDER BY choice on the migrated lineitem table. +SELECT + l_shipdate, + o_orderpriority, + count(*) AS line_count, + sum(l_extendedprice * (1 - l_discount)) AS revenue +FROM migration_demo.tpch.lineitem l +JOIN migration_demo.tpch.orders o ON l.l_orderkey = o.o_orderkey +WHERE l_shipdate BETWEEN DATE '1994-01-01' AND DATE '1994-12-31' +GROUP BY l_shipdate, o_orderpriority +ORDER BY l_shipdate, o_orderpriority; + +-- 2. Top customers by lifetime revenue — multi-table join where the +-- GROUP BY columns live in dimensions, not the fact table. +SELECT + c.c_custkey, + c.c_name, + n.n_name AS nation, + count(o.o_orderkey) AS order_count, + sum(o.o_totalprice) AS lifetime_revenue +FROM migration_demo.tpch.orders o +JOIN migration_demo.tpch.customer c ON o.o_custkey = c.c_custkey +JOIN migration_demo.tpch.nation n ON c.c_nationkey = n.n_nationkey +GROUP BY c.c_custkey, c.c_name, n.n_name +ORDER BY lifetime_revenue DESC +LIMIT 50; + +-- 3. VARIANT extraction — reads the o_metadata column with Databricks' +-- colon path syntax. On ClickHouse this becomes JSONExtract* over a +-- JSON column, or a typed column if the agent extracted hot keys. +SELECT + o_metadata:channel::string AS channel, + o_metadata:fulfilment.warehouse::string AS warehouse, + count(*) AS order_count, + sum(o_totalprice) AS revenue +FROM migration_demo.tpch.orders +WHERE o_orderdate >= DATE '1995-01-01' +GROUP BY channel, warehouse +ORDER BY revenue DESC; + +-- 4. Nested-type access — explodes the ARRAY shipping events and +-- reads the MAP. `explode` becomes arrayJoin on ClickHouse; the MAP +-- subscript becomes a Map(String, String) lookup. +SELECT + event.status AS status, + l.l_shipmode AS ship_mode, + l.l_attributes['carrier'] AS carrier, + count(*) AS event_count +FROM migration_demo.tpch.lineitem l +LATERAL VIEW explode(l.l_shipping_events) AS event +GROUP BY event.status, l.l_shipmode, l.l_attributes['carrier'] +ORDER BY event_count DESC; + +-- 5. Window function with QUALIFY — Databricks supports QUALIFY, ClickHouse +-- does not, so this must be rewritten as a subquery with WHERE. +-- Also reads o_orderyear, the GENERATED ALWAYS AS column. +SELECT + o_orderyear, + o_orderpriority, + sum(o_totalprice) AS revenue, + rank() OVER ( + PARTITION BY o_orderyear + ORDER BY sum(o_totalprice) DESC + ) AS revenue_rank +FROM migration_demo.tpch.orders +GROUP BY o_orderyear, o_orderpriority +QUALIFY revenue_rank <= 3 +ORDER BY o_orderyear, revenue_rank; + +-- 6. Higher-order function — `aggregate` and `filter` over an array have no +-- direct ClickHouse syntax; they map to arrayReduce / arrayFilter. +SELECT + l_shipmode, + count(*) AS line_count, + avg( + aggregate( + filter(l_shipping_events, e -> e.status <> 'CANCELLED'), + 0, + (acc, e) -> acc + 1 + ) + ) AS avg_live_events +FROM migration_demo.tpch.lineitem +GROUP BY l_shipmode +ORDER BY line_count DESC; + +-- 7. Pre-aggregated read against the materialized view (skipped by the +-- setup script on non-serverless warehouses — if daily_order_summary +-- does not exist, drop this query). +SELECT + order_day, + o_orderpriority, + order_count, + daily_revenue +FROM migration_demo.tpch.daily_order_summary +WHERE order_day BETWEEN DATE '1995-01-01' AND DATE '1995-12-31' +ORDER BY order_day, o_orderpriority; From 905a84a87d673335d5bf2226c76a15957f4293f2 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 17:00:20 +0800 Subject: [PATCH 11/21] feat(databricks): TPC-H workload setup with Delta augmentations Adds setup_workload.sql (TPC-H copy from samples.tpch plus VARIANT, generated column, nested ARRAY/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. --- sources/databricks/scripts/requirements.txt | 1 + sources/databricks/scripts/setup_workload.py | 137 ++++++++++++++++ sources/databricks/scripts/setup_workload.sql | 152 ++++++++++++++++++ tests/test_setup_workload_parser.py | 66 ++++++++ 4 files changed, 356 insertions(+) create mode 100644 sources/databricks/scripts/requirements.txt create mode 100644 sources/databricks/scripts/setup_workload.py create mode 100644 sources/databricks/scripts/setup_workload.sql create mode 100644 tests/test_setup_workload_parser.py diff --git a/sources/databricks/scripts/requirements.txt b/sources/databricks/scripts/requirements.txt new file mode 100644 index 0000000..944e7fc --- /dev/null +++ b/sources/databricks/scripts/requirements.txt @@ -0,0 +1 @@ +databricks-sql-connector>=4.0 diff --git a/sources/databricks/scripts/setup_workload.py b/sources/databricks/scripts/setup_workload.py new file mode 100644 index 0000000..9f852a7 --- /dev/null +++ b/sources/databricks/scripts/setup_workload.py @@ -0,0 +1,137 @@ +""" +MigrationRoom — Databricks workload setup. + +Reads DATABRICKS_* credentials from the environment and executes +setup_workload.sql against the partner's SQL warehouse. Copies the built-in +samples.tpch tables into migration_demo.tpch and adds Databricks-specific +decorations (VARIANT, generated column, ARRAY, MAP, +TIMESTAMP/TIMESTAMP_NTZ, liquid clustering, deletion vectors, +materialized view). + +No data download — TPC-H is already inside every Databricks workspace. + +Usage: + pip install -r sources/databricks/scripts/requirements.txt + set -a; source .env; set +a + python3 sources/databricks/scripts/setup_workload.py + +Environment: + DATABRICKS_HOST required — workspace URL or bare hostname + DATABRICKS_HTTP_PATH required — e.g. /sql/1.0/warehouses/abc123 + DATABRICKS_TOKEN required +""" +import os +import re +import sys +from pathlib import Path + +SQL_FILE = Path(__file__).parent / "setup_workload.sql" + +DIRECTIVE = re.compile(r"^\s*--\s*@(optional|requires)\b[:\s]*(.*)$", re.I) +COMMENT = re.compile(r"^\s*--") + + +def require_env(key: str) -> str: + value = os.environ.get(key) + if not value: + print(f"❌ Missing required env var: {key}", file=sys.stderr) + sys.exit(2) + return value + + +def normalize_host(raw: str) -> str: + host = raw.strip() + return host.removeprefix("https://").removeprefix("http://").rstrip("/") + + +def parse_statements(sql: str) -> list[tuple[str, str, str]]: + """Split the script into (statement, kind, hint) triples. + + `kind` is 'required' when the statement was preceded by a + `-- @requires:` directive, 'optional' for `-- @optional:`, else + 'plain'. A directive applies to the next statement only. + + Naive by design: the setup script has no embedded semicolons in + string literals and no procedural blocks. + """ + statements: list[tuple[str, str, str]] = [] + buffer: list[str] = [] + kind, hint = "plain", "" + + for line in sql.splitlines(): + directive = DIRECTIVE.match(line) + if directive: + kind = "optional" if directive.group(1).lower() == "optional" else "required" + hint = directive.group(2).strip() + continue + if COMMENT.match(line): + continue + buffer.append(line) + if line.rstrip().endswith(";"): + statement = "\n".join(buffer).strip().rstrip(";").strip() + if statement: + statements.append((statement, kind, hint)) + buffer, kind, hint = [], "plain", "" + + tail = "\n".join(buffer).strip().rstrip(";").strip() + if tail: + statements.append((tail, kind, hint)) + return statements + + +def first_line(statement: str, limit: int = 80) -> str: + for line in statement.splitlines(): + line = line.strip() + if line: + return line[:limit] + ("…" if len(line) > limit else "") + return "" + + +def main() -> int: + from databricks import sql as dbsql + + host = normalize_host(require_env("DATABRICKS_HOST")) + http_path = require_env("DATABRICKS_HTTP_PATH") + token = require_env("DATABRICKS_TOKEN") + + statements = parse_statements(SQL_FILE.read_text(encoding="utf-8")) + print(f"Loaded {len(statements)} statements from {SQL_FILE.name}.") + print(f"Connecting to {host} ({http_path})…") + + skipped: list[str] = [] + with dbsql.connect( + server_hostname=host, http_path=http_path, access_token=token + ) as conn: + cur = conn.cursor() + try: + for i, (statement, kind, hint) in enumerate(statements, 1): + print(f" [{i:>2}/{len(statements)}] {first_line(statement)}", flush=True) + try: + cur.execute(statement) + except Exception as exc: + if kind == "optional": + print(f" ⚠️ skipped: {hint or exc}", flush=True) + skipped.append(first_line(statement, 60)) + continue + if kind == "required": + print(f"\n❌ {hint}", file=sys.stderr) + print(f" Statement: {first_line(statement)}", file=sys.stderr) + print(f" Error: {exc}", file=sys.stderr) + return 3 + print(f"\n❌ Statement failed: {first_line(statement)}", file=sys.stderr) + print(f" Error: {exc}", file=sys.stderr) + return 1 + finally: + cur.close() + + print("\n✅ Workload setup complete. migration_demo.tpch is ready.") + if skipped: + print(f"⚠️ {len(skipped)} optional statement(s) skipped:") + for label in skipped: + print(f" - {label}") + print(" Set DATABRICKS_NAMESPACE=migration_demo.tpch in .env.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/sources/databricks/scripts/setup_workload.sql b/sources/databricks/scripts/setup_workload.sql new file mode 100644 index 0000000..71b6e0a --- /dev/null +++ b/sources/databricks/scripts/setup_workload.sql @@ -0,0 +1,152 @@ +-- MigrationRoom — Databricks demo workload. +-- +-- Builds migration_demo.tpch from the read-only `samples.tpch` catalog that +-- ships with every Databricks workspace (~6M rows, no download), then adds +-- Databricks-specific decoration so the migration agent has to make real +-- decisions rather than a mechanical type-for-type copy. +-- +-- Directives read by setup_workload.py: +-- -- @requires: abort with if the next statement fails +-- -- @optional: warn and continue if the next statement fails +-- +-- Why TPC-H rather than a bespoke workload: see docs/adding-a-source.md. +-- Reusing it lets partners compare Databricks, Snowflake, and BigQuery +-- migrations side by side — same tables, only the decoration differs. + +CREATE CATALOG IF NOT EXISTS migration_demo; + +CREATE SCHEMA IF NOT EXISTS migration_demo.tpch; + +-- ── The 8 TPC-H tables, copied from the built-in samples catalog ────── + +CREATE OR REPLACE TABLE migration_demo.tpch.region AS SELECT * FROM samples.tpch.region; +CREATE OR REPLACE TABLE migration_demo.tpch.nation AS SELECT * FROM samples.tpch.nation; +CREATE OR REPLACE TABLE migration_demo.tpch.supplier AS SELECT * FROM samples.tpch.supplier; +CREATE OR REPLACE TABLE migration_demo.tpch.customer AS SELECT * FROM samples.tpch.customer; +CREATE OR REPLACE TABLE migration_demo.tpch.part AS SELECT * FROM samples.tpch.part; +CREATE OR REPLACE TABLE migration_demo.tpch.partsupp AS SELECT * FROM samples.tpch.partsupp; +CREATE OR REPLACE TABLE migration_demo.tpch.orders AS SELECT * FROM samples.tpch.orders; +CREATE OR REPLACE TABLE migration_demo.tpch.lineitem AS SELECT * FROM samples.tpch.lineitem; + +-- ── Augmentation 1: VARIANT on orders ──────────────────────────────── +-- Forces a decision: map to ClickHouse JSON, or extract hot keys into +-- typed columns? VARIANT is native from DBSQL 2024.35 / DBR 15.3. + +-- @requires: VARIANT requires DBSQL 2024.35+ or DBR 15.3+. Upgrade the SQL warehouse channel to Current, or use a newer runtime. +ALTER TABLE migration_demo.tpch.orders ADD COLUMN o_metadata VARIANT; + +UPDATE migration_demo.tpch.orders +SET o_metadata = parse_json( + concat( + '{"channel":"', + element_at(array('web', 'retail', 'partner', 'phone'), cast(pmod(o_orderkey, 4) + 1 AS INT)), + '","fulfilment":{"warehouse":"WH-', + cast(pmod(o_orderkey, 7) + 1 AS STRING), + '","expedited":', + CASE WHEN o_orderpriority LIKE '1-URGENT%' THEN 'true' ELSE 'false' END, + '},"discount_codes":["', + element_at(array('NONE', 'SPRING10', 'LOYALTY5'), cast(pmod(o_orderkey, 3) + 1 AS INT)), + '"]}' + ) +); + +-- ── Augmentation 2: generated column on orders ──────────────────────── +-- Forces a decision: MATERIALIZED column or ALIAS column on ClickHouse? + +-- @requires: GENERATED ALWAYS AS requires Delta Lake with column generation support (DBR 8.3+). +ALTER TABLE migration_demo.tpch.orders +ADD COLUMN o_orderyear INT GENERATED ALWAYS AS (year(o_orderdate)); + +-- ── Augmentation 3: nested types on lineitem ────────────────────────── +-- ARRAY and MAP. Forces a decision between ClickHouse Nested, +-- Array(Tuple(...)), and Map(String, String). + +ALTER TABLE migration_demo.tpch.lineitem +ADD COLUMNS ( + l_shipping_events ARRAY>, + l_attributes MAP +); + +UPDATE migration_demo.tpch.lineitem +SET l_shipping_events = array( + named_struct( + 'status', 'PACKED', + 'event_ts', cast(l_shipdate AS TIMESTAMP), + 'location', concat('WH-', cast(pmod(l_orderkey, 7) + 1 AS STRING)) + ), + named_struct( + 'status', CASE + WHEN pmod(l_orderkey, 10) = 0 THEN 'CANCELLED' + WHEN l_returnflag = 'R' THEN 'RETURNED' + ELSE 'DELIVERED' + END, + 'event_ts', cast(l_receiptdate AS TIMESTAMP), + 'location', concat('DC-', cast(pmod(l_partkey, 5) + 1 AS STRING)) + ) + ), + l_attributes = map( + 'carrier', element_at(array('UPS', 'FEDEX', 'DHL', 'USPS'), cast(pmod(l_orderkey, 4) + 1 AS INT)), + 'fragile', CASE WHEN pmod(l_partkey, 11) = 0 THEN 'true' ELSE 'false' END + ); + +-- ── Augmentation 4: TIMESTAMP vs TIMESTAMP_NTZ on lineitem ──────────── +-- Forces UTC normalisation and a DateTime64 precision choice. + +-- @requires: TIMESTAMP_NTZ requires DBSQL 2023.35+ or DBR 13.3+. +ALTER TABLE migration_demo.tpch.lineitem +ADD COLUMNS ( + l_committed_at TIMESTAMP, + l_committed_at_ntz TIMESTAMP_NTZ +); + +UPDATE migration_demo.tpch.lineitem +SET l_committed_at = cast(l_commitdate AS TIMESTAMP), + l_committed_at_ntz = cast(cast(l_commitdate AS TIMESTAMP) AS TIMESTAMP_NTZ); + +-- ── Augmentation 5: liquid clustering on lineitem ───────────────────── +-- Forces a deliberate ClickHouse ORDER BY choice rather than copying a key. + +-- @requires: Liquid clustering (CLUSTER BY) requires DBR 13.3+ / DBSQL 2023.40+. +ALTER TABLE migration_demo.tpch.lineitem +CLUSTER BY (l_shipdate, l_suppkey); + +-- ── Augmentation 6: deletion vectors + a real history to time-travel ── +-- No ClickHouse equivalent — the agent has to reason about +-- ReplacingMergeTree, ClickPipes, or deferring CDC entirely. + +-- @requires: Deletion vectors require DBR 12.2+ / DBSQL 2023.10+. +ALTER TABLE migration_demo.tpch.lineitem +SET TBLPROPERTIES (delta.enableDeletionVectors = true); + +DELETE FROM migration_demo.tpch.lineitem +WHERE l_orderkey IN ( + SELECT l_orderkey FROM migration_demo.tpch.lineitem LIMIT 500 +); + +-- ── Augmentation 7: materialized view (serverless only) ─────────────── +-- Recreated as a ClickHouse Materialized View on AggregatingMergeTree. +-- Skipped rather than fatal: materialized views need serverless compute, +-- and a classic warehouse is a perfectly reasonable demo environment. + +-- @optional: Materialized views require a serverless SQL warehouse. Skipping — the demo works without it, and sample query 7 should be removed if absent. +CREATE OR REPLACE MATERIALIZED VIEW migration_demo.tpch.daily_order_summary AS +SELECT + o_orderdate AS order_day, + o_orderpriority, + count(*) AS order_count, + sum(o_totalprice) AS daily_revenue +FROM migration_demo.tpch.orders +GROUP BY o_orderdate, o_orderpriority; + +-- ── Table comments, so the agent's discovery step has something to read ── + +COMMENT ON TABLE migration_demo.tpch.orders IS + 'TPC-H orders, augmented with a VARIANT metadata column and a generated year column.'; + +COMMENT ON TABLE migration_demo.tpch.lineitem IS + 'TPC-H lineitem, augmented with nested shipping events, a MAP of attributes, TIMESTAMP/TIMESTAMP_NTZ pair, liquid clustering, and deletion vectors.'; + +-- ── Recompute statistics so DESCRIBE DETAIL reports useful sizes ────── + +ANALYZE TABLE migration_demo.tpch.orders COMPUTE STATISTICS; +ANALYZE TABLE migration_demo.tpch.lineitem COMPUTE STATISTICS; diff --git a/tests/test_setup_workload_parser.py b/tests/test_setup_workload_parser.py new file mode 100644 index 0000000..b4e8e59 --- /dev/null +++ b/tests/test_setup_workload_parser.py @@ -0,0 +1,66 @@ +"""Unit tests for the setup-workload SQL statement/directive parser. + +Run from the repo root: + python3 -m pytest tests/test_setup_workload_parser.py -v +""" +import sys +from pathlib import Path + +sys.path.insert( + 0, str(Path(__file__).resolve().parents[1] / "sources" / "databricks" / "scripts") +) + +from setup_workload import parse_statements # noqa: E402 + + +def test_plain_statements_are_parsed(): + parsed = parse_statements("SELECT 1;\nSELECT 2;\n") + assert [(s, k) for s, k, _ in parsed] == [("SELECT 1", "plain"), ("SELECT 2", "plain")] + + +def test_comments_are_dropped(): + parsed = parse_statements("-- a comment\nSELECT 1;\n") + assert len(parsed) == 1 + assert parsed[0][0] == "SELECT 1" + + +def test_requires_directive_attaches_hint(): + sql = "-- @requires: VARIANT needs DBSQL 2024.35+\nALTER TABLE t ADD COLUMN v VARIANT;\n" + statement, kind, hint = parse_statements(sql)[0] + assert kind == "required" + assert statement.startswith("ALTER TABLE") + assert "2024.35" in hint + + +def test_optional_directive_attaches_hint(): + sql = "-- @optional: needs serverless\nCREATE MATERIALIZED VIEW v AS SELECT 1;\n" + statement, kind, hint = parse_statements(sql)[0] + assert kind == "optional" + assert "serverless" in hint + + +def test_directive_applies_only_to_the_next_statement(): + sql = ( + "-- @optional: first only\n" + "SELECT 1;\n" + "SELECT 2;\n" + ) + parsed = parse_statements(sql) + assert parsed[0][1] == "optional" + assert parsed[1][1] == "plain" + + +def test_multiline_statement_is_kept_together(): + sql = "CREATE TABLE t (\n a INT,\n b INT\n);\n" + parsed = parse_statements(sql) + assert len(parsed) == 1 + assert "a INT" in parsed[0][0] and "b INT" in parsed[0][0] + + +def test_final_statement_without_semicolon_is_kept(): + parsed = parse_statements("SELECT 1") + assert parsed[0][0] == "SELECT 1" + + +def test_empty_input_yields_nothing(): + assert parse_statements("\n-- only a comment\n") == [] From 0c39592e82497d7df3f0bd43489b9152046a95e7 Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 17:30:55 +0800 Subject: [PATCH 12/21] fix(databricks): declare o_orderyear at CREATE TABLE, downsample SF1000 source MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- sources/databricks/scripts/setup_workload.sql | 99 ++++++++++++++----- 1 file changed, 75 insertions(+), 24 deletions(-) diff --git a/sources/databricks/scripts/setup_workload.sql b/sources/databricks/scripts/setup_workload.sql index 71b6e0a..0b8ccf6 100644 --- a/sources/databricks/scripts/setup_workload.sql +++ b/sources/databricks/scripts/setup_workload.sql @@ -1,9 +1,14 @@ -- MigrationRoom — Databricks demo workload. -- -- Builds migration_demo.tpch from the read-only `samples.tpch` catalog that --- ships with every Databricks workspace (~6M rows, no download), then adds --- Databricks-specific decoration so the migration agent has to make real --- decisions rather than a mechanical type-for-type copy. +-- ships with every Databricks workspace, then adds Databricks-specific +-- decoration so the migration agent has to make real decisions rather than +-- a mechanical type-for-type copy. `samples.tpch` itself is scale factor +-- 1000 (~1 TB, ~6 billion lineitem rows, ~1.5 billion orders rows); the +-- table-population predicates below downsample it to SF1-equivalent +-- cardinality (~6M lineitem rows, no download) so this script runs in +-- minutes and the eventual ClickHouse Cloud migration moves megabytes, not +-- terabytes. -- -- Directives read by setup_workload.py: -- -- @requires: abort with if the next statement fails @@ -17,16 +22,69 @@ CREATE CATALOG IF NOT EXISTS migration_demo; CREATE SCHEMA IF NOT EXISTS migration_demo.tpch; --- ── The 8 TPC-H tables, copied from the built-in samples catalog ────── +-- ── The 8 TPC-H tables, downsampled from the built-in samples catalog ── +-- +-- samples.tpch is scale factor 1000 (~1 TB total). The predicates below cut +-- every table down to SF1-equivalent cardinality so this script runs in +-- minutes and the demo migrates megabytes, not terabytes, into ClickHouse +-- Cloud. Do NOT remove them: +-- - orders and lineitem share the same key range (o_orderkey / +-- l_orderkey <= 6000000) on purpose, so every kept line item's order +-- still exists — a LIMIT or TABLESAMPLE on either would orphan rows. +-- - customer is derived from the orders that were actually kept, so +-- every kept order's customer still exists too. +-- - region, nation are tiny at any scale factor and are copied whole. +-- - part, partsupp, supplier are not joined by any demo query, so a +-- cheap independent key-range slice is fine. + +CREATE OR REPLACE TABLE migration_demo.tpch.region AS SELECT * FROM samples.tpch.region; +CREATE OR REPLACE TABLE migration_demo.tpch.nation AS SELECT * FROM samples.tpch.nation; + +CREATE OR REPLACE TABLE migration_demo.tpch.supplier AS +SELECT * FROM samples.tpch.supplier WHERE s_suppkey <= 10000; + +CREATE OR REPLACE TABLE migration_demo.tpch.part AS +SELECT * FROM samples.tpch.part WHERE p_partkey <= 200000; + +CREATE OR REPLACE TABLE migration_demo.tpch.partsupp AS +SELECT * FROM samples.tpch.partsupp WHERE ps_partkey <= 200000; + +-- orders is declared explicitly (rather than `CREATE TABLE ... AS SELECT`) +-- because o_orderyear is a generated column, and Databricks' ALTER TABLE +-- ADD COLUMN grammar has no GENERATED clause — only CREATE TABLE's +-- column_properties does. This is also augmentation "generated column on +-- orders": it forces a decision between a MATERIALIZED column and an ALIAS +-- column on ClickHouse. + +-- @requires: This hard-coded orders schema may not exactly match samples.tpch.orders's column types on this workspace (untested without a live workspace). If this fails, run DESCRIBE TABLE samples.tpch.orders and adjust the column list below to match, then rerun. +CREATE OR REPLACE TABLE migration_demo.tpch.orders ( + o_orderkey BIGINT, + o_custkey BIGINT, + o_orderstatus STRING, + o_totalprice DECIMAL(18,2), + o_orderdate DATE, + o_orderpriority STRING, + o_clerk STRING, + o_shippriority INT, + o_comment STRING, + o_orderyear INT GENERATED ALWAYS AS (year(o_orderdate)) +); + +INSERT INTO migration_demo.tpch.orders ( + o_orderkey, o_custkey, o_orderstatus, o_totalprice, o_orderdate, + o_orderpriority, o_clerk, o_shippriority, o_comment +) +SELECT o_orderkey, o_custkey, o_orderstatus, o_totalprice, o_orderdate, + o_orderpriority, o_clerk, o_shippriority, o_comment +FROM samples.tpch.orders +WHERE o_orderkey <= 6000000; -CREATE OR REPLACE TABLE migration_demo.tpch.region AS SELECT * FROM samples.tpch.region; -CREATE OR REPLACE TABLE migration_demo.tpch.nation AS SELECT * FROM samples.tpch.nation; -CREATE OR REPLACE TABLE migration_demo.tpch.supplier AS SELECT * FROM samples.tpch.supplier; -CREATE OR REPLACE TABLE migration_demo.tpch.customer AS SELECT * FROM samples.tpch.customer; -CREATE OR REPLACE TABLE migration_demo.tpch.part AS SELECT * FROM samples.tpch.part; -CREATE OR REPLACE TABLE migration_demo.tpch.partsupp AS SELECT * FROM samples.tpch.partsupp; -CREATE OR REPLACE TABLE migration_demo.tpch.orders AS SELECT * FROM samples.tpch.orders; -CREATE OR REPLACE TABLE migration_demo.tpch.lineitem AS SELECT * FROM samples.tpch.lineitem; +CREATE OR REPLACE TABLE migration_demo.tpch.customer AS +SELECT * FROM samples.tpch.customer +WHERE c_custkey IN (SELECT o_custkey FROM migration_demo.tpch.orders); + +CREATE OR REPLACE TABLE migration_demo.tpch.lineitem AS +SELECT * FROM samples.tpch.lineitem WHERE l_orderkey <= 6000000; -- ── Augmentation 1: VARIANT on orders ──────────────────────────────── -- Forces a decision: map to ClickHouse JSON, or extract hot keys into @@ -50,14 +108,7 @@ SET o_metadata = parse_json( ) ); --- ── Augmentation 2: generated column on orders ──────────────────────── --- Forces a decision: MATERIALIZED column or ALIAS column on ClickHouse? - --- @requires: GENERATED ALWAYS AS requires Delta Lake with column generation support (DBR 8.3+). -ALTER TABLE migration_demo.tpch.orders -ADD COLUMN o_orderyear INT GENERATED ALWAYS AS (year(o_orderdate)); - --- ── Augmentation 3: nested types on lineitem ────────────────────────── +-- ── Augmentation 2: nested types on lineitem ────────────────────────── -- ARRAY and MAP. Forces a decision between ClickHouse Nested, -- Array(Tuple(...)), and Map(String, String). @@ -89,7 +140,7 @@ SET l_shipping_events = array( 'fragile', CASE WHEN pmod(l_partkey, 11) = 0 THEN 'true' ELSE 'false' END ); --- ── Augmentation 4: TIMESTAMP vs TIMESTAMP_NTZ on lineitem ──────────── +-- ── Augmentation 3: TIMESTAMP vs TIMESTAMP_NTZ on lineitem ──────────── -- Forces UTC normalisation and a DateTime64 precision choice. -- @requires: TIMESTAMP_NTZ requires DBSQL 2023.35+ or DBR 13.3+. @@ -103,14 +154,14 @@ UPDATE migration_demo.tpch.lineitem SET l_committed_at = cast(l_commitdate AS TIMESTAMP), l_committed_at_ntz = cast(cast(l_commitdate AS TIMESTAMP) AS TIMESTAMP_NTZ); --- ── Augmentation 5: liquid clustering on lineitem ───────────────────── +-- ── Augmentation 4: liquid clustering on lineitem ───────────────────── -- Forces a deliberate ClickHouse ORDER BY choice rather than copying a key. -- @requires: Liquid clustering (CLUSTER BY) requires DBR 13.3+ / DBSQL 2023.40+. ALTER TABLE migration_demo.tpch.lineitem CLUSTER BY (l_shipdate, l_suppkey); --- ── Augmentation 6: deletion vectors + a real history to time-travel ── +-- ── Augmentation 5: deletion vectors + a real history to time-travel ── -- No ClickHouse equivalent — the agent has to reason about -- ReplacingMergeTree, ClickPipes, or deferring CDC entirely. @@ -123,7 +174,7 @@ WHERE l_orderkey IN ( SELECT l_orderkey FROM migration_demo.tpch.lineitem LIMIT 500 ); --- ── Augmentation 7: materialized view (serverless only) ─────────────── +-- ── Augmentation 6: materialized view (serverless only) ─────────────── -- Recreated as a ClickHouse Materialized View on AggregatingMergeTree. -- Skipped rather than fatal: materialized views need serverless compute, -- and a classic warehouse is a perfectly reasonable demo environment. From 7a022aeec88a4176b6b240e3c29c068ce63cc13d Mon Sep 17 00:00:00 2001 From: sishuo Date: Fri, 7 Aug 2026 17:41:09 +0800 Subject: [PATCH 13/21] feat(databricks): six dashboard step prompts 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) --- .../prompts/01-discover-and-design.md | 88 ++++++++++++++ sources/databricks/prompts/02-migrate-data.md | 115 ++++++++++++++++++ sources/databricks/prompts/03-validate.md | 44 +++++++ .../databricks/prompts/04-rewrite-queries.md | 56 +++++++++ sources/databricks/prompts/05-benchmark.md | 52 ++++++++ sources/databricks/prompts/06-optimize.md | 20 +++ 6 files changed, 375 insertions(+) create mode 100644 sources/databricks/prompts/01-discover-and-design.md create mode 100644 sources/databricks/prompts/02-migrate-data.md create mode 100644 sources/databricks/prompts/03-validate.md create mode 100644 sources/databricks/prompts/04-rewrite-queries.md create mode 100644 sources/databricks/prompts/05-benchmark.md create mode 100644 sources/databricks/prompts/06-optimize.md 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=