This is an automated email from the ASF dual-hosted git repository. cgivre pushed a commit to branch feat/drill-mcp-server in repository https://gitbox.apache.org/repos/asf/drill-mcp.git
commit 01eac1813273adccf5a1a05364c474c7c6b32c30 Author: cgivre <[email protected]> AuthorDate: Tue Aug 11 17:04:09 2026 -0400 fix: switch SQL guard from Postgres to native Drill dialect sqlglot ships a Drill dialect; using Postgres as a stand-in wrongly rejected idiomatic Drill queries that sqlglot's Postgres grammar cannot parse, e.g. backtick-quoted file-path queries and INFORMATION_SCHEMA lookups. --- .../plans/2026-08-11-drill-mcp-server.md | 27 ++++++++----- .../specs/2026-08-11-drill-mcp-design.md | 34 ++++++++++++++--- drill_mcp/guard.py | 8 ++-- tests/test_guard.py | 44 +++++++++++++++++++--- 4 files changed, 87 insertions(+), 26 deletions(-) diff --git a/docs/superpowers/plans/2026-08-11-drill-mcp-server.md b/docs/superpowers/plans/2026-08-11-drill-mcp-server.md index 2b74791..8410c21 100644 --- a/docs/superpowers/plans/2026-08-11-drill-mcp-server.md +++ b/docs/superpowers/plans/2026-08-11-drill-mcp-server.md @@ -331,7 +331,7 @@ git commit -m "feat: package skeleton and configuration loading" - `drill_mcp.guard.check(sql: str, policy: Policy) -> None` — returns `None` if permitted, raises `PolicyError` otherwise - `drill_mcp.guard.matches_prefix(qualified: str, entries: Iterable[str]) -> bool` -**Background for the implementer:** Drill's SQL is Apache Calcite–based. `sqlglot` has no Drill dialect; the Postgres dialect is the closest fit and is what this module uses. `sqlglot.parse()` returns a *list* of statements, which is how statement-stacking is detected. Constructs sqlglot does not recognize become `exp.Command`, a catch-all node holding the leading keyword and the raw remainder — so `exp.Command` must never be blanket-allowed. +**Background for the implementer:** Drill's SQL is Apache Calcite–based, and `sqlglot` ships a native `Drill` dialect — parse with `read="drill"`. Do not substitute a near neighbour: under the Postgres dialect `sqlglot` raises `ParseError` on backtick-quoted identifiers, and because the guard rejects whatever it cannot parse, that silently turns `SELECT * FROM dfs.`/path/file.csv`` and `SELECT * FROM INFORMATION_SCHEMA.`TABLES`` into policy rejections. `sqlglot.parse()` returns a *list* [...] - [ ] **Step 1: Write a characterization test for sqlglot's parse tree** @@ -351,31 +351,31 @@ class TestSqlglotAssumptions: """Characterization tests: what the guard relies on sqlglot doing.""" def test_parse_returns_one_statement_per_semicolon(self): - assert len(sqlglot.parse("SELECT 1; SELECT 2", read="postgres")) == 2 + assert len(sqlglot.parse("SELECT 1; SELECT 2", read="drill")) == 2 def test_select_parses_to_select(self): - stmt = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="postgres") + stmt = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="drill") assert isinstance(stmt, exp.Select) def test_table_exposes_catalog_db_name(self): - table = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="postgres").find(exp.Table) + table = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="drill").find(exp.Table) assert table.catalog == "dfs" assert table.db == "tmp" assert table.name == "foo" def test_two_part_name_populates_db_not_catalog(self): - table = sqlglot.parse_one("SELECT * FROM sys.options", read="postgres").find(exp.Table) + table = sqlglot.parse_one("SELECT * FROM sys.options", read="drill").find(exp.Table) assert table.catalog == "" assert table.db == "sys" assert table.name == "options" def test_comments_are_stripped_by_the_tokenizer(self): - stmt = sqlglot.parse_one("-- CREATE TABLE evil\nSELECT 1", read="postgres") + stmt = sqlglot.parse_one("-- CREATE TABLE evil\nSELECT 1", read="drill") assert isinstance(stmt, exp.Select) def test_ctas_target_is_reachable_from_this(self): stmt = sqlglot.parse_one( - "CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.raw.src", read="postgres" + "CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.raw.src", read="drill" ) assert isinstance(stmt, exp.Create) target = stmt.this.this if isinstance(stmt.this, exp.Schema) else stmt.this @@ -550,7 +550,7 @@ from dataclasses import dataclass import sqlglot from sqlglot import exp -DIALECT = "postgres" # closest available fit for Drill's Calcite SQL +DIALECT = "drill" # sqlglot ships a native Drill dialect — do not substitute a near neighbour # Commands sqlglot does not model as expressions, but which cannot write. _SAFE_COMMANDS = {"SHOW", "DESCRIBE", "DESC", "EXPLAIN"} @@ -824,8 +824,15 @@ def test_leaves_innocuous_keys_alone(): def test_recurses_into_nested_dicts(): - source = {"config": {"credentialsProvider": {"awsSecretAccessKey": "s"}}} - assert redact(source)["config"]["credentialsProvider"]["awsSecretAccessKey"] == REDACTED + # Nested under a neutral container on purpose: a key like `credentialsProvider` + # is itself sensitive-named and gets blanked wholesale, so it cannot double as + # the vehicle for proving recursion. + source = {"config": {"aws": {"awsSecretAccessKey": "s"}}} + assert redact(source)["config"]["aws"]["awsSecretAccessKey"] == REDACTED + + +def test_blanks_a_credentials_provider_block_wholesale(): + assert redact({"credentialsProvider": {"clientID": "x"}})["credentialsProvider"] == REDACTED def test_recurses_into_lists(): diff --git a/docs/superpowers/specs/2026-08-11-drill-mcp-design.md b/docs/superpowers/specs/2026-08-11-drill-mcp-design.md index 7196861..56f5b75 100644 --- a/docs/superpowers/specs/2026-08-11-drill-mcp-design.md +++ b/docs/superpowers/specs/2026-08-11-drill-mcp-design.md @@ -88,9 +88,25 @@ window. access keys, JDBC passwords, and OAuth tokens. Before returning, the server walks the config tree and replaces the value of any key matching a redaction pattern (`password`, `secret`, `accessKey`, `access_key`, `token`, `credential`, -`privateKey`, case-insensitive) with `"***REDACTED***"`. Redaction is applied -recursively, including inside `credentialsProvider` blocks and nested -`workspaces` entries. +`privateKey`, `apiKey`, `authorization`, `passphrase`, `keytab`, `principal`, +case-insensitive) with `"***REDACTED***"`. Redaction recurses through nested +dicts, lists, and tuples — `workspaces` entries included. + +A key whose *name* matches is blanked wholesale rather than walked, so an +unrecognized child key under a recognized parent cannot leak. That means a +`credentialsProvider` block is replaced entirely rather than descended into: it +is by definition all credentials, and walking it would make safety depend on +every child key name matching, which nothing enforces across third-party +provider implementations. `authorization` is on the list because Drill's HTTP +plugin carries bearer tokens in a `headers` block. + +The pattern is matched as a substring of the key, so `fs.s3a.secret.key` and +`awsSecretAccessKey` are both caught. Resist narrowing it with lookaheads: a +lookahead silently generalizes to key names nobody enumerated. An early +implementation excluded `credentialsProvider` with a `(?![a-z])` lookahead and +thereby also un-redacted `credentialsJson`, `credentialsB64`, and +`serviceAccountCredentialsJson` — the standard shapes for an inlined GCP +service-account private key. This is a trust boundary: MCP tool output goes to a model and often to a third-party API. Redaction is not configurable off. @@ -132,9 +148,15 @@ writable_plugins: [] # e.g. [dfs.tmp] statement stacking. The guard is the only thing standing between a model and the user's data, so it gets a real parser. -Drill's SQL is Calcite-based; `sqlglot`'s Postgres dialect is the closest fit. -Where a legitimate Drill query fails to parse, the failure is a rejection with a -clear message naming `guard.py` — a false negative that blocks a read is +`sqlglot` ships a native `Drill` dialect, and the guard parses with +`read="drill"`. Use it rather than approximating with a near neighbour: under +the Postgres dialect, `sqlglot` raises `ParseError` on backtick-quoted +identifiers, and since the guard rejects whatever it cannot parse, that turned +`SELECT * FROM dfs.`/path/file.csv`` — the most idiomatic Drill query there is — +into a policy rejection. + +Where a legitimate Drill query still fails to parse, the failure is a rejection +with a clear message naming `guard.py` — a false negative that blocks a read is acceptable; a false positive that permits a write is not. ## Hidden schemas diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index 58d1596..9eb0066 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -37,7 +37,7 @@ from dataclasses import dataclass import sqlglot from sqlglot import exp -DIALECT = "postgres" # closest available fit for Drill's Calcite SQL +DIALECT = "drill" # sqlglot ships a native Drill dialect (sqlglot.dialects.drill.Drill). # Commands sqlglot does not model as expressions, but which cannot write. # EXPLAIN is handled separately (see _check_write): it is not blanket-safe @@ -48,9 +48,9 @@ _READ_TYPES = (exp.Select, exp.Union, exp.Intersect, exp.Except, exp.Subquery, e # Node types that indicate a write is embedded somewhere inside a statement # whose root node is a read type (e.g. `WITH x AS (INSERT ...) SELECT * FROM x`, -# or Postgres-dialect `SELECT ... INTO`). Checking only the root type is not -# enough: the safety property must not depend on Drill's parser being any -# narrower than sqlglot's Postgres dialect. +# or `SELECT ... INTO ...`). Checking only the root type is not enough: the +# safety property must not depend on sqlglot's Drill grammar rejecting these +# forms outright — a write hidden deeper in the tree must still be caught. _EMBEDDED_WRITE_TYPES = (exp.Insert, exp.Update, exp.Delete, exp.Merge, exp.Create, exp.Drop, exp.Into) # EXPLAIN unwraps its body and re-checks it recursively; this bounds diff --git a/tests/test_guard.py b/tests/test_guard.py index a059b6f..705ea90 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -31,31 +31,31 @@ class TestSqlglotAssumptions: """Characterization tests: what the guard relies on sqlglot doing.""" def test_parse_returns_one_statement_per_semicolon(self): - assert len(sqlglot.parse("SELECT 1; SELECT 2", read="postgres")) == 2 + assert len(sqlglot.parse("SELECT 1; SELECT 2", read="drill")) == 2 def test_select_parses_to_select(self): - stmt = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="postgres") + stmt = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="drill") assert isinstance(stmt, exp.Select) def test_table_exposes_catalog_db_name(self): - table = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="postgres").find(exp.Table) + table = sqlglot.parse_one("SELECT * FROM dfs.tmp.foo", read="drill").find(exp.Table) assert table.catalog == "dfs" assert table.db == "tmp" assert table.name == "foo" def test_two_part_name_populates_db_not_catalog(self): - table = sqlglot.parse_one("SELECT * FROM sys.options", read="postgres").find(exp.Table) + table = sqlglot.parse_one("SELECT * FROM sys.options", read="drill").find(exp.Table) assert table.catalog == "" assert table.db == "sys" assert table.name == "options" def test_comments_are_stripped_by_the_tokenizer(self): - stmt = sqlglot.parse_one("-- CREATE TABLE evil\nSELECT 1", read="postgres") + stmt = sqlglot.parse_one("-- CREATE TABLE evil\nSELECT 1", read="drill") assert isinstance(stmt, exp.Select) def test_ctas_target_is_reachable_from_this(self): stmt = sqlglot.parse_one( - "CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.raw.src", read="postgres" + "CREATE TABLE dfs.tmp.out AS SELECT * FROM dfs.raw.src", read="drill" ) assert isinstance(stmt, exp.Create) target = stmt.this.this if isinstance(stmt.this, exp.Schema) else stmt.this @@ -292,6 +292,34 @@ class TestHiddenSchemas: check("SELECT 'sys.options' AS note FROM dfs.tmp.a", HIDDEN) +class TestDrillDialectRegressions: + """Regression coverage for the Postgres -> Drill dialect switch. + + Under the Postgres dialect these idiomatic Drill query forms failed to + parse and were wrongly rejected: backtick identifier quoting (Drill's + standard quoting) and querying a file directly by path, the single most + common Drill query shape. The Drill dialect parses both. + """ + + def test_file_path_query_with_backtick_quoting_is_permitted(self): + check("SELECT * FROM dfs.`/path/to/file.csv`", CLOSED) + + def test_information_schema_with_backtick_quoting_is_permitted(self): + check("SELECT * FROM INFORMATION_SCHEMA.`TABLES`", CLOSED) + + def test_backtick_quoted_write_target_rejected_by_default(self): + with pytest.raises(PolicyError): + check("CREATE TABLE dfs.tmp.`out` AS SELECT 1", CLOSED) + + def test_backtick_quoted_write_target_permitted_with_allowlist(self): + check("CREATE TABLE dfs.tmp.`out` AS SELECT 1", OPEN) + + def test_backtick_quoted_hidden_schema_still_caught(self): + hidden = Policy(hidden_schemas=("sys",)) + with pytest.raises(PolicyError, match="hidden"): + check("SELECT * FROM `sys`.options", hidden) + + class TestCoverageGaps: """Exercises branches not reached by the scenarios above, so guard.py stays at 100% line coverage without weakening any other test. @@ -316,6 +344,10 @@ class TestCoverageGaps: with pytest.raises(PolicyError, match="not permitted"): check("CREATE SCHEMA foo", CLOSED) + def test_drop_of_unsupported_kind_rejected(self): + with pytest.raises(PolicyError, match="not permitted"): + check("DROP SCHEMA foo", CLOSED) + def test_ctas_with_explicit_column_list_permitted(self): # The parenthesized column list wraps the target table in an # exp.Schema node; _write_target must unwrap it to find the table.
