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 cedd93d7f80e158401de55cc621c82815c6d903e Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 09:48:36 2026 -0400 feat: management tools and parser-based SHOW SCHEMAS filtering SHOW SCHEMAS/DATABASES rows are filtered on return using a parsed-statement check (guard.is_show_schemas) rather than a regex over the raw SQL, since a leading comment defeats a regex anchor but not the tokenizer. --- drill_mcp/guard.py | 31 ++++++++++ drill_mcp/server.py | 70 ++++++++++++++++++++++- tests/test_guard.py | 44 +++++++++++++- tests/test_server.py | 159 +++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 301 insertions(+), 3 deletions(-) diff --git a/drill_mcp/guard.py b/drill_mcp/guard.py index ce7b8a3..f4a1eb1 100644 --- a/drill_mcp/guard.py +++ b/drill_mcp/guard.py @@ -90,6 +90,37 @@ 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`. + + 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'))`). + + 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. + """ + try: + statements = [s for s in sqlglot.parse(sql, read=DIALECT) if s is not None] + except Exception: + return False + 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 "") + return str(target or "").strip().upper() in {"SCHEMAS", "DATABASES"} + + def check(sql: str, policy: Policy) -> None: """Return None if `sql` is permitted under `policy`; raise PolicyError otherwise.""" _check(sql, policy, depth=0) diff --git a/drill_mcp/server.py b/drill_mcp/server.py index 5576ebf..b815e12 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -30,13 +30,19 @@ from typing import Any from .client_rest import DrillError from .config import Config -from .guard import Policy, PolicyError, check, matches_prefix +from .guard import Policy, PolicyError, check, is_show_schemas, matches_prefix class ToolError(Exception): """The single error type surfaced to MCP clients. Never carries a traceback.""" +def _first_value(row: dict[str, Any]) -> str | None: + for value in row.values(): + return str(value) if value is not None else None + return None + + class DrillTools: def __init__(self, config: Config, client: Any) -> None: self._config = config @@ -86,9 +92,16 @@ 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. + rows = result.rows + if self._policy.hidden_schemas and is_show_schemas(sql): + rows = [row for row in rows if self._visible(_first_value(row))] + payload: dict[str, Any] = { "columns": result.columns, - "rows": result.rows, + "rows": rows, "query_id": result.query_id, "truncated": result.truncated, } @@ -122,3 +135,56 @@ class DrillTools: return self._client.columns(schema, table) except DrillError as exc: raise ToolError(str(exc)) from exc + + # -- management tools --------------------------------------------------- + + def _require_management(self, name: str) -> Any: + method = getattr(self._client, name, None) + if method is None: + raise ToolError( + f"'{name}' needs a REST connection to Drill; the JDBC backend " + "does not expose management endpoints" + ) + return method + + def list_storage_plugins(self) -> list[dict[str, Any]]: + """List storage plugin configurations, with all secrets redacted.""" + try: + 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"))] + + def cluster_status(self) -> dict[str, Any]: + """Report Drillbit membership and overall cluster status.""" + try: + return self._require_management("cluster_status")() + except DrillError as exc: + raise ToolError(str(exc)) from exc + + def list_profiles(self, limit: int = 20) -> list[dict[str, Any]]: + """List recent and running query profiles, newest first.""" + if not isinstance(limit, int) or isinstance(limit, bool): + raise ToolError("limit must be an integer") + try: + return self._require_management("profiles")(limit=limit) + except DrillError as exc: + raise ToolError(str(exc)) from exc + + def get_profile(self, query_id: str) -> dict[str, Any]: + """Fetch the full profile for one query id.""" + if not isinstance(query_id, str): + raise ToolError("query_id must be a string") + try: + return self._require_management("profile")(query_id) + except DrillError as exc: + raise ToolError(str(exc)) from exc + + def cancel_query(self, query_id: str) -> str: + """Cancel a running query by its query id.""" + if not isinstance(query_id, str): + raise ToolError("query_id must be a string") + try: + return self._require_management("cancel_query")(query_id) + except DrillError as exc: + raise ToolError(str(exc)) from exc diff --git a/tests/test_guard.py b/tests/test_guard.py index 73fa613..0b5f32a 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, matches_prefix +from drill_mcp.guard import Policy, PolicyError, check, is_show_schemas, matches_prefix class TestSqlglotAssumptions: @@ -328,6 +328,48 @@ 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. + """ + + def test_show_schemas_matches(self): + assert is_show_schemas("SHOW SCHEMAS") is True + + def test_show_databases_matches(self): + assert is_show_schemas("SHOW DATABASES") is True + + def test_lowercase_matches(self): + assert is_show_schemas("show schemas") is True + + def test_mixed_case_matches(self): + assert is_show_schemas("ShOw DaTaBaSeS") is True + + def test_leading_block_comment_still_matches(self): + assert is_show_schemas("/* x */ SHOW SCHEMAS") is True + + def test_leading_line_comment_still_matches(self): + assert is_show_schemas("-- comment\nSHOW DATABASES") is True + + def test_show_tables_does_not_match(self): + assert is_show_schemas("SHOW TABLES") is False + + def test_show_files_does_not_match(self): + assert is_show_schemas("SHOW FILES IN dfs.tmp") is False + + def test_ordinary_select_does_not_match(self): + assert is_show_schemas("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 + + def test_unparseable_sql_returns_false_rather_than_raising(self): + assert is_show_schemas("((((((") is False + + def test_empty_sql_returns_false(self): + assert is_show_schemas("") is False + + class TestCoverageGaps: """Exercises branches not reached by the scenarios above, so guard.py stays at 100% line coverage without weakening any other test. diff --git a/tests/test_server.py b/tests/test_server.py index ca2dd89..38479f6 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -184,3 +184,162 @@ class TestDescribeTable: client.columns.side_effect = DrillError("invalid identifier") with pytest.raises(ToolError, match="invalid identifier"): make_tools(client).describe_table("dfs.tmp", "nope") + + +class TestManagementTools: + def test_list_storage_plugins_passes_through_redacted_output(self): + client = MagicMock() + client.storage_plugins.return_value = [ + {"name": "s3", "config": {"secret": "***REDACTED***"}} + ] + assert make_tools(client).list_storage_plugins()[0]["name"] == "s3" + + def test_list_storage_plugins_hides_plugins_backing_hidden_schemas(self): + client = MagicMock() + client.storage_plugins.return_value = [{"name": "sys"}, {"name": "dfs"}] + result = make_tools(client, hidden_schemas=["sys"]).list_storage_plugins() + assert [p["name"] for p in result] == ["dfs"] + + def test_cluster_status(self): + client = MagicMock() + client.cluster_status.return_value = {"status": "Running!"} + assert make_tools(client).cluster_status()["status"] == "Running!" + + def test_list_profiles_uses_the_default_limit(self): + client = MagicMock() + client.profiles.return_value = [] + make_tools(client).list_profiles() + assert client.profiles.call_args.kwargs["limit"] == 20 + + def test_list_profiles_honours_an_explicit_limit(self): + client = MagicMock() + client.profiles.return_value = [] + make_tools(client).list_profiles(limit=5) + assert client.profiles.call_args.kwargs["limit"] == 5 + + def test_get_profile(self): + client = MagicMock() + client.profile.return_value = {"queryId": "abc"} + assert make_tools(client).get_profile("abc")["queryId"] == "abc" + + def test_cancel_query(self): + client = MagicMock() + client.cancel_query.return_value = "Cancelled" + assert make_tools(client).cancel_query("abc") == "Cancelled" + + def test_management_tools_are_unavailable_on_a_client_without_them(self): + client = MagicMock(spec=["query", "schemas", "tables", "columns"]) + with pytest.raises(ToolError, match="REST"): + make_tools(client).cluster_status() + + def test_drill_errors_become_tool_errors(self): + client = MagicMock() + client.profile.side_effect = DrillError("no such query") + with pytest.raises(ToolError, match="no such query"): + make_tools(client).get_profile("abc") + + def test_get_profile_rejects_non_string_query_id(self): + client = MagicMock() + with pytest.raises(ToolError): + make_tools(client).get_profile(123) + + def test_cancel_query_rejects_non_string_query_id(self): + client = MagicMock() + with pytest.raises(ToolError): + make_tools(client).cancel_query(123) + + def test_list_profiles_rejects_non_integer_limit(self): + client = MagicMock() + with pytest.raises(ToolError): + make_tools(client).list_profiles(limit="20") + + +class TestShowFiltering: + """SHOW is evaluated server-side by Drill, so rows are filtered on return.""" + + def test_show_schemas_rows_are_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_databases_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" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs"}] + + def test_ordinary_select_rows_are_not_filtered(self): + client = MagicMock() + client.query.return_value = QueryResult( + ["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}] + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query( + "SELECT SCHEMA_NAME FROM dfs.tmp.notes" + ) + assert result["rows"] == [{"SCHEMA_NAME": "sys"}] + + def test_show_filtering_is_a_no_op_without_hidden_schemas(self): + client = MagicMock() + client.query.return_value = QueryResult(["SCHEMA_NAME"], [{"SCHEMA_NAME": "sys"}]) + result = make_tools(client).run_query("SHOW SCHEMAS") + assert result["rows"] == [{"SCHEMA_NAME": "sys"}] + + def test_show_schemas_with_leading_block_comment_is_still_filtered(self): + """Regression test for the raw-regex bypass: a leading comment defeats + a `^\\s*SHOW` anchor but not the parser, since the tokenizer strips + comments before the guard or this filter ever sees the text.""" + 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( + "/* x */ SHOW SCHEMAS" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs.tmp"}] + + def test_show_databases_with_leading_line_comment_is_still_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( + "-- comment\nSHOW DATABASES" + ) + assert result["rows"] == [{"SCHEMA_NAME": "dfs"}] + + def test_show_filtering_is_case_insensitive(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"}] + + 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_tables_rows_are_not_filtered(self): + """SHOW TABLES rows are table names, not schema names; filtering them + against hidden_schemas would be wrong.""" + client = MagicMock() + client.query.return_value = QueryResult( + ["TABLE_NAME"], [{"TABLE_NAME": "sys"}] + ) + result = make_tools(client, hidden_schemas=["sys"]).run_query("SHOW TABLES") + assert result["rows"] == [{"TABLE_NAME": "sys"}]
