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 c92b58b09a8b16a38c1452822e52a5764911b3c3 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 14:21:12 2026 -0400 feat: server wiring, CLI entry point, and documentation Wires DrillTools' nine bound methods into the installed mcp SDK's MCPServer (mcp.server.fastmcp.FastMCP no longer exists as of mcp 2.0.0, which is what this environment has installed; MCPServer is its replacement and keeps the same add_tool/_tool_manager/run shape), adds the drill-mcp CLI entry point with eager config validation, and documents actual server behavior in README.md and drill.example.yaml. --- README.md | 159 ++++++++++++++++++++++++++++++++++++++++++++++++++- drill.example.yaml | 43 ++++++++++++++ drill_mcp/server.py | 102 ++++++++++++++++++++++++++++++++- tests/test_server.py | 74 ++++++++++++++++++++++++ 4 files changed, 374 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 7c493cb..777f153 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,160 @@ # drill-mcp -MCP server for Apache Drill. +An MCP server for [Apache Drill](https://drill.apache.org/). Lets an MCP client +enumerate schemata, run SQL, and inspect storage plugin and cluster state. + +## Install + +```bash +pip install drill-mcp # REST backend +pip install drill-mcp[jdbc] # adds the JDBC backend (needs a JVM) +pip install drill-mcp[kerberos] # adds SPNEGO for the REST backend +``` + +## Run + +```bash +drill-mcp --url http://localhost:8047 +drill-mcp --config drill.yaml +``` + +Register it with an MCP client: + +```json +{ + "mcpServers": { + "drill": { + "command": "drill-mcp", + "args": ["--config", "/etc/drill-mcp/drill.yaml"] + } + } +} +``` + +`drill-mcp --help` lists every flag; `drill.example.yaml` documents every +config key. CLI flags override the config file, which overrides the +`DRILL_*` environment variables, which override the built-in defaults. +Credentials (`user`/`password`) are read from the config file or from +`DRILL_USER`/`DRILL_PASSWORD` only -- there is no `--user`/`--password` flag, +and no tool accepts a credential as an argument. + +## Tools + +| Tool | Description | +|---|---| +| `run_query` | Run one SQL statement | +| `list_schemas` | List visible schemas | +| `list_tables` | List tables in a schema | +| `describe_table` | Column names and types | +| `list_storage_plugins` | Plugin configs, secrets redacted | +| `cluster_status` | Drillbit membership and status | +| `list_profiles` | Recent and running queries | +| `get_profile` | Full profile for one query id | +| `cancel_query` | Cancel a running query | + +The management tools (`list_storage_plugins`, `cluster_status`, +`list_profiles`, `get_profile`, `cancel_query`) need Drill's REST management +endpoints; they raise a clear `ToolError` when the JDBC backend is in use, +rather than failing silently. + +## Safety model + +- **Writes are denied by default.** `CREATE TABLE AS`, `CREATE VIEW`, and + `DROP` are permitted only into plugins listed in `writable_plugins`. + `INSERT`, `ALTER`, and `USE` are always rejected, no matter what + `writable_plugins` contains. +- **`ALTER SYSTEM` and storage plugin create/update/delete are not + implemented at all.** There is no flag, config key, or code path that + turns them on -- there is nothing for such a tool to call. +- **`EXPLAIN` recurses.** The guard strips a leading `EXPLAIN` or + `EXPLAIN PLAN FOR` and re-checks the remaining statement against the same + allowlist and hidden-schema rules, so `EXPLAIN` cannot be used to peek at + or execute a statement that would otherwise be rejected. Recursion is + bounded so a chain of nested `EXPLAIN EXPLAIN ...` cannot exhaust the + stack. +- **Every statement is checked with a real SQL parser** (sqlglot's Drill + dialect), not a regex, so a write hidden in a comment, a string literal, + or a stacked statement does not slip through. +- **Secrets are always redacted** from storage plugin output. This is not + configurable off. + +## Hidden schemas + +`hidden_schemas` removes schemas (and their children -- hiding `sys` also +hides `sys.mem`, `sys.options`, etc.) from `list_schemas` and +`list_storage_plugins`, and `list_tables`/`describe_table` refuse to operate +on a hidden schema at all. + +`SHOW` commands (`SHOW SCHEMAS`, `SHOW DATABASES`, `SHOW TABLES`, `SHOW +FILES`, ...) are evaluated server-side by Drill, so the guard cannot filter +them by rewriting the query; instead, `run_query` filters the *first column +of every row* of *every* `SHOW` command's result against `hidden_schemas`. + +This is broader than filtering just `SHOW SCHEMAS`/`SHOW DATABASES`, and +that is deliberate. Drill has no single, reliably-recognizable spelling for +"a `SHOW` that lists schemas" -- three narrower approaches (a regex over the +raw SQL text, exact matching against sqlglot's parsed literal, a +comment-stripping regex over that literal) were each defeated by some +spelling (a leading/trailing comment, `SHOW SCHEMAS LIKE '...'`, a nested or +unbalanced comment token) and leaked hidden schema names through. Filtering +every `SHOW` result's first column closes all of those gaps at once, at the +cost of also filtering `SHOW TABLES`/`SHOW FILES` rows: a table or file +whose name happens to match a hidden-schema prefix is hidden too, even +though it is not itself a schema. A false positive here (a real table +briefly missing from a listing) is an acceptable trade for never leaking a +schema name the operator asked to hide. + +**Hiding `INFORMATION_SCHEMA` does not break the metadata tools.** +`list_schemas`, `list_tables`, and `describe_table` query +`INFORMATION_SCHEMA` internally regardless of what is hidden; they simply +omit hidden schemas from what they return to the caller. + +## Column discovery + +`describe_table`'s strategy depends on how the underlying storage plugin +registers its schema: + +- For plugins with a schema registered ahead of time (most JDBC-style and + relational-style plugins), `describe_table` uses Drill's `DESCRIBE`, which + is metadata-only and never reads user data. +- For plugins whose schema is only known at read time (`file`, `mongo`, + `splunk`), `DESCRIBE` cannot answer, so `describe_table` probes with + `SELECT ... LIMIT 1` against the target instead. It returns **only column + names and types** -- the sampled row itself is never included in the + result. +- **HTTP plugins cannot report columns until a query has been run** against + the endpoint; Drill has no schema for an HTTP source ahead of a real + request. `describe_table` says this explicitly in its error rather than + returning an empty column list. + +`list_tables` uses `SHOW FILES` instead of `INFORMATION_SCHEMA.TABLES` for +file plugins, since file-based storage plugins do not register their +contents in `INFORMATION_SCHEMA`. + +## Backends + +- **REST** (default): talks to a Drillbit's HTTP endpoints over `httpx`. + Works out of the box, no JVM required. +- **JDBC** (`backend: jdbc`, the `[jdbc]` extra): uses `jaydebeapi`/`JPype1` + against Drill's JDBC driver. Exists mainly for Kerberos environments where + the REST endpoint's SPNEGO support is impractical. Needs a JVM and + `jdbc_driver_path` pointing at `drill-jdbc-all.jar`. The management tools + are unavailable on this backend, since Drill's management API is + REST-only. + +Wire-level behavior (identifier quoting, `INFORMATION_SCHEMA` queries, +per-plugin-type column discovery) follows +[`sqlalchemy-drill`](https://github.com/JohnOmernik/sqlalchemy-drill), the +most complete and maintained reference for talking to Drill from Python. +`PyDrill` was evaluated and not adopted: it has no form-based login, no +Kerberos support, and its last release was in 2018. + +See `drill.example.yaml` for the full configuration reference. + +## Development + +```bash +pip install -e ".[dev]" +pytest # unit tests, no cluster or JVM needed +pytest -m integration # requires a live Drill at DRILL_URL +``` diff --git a/drill.example.yaml b/drill.example.yaml new file mode 100644 index 0000000..244b7e1 --- /dev/null +++ b/drill.example.yaml @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +# either express or implied. See the License for the specific +# language governing permissions and limitations under the +# License. +# + +# Apache Drill MCP server configuration. +# Every value below is a default; delete what you do not need to change. + +url: http://localhost:8047 +backend: rest # rest | jdbc +auth: none # none | basic | kerberos + +# Credentials may also come from DRILL_USER / DRILL_PASSWORD. +# user: alice +# password: s3cret + +max_rows: 1000 +timeout_seconds: 60 + +# Plugins permitted to accept CTAS / CREATE VIEW / DROP. Empty means no writes. +writable_plugins: [] +# writable_plugins: [dfs.tmp] + +# Schemas hidden from listings and rejected in queries. +hidden_schemas: [] +# hidden_schemas: [sys, INFORMATION_SCHEMA] + +# Required when backend is jdbc. +# jdbc_driver_path: /opt/drill/jars/jdbc-driver/drill-jdbc-all.jar diff --git a/drill_mcp/server.py b/drill_mcp/server.py index d9262c9..6f3700d 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -21,15 +21,23 @@ Tool bodies live on `DrillTools` as plain methods so they can be unit-tested without standing up an MCP session; `build_server` (Task 9/10) registers the -bound methods with FastMCP. +bound methods with the MCP server (`mcp.server.mcpserver.MCPServer` -- this +package's `mcp` dependency renamed `FastMCP` to `MCPServer` as of mcp 2.0.0; +the internal shape used here, `add_tool`/`_tool_manager.list_tools`/`run`, +is unchanged). """ from __future__ import annotations +import argparse +import logging +import sys from typing import Any -from .client_rest import DrillError -from .config import Config +from mcp.server.mcpserver import MCPServer + +from .client_rest import DrillError, RestClient +from .config import Config, ConfigError, load_config from .guard import Policy, PolicyError, check, is_show_command, matches_prefix @@ -207,3 +215,91 @@ class DrillTools: return self._require_management("cancel_query")(query_id) except DrillError as exc: raise ToolError(str(exc)) from exc + + +def build_client(config: Config) -> Any: + """Construct the wire client the configured backend calls for.""" + if config.backend == "jdbc": + from .client_jdbc import JdbcClient + + return JdbcClient(config) + return RestClient(config) + + +def build_server(config: Config) -> MCPServer: + """Build an MCP server with every read/metadata tool registered. + + No write- or mutation-capable tool is ever registered here: storage + plugin create/update/delete and `ALTER SYSTEM` have no implementation + anywhere in this package, so there is nothing such a tool could call. + """ + client = build_client(config) + tools = DrillTools(config, client) + server = MCPServer("drill") + for method in ( + tools.run_query, + tools.list_schemas, + tools.list_tables, + tools.describe_table, + tools.list_storage_plugins, + tools.cluster_status, + tools.list_profiles, + tools.get_profile, + tools.cancel_query, + ): + server.add_tool(method, name=method.__name__, description=method.__doc__) + return server + + +def _parse_args(argv: list[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser(prog="drill-mcp", description="MCP server for Apache Drill") + parser.add_argument("--config", help="path to a YAML config file") + parser.add_argument("--url", help="Drill HTTP endpoint, e.g. http://localhost:8047";) + parser.add_argument("--backend", choices=["rest", "jdbc"]) + parser.add_argument("--auth", choices=["none", "basic", "kerberos"]) + parser.add_argument("--max-rows", type=int, dest="max_rows") + parser.add_argument( + "--writable-plugin", + action="append", + dest="writable_plugins", + metavar="PLUGIN", + help="permit data writes into this plugin; repeatable, empty by default", + ) + parser.add_argument( + "--hidden-schema", + action="append", + dest="hidden_schemas", + metavar="SCHEMA", + help="hide this schema from listings and queries; repeatable", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + # sqlglot logs a WARNING-level "Falling back to parsing as a 'Command'" + # message for every SHOW/EXPLAIN/ALTER statement the guard parses (twice + # for SHOW, since the guard parses it twice). Python's logging module + # writes unconfigured loggers to stderr, never stdout, so this cannot + # corrupt the JSON-RPC session on stdio -- it is only quieted here so + # operators are not spammed. Configured here, not at import time: a + # library that reconfigures logging as a side effect of being imported + # is bad manners. + logging.getLogger("sqlglot").setLevel(logging.ERROR) + + args = _parse_args(argv) + overrides = { + key: value + for key, value in vars(args).items() + if key != "config" and value is not None + } + try: + config = load_config(args.config, overrides=overrides) + except ConfigError as exc: + print(f"drill-mcp: configuration error: {exc}", file=sys.stderr) + return 1 + build_server(config).run() + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/tests/test_server.py b/tests/test_server.py index 8139ca3..55fcf16 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -535,3 +535,77 @@ class TestShowFiltering: 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"}] + + +from drill_mcp.client_jdbc import JdbcClient +from drill_mcp.client_rest import RestClient +from drill_mcp.server import build_client, build_server + + +class TestWiring: + def test_rest_backend_builds_a_rest_client(self): + assert isinstance(build_client(load_config()), RestClient) + + def test_jdbc_backend_builds_a_jdbc_client(self): + cfg = load_config(overrides={"backend": "jdbc", "jdbc_driver_path": "/x.jar"}) + assert isinstance(build_client(cfg), JdbcClient) + + def test_all_tools_are_registered(self): + server = build_server(load_config()) + names = {tool.name for tool in server._tool_manager.list_tools()} + assert names == { + "run_query", + "list_schemas", + "list_tables", + "describe_table", + "list_storage_plugins", + "cluster_status", + "list_profiles", + "get_profile", + "cancel_query", + } + + def test_every_tool_has_a_description(self): + server = build_server(load_config()) + assert all(tool.description for tool in server._tool_manager.list_tools()) + + def test_no_write_or_mutation_tools_are_registered(self): + server = build_server(load_config()) + names = {tool.name for tool in server._tool_manager.list_tools()} + forbidden = {"create_storage_plugin", "update_storage_plugin", + "delete_storage_plugin", "set_option", "alter_system"} + assert not (names & forbidden) + + def test_no_registered_tool_accepts_a_credential_argument(self): + """Credentials come from config or environment only, never a tool argument.""" + server = build_server(load_config()) + credential_words = {"user", "password", "username", "passwd", "secret", "token", "credential"} + for tool in server._tool_manager.list_tools(): + params = set(tool.parameters.get("properties", {})) + assert not (params & credential_words), f"{tool.name} accepts {params & credential_words}" + + +class TestMain: + def test_config_error_exits_nonzero_with_a_message(self, capsys): + from drill_mcp.server import main + + assert main(["--config", "/nonexistent.yaml"]) == 1 + assert "not found" in capsys.readouterr().err + + def test_cli_flags_reach_the_config(self, monkeypatch): + from drill_mcp import server as server_module + + # Note: `captured.setdefault("cfg", cfg) or MagicMock()` (as drafted + # in the task brief) returns the truthy `cfg` itself rather than the + # MagicMock, so `build_server(cfg).run()` would blow up calling + # `.run()` on a `Config`. Using a real fake avoids that trap. + captured = {} + + def fake_build_server(cfg): + captured["cfg"] = cfg + return MagicMock() + + monkeypatch.setattr(server_module, "build_server", fake_build_server) + server_module.main(["--url", "http://cli:8047";, "--max-rows", "7"]) + assert captured["cfg"].url == "http://cli:8047"; + assert captured["cfg"].max_rows == 7
