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 0d747a21370a62e59b406b3885c991ee4dc1f83b Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 10:04:22 2026 -0400 fix: close SHOW SCHEMAS LIKE bypass and harden filter against bad input is_show_schemas compared the whole Command literal for equality against {SCHEMAS, DATABASES}, but sqlglot's Command fallback swallows the entire remainder of a SHOW statement into that literal, so SHOW SCHEMAS LIKE '...', a trailing block comment, or SHOW/**/SCHEMAS all defeated it and let hidden schemas through unfiltered. Detection now strips block comments and compares only the first token. Also hardens _first_value and list_storage_plugins against non-dict rows/ entries instead of raising, makes _visible fail closed on None, and covers all five management tools (not just cluster_status) against a client missing the REST endpoints. --- drill_mcp/guard.py | 21 ++++++++-- drill_mcp/server.py | 17 ++++++-- tests/test_guard.py | 40 +++++++++++++++++++ tests/test_server.py | 108 ++++++++++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 178 insertions(+), 8 deletions(-) diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index f4a1eb1..cf2595e 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -91,14 +91,25 @@ def matches_prefix(qualified: str, entries: Iterable[str]) -> bool: def is_show_schemas(sql: str) -> bool: - """True if `sql` is `SHOW SCHEMAS` or `SHOW DATABASES`. + """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 target left as a `Literal` in `expression` - (e.g. `Command(this='SHOW', expression=Literal(this='SCHEMAS'))`). + 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. Returns False rather than raising if `sql` does not parse: by the time this is called, `check()` has already accepted the statement, so a @@ -118,7 +129,9 @@ def is_show_schemas(sql: str) -> bool: return False expression = statement.args.get("expression") target = expression.this if isinstance(expression, exp.Literal) else str(expression or "") - return str(target or "").strip().upper() in {"SCHEMAS", "DATABASES"} + 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"} def check(sql: str, policy: Policy) -> None: diff --git a/drill_mcp/server.py b/drill_mcp/server.py index b815e12..b5403f3 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -37,7 +37,9 @@ class ToolError(Exception): """The single error type surfaced to MCP clients. Never carries a traceback.""" -def _first_value(row: dict[str, Any]) -> str | None: +def _first_value(row: Any) -> str | None: + if not isinstance(row, dict): + return None for value in row.values(): return str(value) if value is not None else None return None @@ -62,7 +64,14 @@ class DrillTools: raise ToolError(f"schema '{schema}' is hidden by configuration") def _visible(self, schema: str | None) -> bool: - return not matches_prefix(schema or "", self._policy.hidden_schemas) + # Fail closed, not open: an item this function cannot identify (no + # name at all) is filtered out rather than shown by default. Drill + # never actually returns a null schema/plugin name, but this + # function's whole job is filtering, so its default on a value it + # cannot classify must not be "keep it". + if schema is None: + return False + return not matches_prefix(schema, self._policy.hidden_schemas) # -- query and metadata tools ------------------------------------------- @@ -153,7 +162,9 @@ class DrillTools: plugins = self._require_management("storage_plugins")() except DrillError as exc: raise ToolError(str(exc)) from exc - return [p for p in plugins if self._visible(p.get("name"))] + return [ + p for p in plugins if isinstance(p, dict) and self._visible(p.get("name")) + ] def cluster_status(self) -> dict[str, Any]: """Report Drillbit membership and overall cluster status.""" diff --git a/tests/test_guard.py b/tests/test_guard.py index 0b5f32a..4ca20d8 100644 --- a/tests/test_guard.py +++ b/tests/test_guard.py @@ -369,6 +369,46 @@ class TestIsShowSchemas: 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 + class TestCoverageGaps: """Exercises branches not reached by the scenarios above, so guard.py stays diff --git a/tests/test_server.py b/tests/test_server.py index 38479f6..a904510 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -232,6 +232,38 @@ class TestManagementTools: with pytest.raises(ToolError, match="REST"): make_tools(client).cluster_status() + @pytest.mark.parametrize( + "call", + [ + lambda tools: tools.list_storage_plugins(), + lambda tools: tools.cluster_status(), + lambda tools: tools.list_profiles(), + lambda tools: tools.get_profile("abc"), + lambda tools: tools.cancel_query("abc"), + ], + ids=[ + "list_storage_plugins", + "cluster_status", + "list_profiles", + "get_profile", + "cancel_query", + ], + ) + def test_every_management_tool_is_unavailable_on_a_client_without_it(self, call): + # get_profile/cancel_query run their own isinstance validation on + # query_id before touching the client, so a valid string argument is + # used here to make sure that validation doesn't mask the missing + # REST endpoint being detected first. + client = MagicMock(spec=["query", "schemas", "tables", "columns"]) + with pytest.raises(ToolError, match="REST"): + call(make_tools(client)) + + def test_list_storage_plugins_skips_non_dict_entries_rather_than_crashing(self): + client = MagicMock() + client.storage_plugins.return_value = ["not-a-dict", {"name": "dfs"}] + result = make_tools(client).list_storage_plugins() + assert [p["name"] for p in result] == ["dfs"] + def test_drill_errors_become_tool_errors(self): client = MagicMock() client.profile.side_effect = DrillError("no such query") @@ -318,7 +350,7 @@ class TestShowFiltering: ) assert result["rows"] == [{"SCHEMA_NAME": "dfs"}] - def test_show_filtering_is_case_insensitive(self): + def test_show_filtering_is_case_insensitive_lowercase(self): client = MagicMock() client.query.return_value = QueryResult( ["SCHEMA_NAME"], @@ -327,6 +359,8 @@ class TestShowFiltering: result = make_tools(client, hidden_schemas=["sys"]).run_query("show schemas") assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + def test_show_filtering_is_case_insensitive_mixed_case(self): + client = MagicMock() client.query.return_value = QueryResult( ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}, {"SCHEMA_NAME": "dfs.tmp"}], @@ -343,3 +377,75 @@ class TestShowFiltering: ) result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW TABLES") assert result["rows"] == [{"TABLE_NAME": "sys"}] + + 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.""" + client = MagicMock() + client.query.return_value = QueryResult( + ["TABLE_NAME"], [{"TABLE_NAME": "sys"}] + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query( + "SHOW TABLES LIKE '%s%'" + ) + assert result["rows"] == [{"TABLE_NAME": "sys"}] + + 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.""" + 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 LIKE '%s%'" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + + def test_show_databases_like_rows_are_filtered(self): + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], + [{"SCHEMA_NAME": "INFORMATION_SCHEMA"}, {"SCHEMA_NAME": "dfs"}], + ) + result = make_tools(client, hidden_schemas=["INFORMATION_SCHEMA"]).run_query( + "SHOW DATABASES LIKE '%y%'" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs"}] + + def test_show_schemas_with_trailing_block_comment_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 /* trailing */" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + + def test_show_schemas_with_no_whitespace_before_comment_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_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 + rather than crashing the whole call.""" + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], ["not-a-dict", {"SCHEMA_NAME": "dfs.tmp"}] + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW SCHEMAS") + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}]
