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 65148076b07da6451087a82e520817d3cf9bb9d2 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 12:57:12 2026 -0400 fix: stop classifying SHOW spellings, filter all SHOW output instead is_show_schemas's comment-stripping regex was itself bypassable: non-greedy /\*.*?\*/ matching stops at the first */, and a stray */ with no opener strips nothing, so 'SHOW /* /* nested */ */ SCHEMAS' and 'SHOW */ SCHEMAS' both leaked hidden schemas unfiltered. This is the third classifier over the SHOW text to fail the same way (raw regex, literal equality, comment stripping) because Drill's SHOW grammar has no dedicated sqlglot node and everything after the keyword collapses into one opaque literal with no reliable positive signal for 'this is a schema listing'. Replace classification with guard.is_show_command(sql), which checks only the already-parsed Command.this == 'SHOW' field and filters every SHOW command's first column when hidden_schemas is configured. SHOW TABLES/SHOW FILES rows are now incidentally filtered too (a table literally named after a hidden schema would be dropped), accepted as the fail-closed trade-off in place of a classifier that can leak. --- drill_mcp/guard.py | 61 +++++++++++++-------------- drill_mcp/server.py | 18 +++++--- tests/test_guard.py | 112 ++++++++++++++++++++++++-------------------------- tests/test_server.py | 114 ++++++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 196 insertions(+), 109 deletions(-) diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index cf2595e..2f0feb8 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -90,31 +90,36 @@ def matches_prefix(qualified: str, entries: Iterable[str]) -> bool: return False -def is_show_schemas(sql: str) -> bool: - """True if `sql` is `SHOW SCHEMAS` or `SHOW DATABASES`, with or without a - trailing `LIKE '...'` clause. - - Detected from the parsed statement, not a regex over the raw text: a - leading comment (`/* x */ SHOW SCHEMAS`) defeats a `^\\s*SHOW` anchor - because comments are only stripped by the tokenizer, not by string - matching. Drill's grammar for both spellings falls back to sqlglot's - generic `Command`, with the *entire remainder* of the statement left as - one `Literal` in `expression` (e.g. `Command(this='SHOW', - expression=Literal(this='SCHEMAS'))`, but equally - `Literal(this="SCHEMAS LIKE '%dfs%'")` or `Literal(this='/**/SCHEMAS')` - for `SHOW/**/SCHEMAS`). Comparing that literal whole, or even - stripped-and-uppercased, only matches the bare two-word spelling and lets - `SHOW SCHEMAS LIKE '...'` — documented Drill syntax — through unfiltered. - Only the first token decides it: strip block comments (the only comment - style that can land inside the remainder; a leading `--` comment is - stripped by the tokenizer before this text is ever assembled, and a - trailing `--` comment lands after the first token so it never affects - the check), then split on whitespace and compare the first word alone. +def is_show_command(sql: str) -> bool: + """True if `sql` is any `SHOW ...` statement (`SCHEMAS`, `DATABASES`, + `TABLES`, `FILES`, ...). + + This used to try to positively identify `SHOW SCHEMAS`/`SHOW DATABASES` + specifically, by inspecting the text of whatever sqlglot left in the + Command's `expression` literal. Three attempts at that (a regex over the + raw SQL, exact-equality against the literal, comparing only its first + token after stripping `/* */` comments) all failed the same way: Drill's + `SHOW` grammar has no dedicated sqlglot node, so everything after the + keyword — including `LIKE '...'` clauses, embedded comments, and + unbalanced comment delimiters like `SHOW */ SCHEMAS` — collapses into one + opaque literal, and any classifier over that text has some spelling it + fails to recognise. A classifier that fails open on the shapes it does + not recognise is the wrong shape entirely for a security filter: it must + fail closed instead. + + So this checks only whether the statement is a `SHOW` command at all — + one already-parsed field (`statement.this == "SHOW"`), no text parsing. + The caller filters the first column of *any* `SHOW` command's result set + when hidden_schemas is configured, accepting that `SHOW TABLES`/`SHOW + FILES` rows are incidentally run through the same filter (a table or file + literally named `sys`, for a `hidden_schemas: [sys]` policy, would be + dropped) in exchange for the property that no `SHOW` spelling can leak a + hidden schema by evading a classifier. Returns False rather than raising if `sql` does not parse: by the time - this is called, `check()` has already accepted the statement, so a - parse failure here would be a bug in this function, not a policy - decision to surface. + this is called, `check()` has already accepted the statement, so a parse + failure here would be a bug in this function, not a policy decision to + surface. """ try: statements = [s for s in sqlglot.parse(sql, read=DIALECT) if s is not None] @@ -123,15 +128,7 @@ def is_show_schemas(sql: str) -> bool: if len(statements) != 1: return False statement = statements[0] - if not isinstance(statement, exp.Command): - return False - if str(statement.this or "").upper() != "SHOW": - return False - expression = statement.args.get("expression") - target = expression.this if isinstance(expression, exp.Literal) else str(expression or "") - without_comments = re.sub(r"/\*.*?\*/", " ", str(target or ""), flags=re.S) - words = without_comments.strip().split(maxsplit=1) - return bool(words) and words[0].upper() in {"SCHEMAS", "DATABASES"} + return isinstance(statement, exp.Command) and str(statement.this or "").upper() == "SHOW" def check(sql: str, policy: Policy) -> None: diff --git a/drill_mcp/server.py b/drill_mcp/server.py index b5403f3..d9262c9 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -30,7 +30,7 @@ from typing import Any from .client_rest import DrillError from .config import Config -from .guard import Policy, PolicyError, check, is_show_schemas, matches_prefix +from .guard import Policy, PolicyError, check, is_show_command, matches_prefix class ToolError(Exception): @@ -101,11 +101,19 @@ class DrillTools: except DrillError as exc: raise ToolError(str(exc)) from exc - # `SHOW SCHEMAS` / `SHOW DATABASES` are evaluated server-side by - # Drill, so the guard cannot filter them by rewriting or rejecting - # the query; their rows are filtered here on the way back instead. + # SHOW commands are evaluated server-side by Drill, so the guard + # cannot filter them by rewriting or rejecting the query; their rows + # are filtered here on the way back instead. This filters *every* + # SHOW command's first column, not just SHOW SCHEMAS/DATABASES: there + # is no reliable way to positively identify which SHOW spelling + # names a schema (see guard.is_show_command's docstring for the + # three narrower approaches that each leaked hidden schemas through + # some spelling). Filtering all SHOW output means SHOW TABLES/SHOW + # FILES rows are incidentally filtered too — a table or file that + # happens to be named after a hidden schema gets dropped — which is + # an acceptable, fail-closed trade-off; a leaked schema is not. rows = result.rows - if self._policy.hidden_schemas and is_show_schemas(sql): + if self._policy.hidden_schemas and is_show_command(sql): rows = [row for row in rows if self._visible(_first_value(row))] payload: dict[str, Any] = { diff --git a/tests/test_guard.py b/tests/test_guard.py index 4ca20d8..a1a6feb 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -24,7 +24,7 @@ import sqlglot from sqlglot import exp import drill_mcp.guard as guard_module -from drill_mcp.guard import Policy, PolicyError, check, is_show_schemas, matches_prefix +from drill_mcp.guard import Policy, PolicyError, check, is_show_command, matches_prefix class TestSqlglotAssumptions: @@ -328,86 +328,82 @@ class TestDrillDialectRegressions: check("SELECT * FROM `sys`.options", hidden) -class TestIsShowSchemas: - """Detection is parser-based, not a regex over the raw SQL string, so a - leading comment (which the tokenizer strips) cannot defeat it. +class TestIsShowCommand: + """`is_show_command` deliberately does not try to distinguish `SHOW + SCHEMAS`/`SHOW DATABASES` from other SHOW spellings: three attempts at + that classification (a raw regex, exact literal equality, a + comment-stripping regex) each leaked hidden schemas through some spelling + the classifier failed to recognise (a LIKE clause, an embedded comment, + an unbalanced `*/`). This is a single already-parsed-field check — + `statement.this == "SHOW"` — with no text parsing left to get wrong, so + it cannot fail open the way a classifier can. """ def test_show_schemas_matches(self): - assert is_show_schemas("SHOW SCHEMAS") is True + assert is_show_command("SHOW SCHEMAS") is True def test_show_databases_matches(self): - assert is_show_schemas("SHOW DATABASES") is True + assert is_show_command("SHOW DATABASES") is True + + def test_show_tables_matches(self): + """SHOW TABLES is also a SHOW command; the caller filters its rows + too, as the fail-closed trade-off (see server.run_query).""" + assert is_show_command("SHOW TABLES") is True + + def test_show_files_matches(self): + assert is_show_command("SHOW FILES IN dfs.tmp") is True def test_lowercase_matches(self): - assert is_show_schemas("show schemas") is True + assert is_show_command("show schemas") is True def test_mixed_case_matches(self): - assert is_show_schemas("ShOw DaTaBaSeS") is True + assert is_show_command("ShOw DaTaBaSeS") is True def test_leading_block_comment_still_matches(self): - assert is_show_schemas("/* x */ SHOW SCHEMAS") is True + assert is_show_command("/* x */ SHOW SCHEMAS") is True def test_leading_line_comment_still_matches(self): - assert is_show_schemas("-- comment\nSHOW DATABASES") is True + assert is_show_command("-- comment\nSHOW DATABASES") is True + + def test_like_clause_still_matches(self): + assert is_show_command("SHOW SCHEMAS LIKE '%dfs%'") is True + + def test_trailing_block_comment_still_matches(self): + assert is_show_command("SHOW SCHEMAS /* trailing */") is True + + def test_no_space_before_comment_still_matches(self): + assert is_show_command("SHOW/**/SCHEMAS") is True + + def test_trailing_semicolon_still_matches(self): + assert is_show_command("SHOW SCHEMAS;") is True + + def test_trailing_line_comment_still_matches(self): + assert is_show_command("SHOW SCHEMAS -- trailing comment") is True - def test_show_tables_does_not_match(self): - assert is_show_schemas("SHOW TABLES") is False + def test_nested_block_comment_still_matches(self): + """This is the specific input that defeated the comment-stripping + regex (`re.sub(r"/\\*.*?\\*/", " ", ...)`): non-greedy matching stops + at the first `*/`, leaving `*/ SCHEMAS` behind. `is_show_command` + never inspects that text at all, so it is unaffected.""" + assert is_show_command("SHOW /* /* nested */ */ SCHEMAS") is True - def test_show_files_does_not_match(self): - assert is_show_schemas("SHOW FILES IN dfs.tmp") is False + def test_unbalanced_comment_delimiter_still_matches(self): + """The other input that defeated the comment-stripping regex: a + stray `*/` with no opener strips nothing, so the first token became + `*/` instead of `SCHEMAS`.""" + assert is_show_command("SHOW */ SCHEMAS") is True def test_ordinary_select_does_not_match(self): - assert is_show_schemas("SELECT * FROM dfs.tmp.notes") is False + assert is_show_command("SELECT * FROM dfs.tmp.notes") is False def test_select_mentioning_show_as_a_string_does_not_match(self): - assert is_show_schemas("SELECT 'SHOW SCHEMAS' AS x") is False + assert is_show_command("SELECT 'SHOW SCHEMAS' AS x") is False def test_unparseable_sql_returns_false_rather_than_raising(self): - assert is_show_schemas("((((((") is False + assert is_show_command("((((((") is False def test_empty_sql_returns_false(self): - assert is_show_schemas("") is False - - def test_show_schemas_with_like_clause_matches(self): - """Regression test: sqlglot's Command fallback swallows the entire - remainder after SHOW into one literal, so `SCHEMAS LIKE '%dfs%'` must - be recognised by its first token, not by comparing the literal as a - whole (which only matches the bare two-word spelling).""" - assert is_show_schemas("SHOW SCHEMAS LIKE '%dfs%'") is True - - def test_show_databases_with_like_clause_matches(self): - assert is_show_schemas("SHOW DATABASES LIKE '%y%'") is True - - def test_show_tables_with_like_clause_does_not_match(self): - assert is_show_schemas("SHOW TABLES LIKE '%x%'") is False - - def test_show_schemas_with_trailing_block_comment_matches(self): - assert is_show_schemas("SHOW SCHEMAS /* trailing */") is True - - def test_show_schemas_with_no_space_before_comment_matches(self): - assert is_show_schemas("SHOW/**/SCHEMAS") is True - - def test_show_schemas_with_trailing_semicolon_matches(self): - assert is_show_schemas("SHOW SCHEMAS;") is True - - def test_show_schemas_with_trailing_line_comment_matches(self): - assert is_show_schemas("SHOW SCHEMAS -- trailing comment") is True - - def test_raw_regex_would_have_missed_the_like_clause(self): - """Documents the exact failure mode this class regression-tests: the - brief's original `^\\s*SHOW\\s+(SCHEMAS|DATABASES)\\b` regex matches - the bare spelling but the naive fix of comparing the whole Command - literal against {"SCHEMAS", "DATABASES"} also fails on this input, - because the literal for `SHOW SCHEMAS LIKE '%dfs%'` is the full - string `"SCHEMAS LIKE '%dfs%'"`, not `"SCHEMAS"`. - """ - whole_literal_equality = "SCHEMAS LIKE '%dfs%'".strip().upper() in { - "SCHEMAS", - "DATABASES", - } - assert whole_literal_equality is False - assert is_show_schemas("SHOW SCHEMAS LIKE '%dfs%'") is True + assert is_show_command("") is False class TestCoverageGaps: diff --git a/tests/test_server.py b/tests/test_server.py index a904510..8139ca3 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -287,7 +287,19 @@ class TestManagementTools: class TestShowFiltering: - """SHOW is evaluated server-side by Drill, so rows are filtered on return.""" + """SHOW is evaluated server-side by Drill, so rows are filtered on return. + + Filtering applies to *every* SHOW command's first column, not just SHOW + SCHEMAS/SHOW DATABASES. Three narrower attempts at recognising only the + schema-listing spellings (a raw regex over the SQL text, exact equality + against the parsed Command's literal, a comment-stripping regex over that + literal) each leaked hidden schemas through some spelling the classifier + failed to recognise. Filtering all SHOW output instead means SHOW + TABLES/SHOW FILES rows are incidentally filtered too, but that direction + of failure — over-filtering a table that happens to share a name with a + hidden schema — is the safe one; leaking the schema list is not. See + guard.is_show_command's docstring for the full history. + """ def test_show_schemas_rows_are_filtered(self): client = MagicMock() @@ -368,34 +380,39 @@ class TestShowFiltering: result = make_tools(client, hidden_schemas=["sys"]).run_query("ShOw ScHeMaS") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] - def test_show_tables_rows_are_not_filtered(self): - """SHOW TABLES rows are table names, not schema names; filtering them - against hidden_schemas would be wrong.""" + def test_show_tables_rows_are_now_filtered_too(self): + """Inverted deliberately: SHOW TABLES rows are table names, not + schema names, so filtering them against hidden_schemas can drop a + table that happens to share a name with a hidden schema. That is the + accepted, fail-closed trade-off — SHOW output is filtered as a whole + because no reliable way exists to single out just the + schema-listing spellings of SHOW without risking a leak (see the + class docstring). A table literally named "sys" is rare; a leaked + hidden schema list is not an acceptable alternative.""" client = MagicMock() client.query.return_value = QueryResult( - ["TABLE_NAME"], [{"TABLE_NAME": "sys"}] + ["TABLE_NAME"], [{"TABLE_NAME": "sys"}, {"TABLE_NAME": "orders"}] ) result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW TABLES") - assert result["rows"] == [{"TABLE_NAME": "sys"}] + assert result["rows"] == [{"TABLE_NAME": "orders"}] - def test_show_tables_like_rows_are_not_filtered(self): - """A trailing LIKE clause must not make SHOW TABLES look like a schema - listing either.""" + def test_show_tables_like_rows_are_filtered_too(self): client = MagicMock() client.query.return_value = QueryResult( - ["TABLE_NAME"], [{"TABLE_NAME": "sys"}] + ["TABLE_NAME"], [{"TABLE_NAME": "sys"}, {"TABLE_NAME": "orders"}] ) result = make_tools(client, hidden_schemas=["sys"]).run_query( "SHOW TABLES LIKE '%s%'" ) - assert result["rows"] == [{"TABLE_NAME": "sys"}] + assert result["rows"] == [{"TABLE_NAME": "orders"}] def test_show_schemas_like_rows_are_filtered(self): """Regression test: `SHOW SCHEMAS LIKE '...'` is documented Drill syntax. sqlglot's Command fallback swallows the whole remainder - (`SCHEMAS LIKE '%dfs%'`) into a single literal, so comparing that - literal whole (or even stripped-and-uppercased) only matches the bare - two-word spelling and lets this form through unfiltered.""" + (`SCHEMAS LIKE '%dfs%'`) into a single literal, so any classifier + that inspects that text specifically has some spelling it misses; + filtering every SHOW command sidesteps the classification problem + entirely.""" client = MagicMock() client.query.return_value = QueryResult( ["SCHEMA_NAME"], @@ -439,6 +456,44 @@ class TestShowFiltering: ) assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + def test_show_schemas_with_trailing_semicolon_is_still_filtered(self): + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], + [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW SCHEMAS;") + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + + def test_show_schemas_with_nested_block_comment_is_still_filtered(self): + """Regression test for the specific input that defeated the + comment-stripping regex fix: non-greedy `/\\*.*?\\*/` matching stops + at the first `*/`, leaving `*/ SCHEMAS` behind, so the "first token" + became `*/` instead of `SCHEMAS` and the row leaked.""" + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], + [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query( + "SHOW /* /* nested */ */ SCHEMAS" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + + def test_show_schemas_with_unbalanced_comment_delimiter_is_still_filtered(self): + """Regression test for the other input that defeated the + comment-stripping regex: a stray `*/` with no opener strips nothing + at all, so the row leaked.""" + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], + [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query( + "SHOW */ SCHEMAS" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + def test_row_that_is_not_a_dict_does_not_crash_filtering(self): """A malformed row must not raise a raw AttributeError out of the filter; it should simply be treated as not identifiable and dropped @@ -449,3 +504,34 @@ class TestShowFiltering: ) result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW SCHEMAS") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + + @pytest.mark.parametrize( + "sql", + [ + "SHOW SCHEMAS", + "SHOW DATABASES", + "/* x */ SHOW SCHEMAS", + "-- c\nSHOW DATABASES", + "SHOW SCHEMAS LIKE '%dfs%'", + "SHOW SCHEMAS /* t */", + "SHOW/**/SCHEMAS", + "SHOW /* /* nested */ */ SCHEMAS", + "SHOW */ SCHEMAS", + "show schemas", + "SHOW SCHEMAS;", + ], + ) + def test_hidden_row_never_appears_in_the_tool_payload(self, sql): + """End-to-end assertion on the tool's actual output, independent of + how detection is implemented: for every spelling that previously + leaked through one of the three narrower classifiers, the hidden + schema must not appear anywhere in the returned payload, not just in + a specific `rows` shape.""" + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], + [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query(sql) + assert {"SCHEMA_NAME": "sys"} not in result["rows"] + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}]
