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 81c02ae3f7c031a928ee166f030b2a50e2fdf534 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 16:24:19 2026 -0400 Various fixes --- README.md | 4 ++-- docs/tools.md | 18 +++++++++++++---- drill_mcp/client_rest.py | 48 +++++++++++++++++++++++++++++++++----------- drill_mcp/config.py | 11 +++++++++- drill_mcp/redact.py | 28 ++++++++++++++++++++++++++ drill_mcp/server.py | 37 ++++++++++++++++++++++++++++++---- pyproject.toml | 2 -- tests/test_client_jdbc.py | 2 +- tests/test_client_rest.py | 41 ++++++++++++++++++++++++++++++++++++- tests/test_config.py | 11 ++++++++++ tests/test_redact.py | 21 +++++++++++++++++++ tests/test_server.py | 51 ++++++++++++++++++++++++++++++++++++++++------- 12 files changed, 240 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index 24f4615..baf20cf 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# drill-mcp +# Drill-MCP Server -An [MCP](https://modelcontextprotocol.io/) server for [Apache Drill](https://drill.apache.org/). +The official [MCP](https://modelcontextprotocol.io/) server for [Apache Drill](https://drill.apache.org/). It lets an MCP client run read-only (and narrowly, explicitly allow-listed write) SQL against a Drill cluster, and inspect schemas, storage plugins, and cluster/query state. It does not implement Drill administration diff --git a/docs/tools.md b/docs/tools.md index 7897218..1c8ff5e 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -397,7 +397,12 @@ client surfaces defaults, since a caller passing no `limit` at all still gets at most 20 profiles back, not an unbounded list. **Returns** a list of profile summary dicts, shape defined by Drill's -`/profiles.json`: +`/profiles.json`, passed through the same secret redaction as +`list_storage_plugins` (profiles are cluster-wide and can carry other +users' connection strings). Any entry whose query text names a +`hidden_schemas` entry is dropped entirely, the same protection +`list_schemas`/`list_tables` apply — otherwise a hidden schema's name would +leak out as data in another user's query text: ```json [ @@ -431,9 +436,12 @@ Fetches the full profile for one query id. |---|---|---|---| | `query_id` | string | no | Drill's query UUID, as returned in `run_query`'s `query_id` field (REST backend only — always `null` on JDBC). Validated against `[A-Za-z0-9-]+`. | -**Returns** Drill's full profile JSON for that query, unmodified — the -complete `/profiles/<query_id>.json` payload (fragments, operator metrics, -timing, etc.), which can be large. +**Returns** Drill's full profile JSON for that query — the complete +`/profiles/<query_id>.json` payload (fragments, operator metrics, timing, +etc.), which can be large — passed through the same secret redaction as +`list_storage_plugins`. A full profile embeds Drill's serialized physical +plan, which for JDBC and HTTP storage plugins can carry plugin +configuration, so this is not returned unmodified. **Errors** @@ -442,6 +450,8 @@ timing, etc.), which can be large. - `query_id must be a string` — wrong type. - `invalid query id: '<x>'` — `query_id` contains characters outside `[A-Za-z0-9-]`. Rejected before any request is made. +- `profile '<id>' references a hidden schema` — the profile's query text + names a `hidden_schemas` entry. - Drill's own error text otherwise (e.g. no profile with that id). --- diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index 2165ba0..110fa43 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -31,6 +31,7 @@ import re from collections.abc import Callable from dataclasses import dataclass, field from typing import Any +from urllib.parse import urlparse, urlunparse import httpx @@ -180,6 +181,22 @@ def _check_query_id(query_id: str) -> str: return query_id +def _safe_url(url: str) -> str: + """Return `url` with any embedded userinfo (e.g. a password) stripped. + + `config.url` is free-form and unvalidated, so nothing stops a value like + `http://alice:s3cret@drill:8047`. Every message that echoes the URL back + to the model must use this instead of the raw config value -- the + `client_jdbc.py` backend already applies this exact defense to its + connection string; the REST backend must not diverge. + """ + parsed = urlparse(url) + if not parsed.hostname: + return url + netloc = parsed.hostname if parsed.port is None else f"{parsed.hostname}:{parsed.port}" + return urlunparse(parsed._replace(netloc=netloc)) + + def _json(response: httpx.Response, url: str) -> Any: """Decode a response body as JSON, converting a decode failure to `DrillError`. @@ -512,23 +529,23 @@ class RestClient: # "/j_security_check", so that check false-positives on success.) if response.status_code >= 400: raise DrillError( - f"authentication endpoint at {self._config.url} returned " + f"authentication endpoint at {_safe_url(self._config.url)} returned " f"HTTP {response.status_code}" ) if _contains_invalid_credentials_marker(response.text): raise DrillError( - f"authentication failed for user {self._config.user!r} at {self._config.url}" + f"authentication failed for user {self._config.user!r} at {_safe_url(self._config.url)}" ) self._authenticated = True def _transport_error(self, exc: httpx.HTTPError) -> DrillError: if isinstance(exc, httpx.TimeoutException): return DrillError( - f"request to {self._config.url} timed out after " + f"request to {_safe_url(self._config.url)} timed out after " f"{self._config.timeout_seconds}s" ) return DrillError( - f"could not reach Drill at {self._config.url} " + f"could not reach Drill at {_safe_url(self._config.url)} " f"(auth mode: {self._config.auth}): {type(exc).__name__}" ) @@ -546,7 +563,7 @@ class RestClient: if response.status_code == 401: raise DrillError( - f"authentication rejected by Drill at {self._config.url} " + f"authentication rejected by Drill at {_safe_url(self._config.url)} " f"for user {self._config.user!r}" ) if response.status_code >= 400: @@ -561,13 +578,20 @@ class RestClient: "/query.json", json={"queryType": "SQL", "query": sql, "autoLimit": max_rows}, ) - payload = _json(response, self._config.url) + payload = _json(response, _safe_url(self._config.url)) rows = payload.get("rows") or [] + truncated = max_rows > 0 and len(rows) >= max_rows + # Defense in depth: `autoLimit` asks Drill to cap rows server-side, + # but the cap must not depend entirely on Drill honoring that field. + # Slice client-side too, exactly like `JdbcClient.query`'s + # `fetchmany(max_rows)` -- the two backends must agree on this. + if max_rows > 0: + rows = rows[:max_rows] return QueryResult( columns=payload.get("columns") or [], rows=rows, query_id=payload.get("queryId"), - truncated=max_rows > 0 and len(rows) >= max_rows, + truncated=truncated, metadata=payload.get("metadata") or [], ) @@ -594,11 +618,11 @@ class RestClient: def storage_plugins(self) -> list[dict[str, Any]]: response = self._request("GET", "/storage.json") - return redact(_json(response, self._config.url)) + return redact(_json(response, _safe_url(self._config.url))) def cluster_status(self) -> dict[str, Any]: - cluster = _json(self._request("GET", "/cluster.json"), self._config.url) - status = _json(self._request("GET", "/status.json"), self._config.url) + cluster = _json(self._request("GET", "/cluster.json"), _safe_url(self._config.url)) + status = _json(self._request("GET", "/status.json"), _safe_url(self._config.url)) merged = dict(cluster) if isinstance(cluster, dict) else {"cluster": cluster} if isinstance(status, dict): merged.update(status) @@ -609,7 +633,7 @@ class RestClient: def profiles(self, limit: int) -> list[dict[str, Any]]: limit = max(limit, 0) response = self._request("GET", "/profiles.json") - payload = _json(response, self._config.url) + payload = _json(response, _safe_url(self._config.url)) if not isinstance(payload, dict): payload = {} running = payload.get("runningQueries") or [] @@ -619,7 +643,7 @@ class RestClient: def profile(self, query_id: str) -> dict[str, Any]: _check_query_id(query_id) response = self._request("GET", f"/profiles/{query_id}.json") - return _json(response, self._config.url) + return _json(response, _safe_url(self._config.url)) def cancel_query(self, query_id: str) -> str: _check_query_id(query_id) diff --git a/drill_mcp/config.py b/drill_mcp/config.py index f413765..721310a 100644 --- a/drill_mcp/config.py +++ b/drill_mcp/config.py @@ -112,4 +112,13 @@ def load_config( try: return Config(**values) except ValidationError as exc: - raise ConfigError(str(exc)) from exc + # `str(exc)` embeds pydantic's `input_value=...` for every error, + # which echoes the offending config value verbatim -- including a + # password typed unquoted in YAML (e.g. `password: 12345`, which + # pydantic reports as `input_value=12345`). That must never reach + # stderr or a log. Rebuild the message from `loc`/`msg` only. + details = "; ".join( + f"{'.'.join(str(p) for p in err['loc'])}: {err['msg']}" if err["loc"] else err["msg"] + for err in exc.errors() + ) + raise ConfigError(details) from exc diff --git a/drill_mcp/redact.py b/drill_mcp/redact.py index f6ac0a4..150ea28 100644 --- a/drill_mcp/redact.py +++ b/drill_mcp/redact.py @@ -39,6 +39,32 @@ _SENSITIVE = re.compile( re.IGNORECASE, ) +# The key-based check above only catches secrets that live at a sensitive +# *key*. A storage-plugin config routinely carries secrets embedded inside an +# ordinary-looking *value* instead -- a JDBC/S3-style URL with userinfo +# (`s3a://AKIA:secret@bucket`), or a connection string with a `password=`/ +# `secret=`/`token=` query parameter. Both shapes are real Drill storage +# plugin configs, so string values are scrubbed too, not just keys. +# +# Matches "scheme://user:pass@" and keeps everything else (scheme, host, +# path) intact -- only the credential pair between "//" and "@" is replaced. +_URL_USERINFO = re.compile(r"(?P<scheme>[A-Za-z][A-Za-z0-9+.-]*://)[^/@\s]+:[^/@\s]*@") + +# Matches a `?key=value` or `&key=value` query parameter whose key looks like +# a secret, and replaces only the value -- the `?`/`&` and key name are kept +# so the rest of the string still parses as the same shape of URL. +_QUERY_SECRET = re.compile( + r"(?P<prefix>[?&](?:password|secret|api_?key|token)=)[^&]*", + re.IGNORECASE, +) + + +def _scrub_value(value: str) -> str: + """Strip credentials embedded inside a string value, not just its key.""" + value = _URL_USERINFO.sub(lambda m: f"{m.group('scheme')}{REDACTED}@", value) + value = _QUERY_SECRET.sub(lambda m: f"{m.group('prefix')}{REDACTED}", value) + return value + def redact(value: Any) -> Any: """Return a copy of `value` with sensitive-looking values replaced.""" @@ -51,4 +77,6 @@ def redact(value: Any) -> Any: return [redact(item) for item in value] if isinstance(value, tuple): return tuple(redact(item) for item in value) + if isinstance(value, str): + return _scrub_value(value) return value diff --git a/drill_mcp/server.py b/drill_mcp/server.py index 2d9cae7..c4c4bf3 100644 --- a/drill_mcp/server.py +++ b/drill_mcp/server.py @@ -39,6 +39,7 @@ 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 +from .redact import redact if TYPE_CHECKING: # Imported only for the type checker: build_client's lazy, in-function @@ -77,6 +78,25 @@ class DrillTools: if matches_prefix(schema, self._policy.hidden_schemas): raise ToolError(f"schema '{schema}' is hidden by configuration") + def _profile_mentions_hidden_schema(self, profile: dict[str, Any]) -> bool: + """True if a profile's query text names a hidden schema. + + Profiles are cluster-wide: `list_profiles`/`get_profile` surface + *other users'* query text, so a hidden schema name can leak out here + as data even though it was never queryable directly -- the one + enumeration path the guard and the metadata-tool filtering were + built to close. A case-insensitive substring match against the + query text is deliberately coarse (over-filtering a profile whose + SQL merely mentions a hidden schema's name in a string literal is an + acceptable false positive; missing one is a leak). + """ + if not self._policy.hidden_schemas: + return False + query_text = str(profile.get("query") or "").lower() + if not query_text: + return False + return any(schema.lower() in query_text for schema in self._policy.hidden_schemas) + def _visible(self, schema: str | None) -> bool: # Fail closed, not open: an item this function cannot identify (no # name at all) is filtered out rather than shown by default. Drill @@ -200,18 +220,27 @@ class DrillTools: if not isinstance(limit, int) or isinstance(limit, bool): raise ToolError("limit must be an integer") try: - return self._require_management("profiles")(limit=limit) + profiles = self._require_management("profiles")(limit=limit) except DrillError as exc: raise ToolError(str(exc)) from exc + profiles = redact(profiles) + return [ + p + for p in profiles + if isinstance(p, dict) and not self._profile_mentions_hidden_schema(p) + ] 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) + profile = self._require_management("profile")(query_id) except DrillError as exc: raise ToolError(str(exc)) from exc + if isinstance(profile, dict) and self._profile_mentions_hidden_schema(profile): + raise ToolError(f"profile {query_id!r} references a hidden schema") + return redact(profile) def cancel_query(self, query_id: str) -> str: """Cancel a running query by its query id.""" @@ -301,10 +330,10 @@ def main(argv: list[str] | None = None) -> int: } try: config = load_config(args.config, overrides=overrides) - except ConfigError as exc: + build_server(config).run() + except (ConfigError, DrillError) as exc: print(f"drill-mcp: configuration error: {exc}", file=sys.stderr) return 1 - build_server(config).run() return 0 diff --git a/pyproject.toml b/pyproject.toml index 4e0e66c..96a5ff5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,5 +53,3 @@ drill-mcp = "drill_mcp.server:main" [tool.pytest.ini_options] testpaths = ["tests"] -markers = ["integration: requires a live Drill cluster (deselected by default)"] -addopts = "-m 'not integration'" diff --git a/tests/test_client_jdbc.py b/tests/test_client_jdbc.py index 02ae59f..29d5958 100644 --- a/tests/test_client_jdbc.py +++ b/tests/test_client_jdbc.py @@ -56,7 +56,7 @@ def fake_jaydebeapi(monkeypatch): def make_client(**overrides): overrides.setdefault("backend", "jdbc") overrides.setdefault("jdbc_driver_path", "/opt/drill-jdbc-all.jar") - return JdbcClient(load_config(overrides=overrides)) + return JdbcClient(load_config(overrides=overrides, env={})) def test_clear_error_when_extra_is_not_installed(monkeypatch): diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 983be61..0a0e939 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -39,7 +39,7 @@ BASE = "http://drill:8047"; def make_client(**overrides): overrides.setdefault("url", BASE) - return RestClient(load_config(overrides=overrides)) + return RestClient(load_config(overrides=overrides, env={})) class TestQuoting: @@ -134,6 +134,24 @@ class TestQuery: ) assert make_client().query("SELECT 1", max_rows=2).truncated is True + @respx.mock + def test_slices_rows_to_max_rows_even_if_drill_ignores_autolimit(self): + # `autoLimit` asks Drill to cap rows server-side, but the cap must + # not depend entirely on Drill honoring that field. Simulate Drill + # returning more rows than requested (e.g. an older Drill version, or + # autoLimit simply not being respected) and confirm the client still + # enforces the cap itself -- exactly like JdbcClient.query's + # fetchmany(max_rows). + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response( + 200, + json={"columns": ["a"], "rows": [{"a": str(i)} for i in range(10)]}, + ) + ) + result = make_client().query("SELECT 1", max_rows=2) + assert len(result.rows) == 2 + assert result.truncated is True + @respx.mock def test_not_truncated_when_max_rows_is_zero(self): # 0 >= 0 would be a false "truncated" without the max_rows > 0 guard. @@ -172,6 +190,27 @@ class TestQuery: assert BASE in str(exc.value) assert "s3cret" not in str(exc.value) + @respx.mock + def test_connection_failure_message_drops_a_password_embedded_in_the_url(self): + # Mirrors client_jdbc.py's + # test_jdbc_url_drops_userinfo_from_a_url_that_embeds_credentials -- + # the REST backend must apply the same defense. config.url is + # free-form and unvalidated, so nothing stops + # DRILL_URL=http://alice:s3cret@drill:8047; every message that + # echoes config.url back to the model must not leak the password + # embedded there. + url = "http://alice:s3cret@drill:8047"; + respx.post(f"{url}/j_security_check").mock(side_effect=httpx.ConnectError("refused")) + respx.post(f"{url}/query.json").mock(side_effect=httpx.ConnectError("refused")) + with pytest.raises(DrillError) as exc: + make_client(url=url, auth="basic", user="alice", password="s3cret").query( + "SELECT 1", max_rows=1 + ) + message = str(exc.value) + assert "s3cret" not in message + assert "alice" not in message + assert "drill:8047" in message + @respx.mock def test_timeout_is_reported_clearly(self): respx.post(f"{BASE}/query.json").mock(side_effect=httpx.ReadTimeout("slow")) diff --git a/tests/test_config.py b/tests/test_config.py index e7c1931..a40e59d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -107,3 +107,14 @@ def test_config_is_immutable(): cfg = load_config(env={}) with pytest.raises(ValidationError, match="frozen"): cfg.url = "http://elsewhere:8047"; + + +def test_a_non_string_password_does_not_appear_in_the_error_message(): + # pydantic's default ValidationError text embeds `input_value=...` for + # every error -- e.g. an unquoted `password: 12345` in YAML produces + # "input_value=12345" verbatim. ConfigError must not echo that: `main()` + # prints it to stderr, and any log capturing stderr would then carry the + # (would-be) password in plaintext. + with pytest.raises(ConfigError) as exc: + load_config(env={}, overrides={"password": 12345}) + assert "12345" not in str(exc.value) diff --git a/tests/test_redact.py b/tests/test_redact.py index 7c0c64b..f195a20 100644 --- a/tests/test_redact.py +++ b/tests/test_redact.py @@ -48,6 +48,27 @@ def test_leaves_innocuous_keys_alone(): } +def test_redacts_userinfo_embedded_in_a_url_shaped_value(): + # A secret can live inside an ordinary-looking value, not just behind a + # sensitive key -- e.g. a `connection` string embedding S3 credentials as + # userinfo. This used to survive untouched (the bug this test replaces + # `test_leaves_innocuous_keys_alone`'s old assertion for): the key-based + # check alone let `redact({"connection": "s3a://AKIA:secret@bucket"})` + # pass the secret straight through. + source = {"connection": "s3a://AKIA:supersecret@bucket"} + result = redact(source) + assert "supersecret" not in result["connection"] + assert "AKIA" not in result["connection"] + assert result["connection"] == "s3a://***REDACTED***@bucket" + + +def test_redacts_password_query_parameter_embedded_in_a_url_shaped_value(): + source = {"url": "jdbc:mysql://host/db?user=root&password=hunter2"} + result = redact(source) + assert "hunter2" not in result["url"] + assert result["url"] == "jdbc:mysql://host/db?user=root&password=***REDACTED***" + + def test_recurses_into_nested_dicts(): source = {"config": {"aws": {"awsSecretAccessKey": "s"}}} assert redact(source)["config"]["aws"]["awsSecretAccessKey"] == REDACTED diff --git a/tests/test_server.py b/tests/test_server.py index 3678132..c323598 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -29,7 +29,7 @@ from drill_mcp.server import DrillTools, ToolError, build_client, build_server def make_tools(client=None, **overrides): client = client or MagicMock() - return DrillTools(load_config(overrides=overrides), client) + return DrillTools(load_config(overrides=overrides, env={}), client) class TestRunQuery: @@ -223,6 +223,43 @@ class TestManagementTools: client.profile.return_value = {"queryId": "abc"} assert make_tools(client).get_profile("abc")["queryId"] == "abc" + def test_get_profile_redacts_secret_looking_keys(self): + # Profiles are cluster-wide: a full profile embeds Drill's + # serialized physical plan, which for JDBC/HTTP plugins can carry + # plugin configuration (passwords, tokens). This must go through the + # same redaction as list_storage_plugins, not be returned unmodified. + client = MagicMock() + client.profile.return_value = {"queryId": "abc", "password": "hunter2"} + result = make_tools(client).get_profile("abc") + assert result["password"] == "***REDACTED***" + assert result["queryId"] == "abc" + + def test_get_profile_is_refused_when_its_query_text_names_a_hidden_schema(self): + # A profile carries the query TEXT of whatever user ran it -- other + # users' queries, not just the caller's own. A hidden schema's name + # can leak out here as data even though it is unreachable directly, + # which is exactly the enumeration path the guard and hidden-schema + # filtering elsewhere were built to close. + client = MagicMock() + client.profile.return_value = {"queryId": "abc", "query": "SELECT * FROM sys.options"} + with pytest.raises(ToolError, match="hidden"): + make_tools(client, hidden_schemas=["sys"]).get_profile("abc") + + def test_list_profiles_redacts_secret_looking_keys(self): + client = MagicMock() + client.profiles.return_value = [{"queryId": "abc", "password": "hunter2"}] + result = make_tools(client).list_profiles() + assert result[0]["password"] == "***REDACTED***" + + def test_list_profiles_drops_entries_whose_query_text_names_a_hidden_schema(self): + client = MagicMock() + client.profiles.return_value = [ + {"queryId": "abc", "query": "SELECT * FROM sys.options"}, + {"queryId": "def", "query": "SELECT * FROM dfs.tmp.x"}, + ] + result = make_tools(client, hidden_schemas=["sys"]).list_profiles() + assert [p["queryId"] for p in result] == ["def"] + def test_cancel_query(self): client = MagicMock() client.cancel_query.return_value = "Cancelled" @@ -540,14 +577,14 @@ class TestShowFiltering: class TestWiring: def test_rest_backend_builds_a_rest_client(self): - assert isinstance(build_client(load_config()), RestClient) + assert isinstance(build_client(load_config(env={})), RestClient) def test_jdbc_backend_builds_a_jdbc_client(self): - cfg = load_config(overrides={"backend": "jdbc", "jdbc_driver_path": "/x.jar"}) + cfg = load_config(overrides={"backend": "jdbc", "jdbc_driver_path": "/x.jar"}, env={}) assert isinstance(build_client(cfg), JdbcClient) def test_all_tools_are_registered(self): - server = build_server(load_config()) + server = build_server(load_config(env={})) names = {tool.name for tool in server._tool_manager.list_tools()} assert names == { "run_query", @@ -562,11 +599,11 @@ class TestWiring: } def test_every_tool_has_a_description(self): - server = build_server(load_config()) + server = build_server(load_config(env={})) 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()) + server = build_server(load_config(env={})) 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"} @@ -574,7 +611,7 @@ class TestWiring: 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()) + server = build_server(load_config(env={})) credential_words = {"user", "password", "username", "passwd", "secret", "token", "credential"} for tool in server._tool_manager.list_tools(): params = set(tool.parameters.get("properties", {}))
