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 f90261d983c5f07f4351ef33f47d8b4dee7f336a Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 00:24:34 2026 -0400 feat: metadata and management endpoints on the REST client --- drill_mcp/client_rest.py | 151 ++++++++++++++++++++++++++++++ tests/test_client_rest.py | 231 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 382 insertions(+) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py index cfda3a0..6e7142d 100644 --- a/drill_mcp/client_rest.py +++ b/drill_mcp/client_rest.py @@ -34,6 +34,7 @@ from typing import Any import httpx from .config import Config +from .redact import redact # -- quoting ----------------------------------------------------------------- # @@ -114,6 +115,27 @@ def quote_literal_path(value: str) -> str: return f"'{value}'" +def quote_identifier_path(value: str) -> str: + """Validate a dotted path and return it backtick-quoted: dfs.tmp -> `dfs`.`tmp`. + + Same trust boundary as quote_literal_path -- these values arrive from the + model and are interpolated into SQL. Reject rather than escape. + """ + parts = value.split(".") + if any(not _IDENTIFIER.fullmatch(part) for part in parts): + raise DrillError(f"invalid identifier: {value!r}") + return ".".join(f"`{part}`" for part in parts) + + +_QUERY_ID = re.compile(r"[A-Za-z0-9-]+") + + +def _check_query_id(query_id: str) -> str: + if not _QUERY_ID.fullmatch(query_id): + raise DrillError(f"invalid query id: {query_id!r}") + return query_id + + def _error_text(response: httpx.Response) -> str: """Drill's own error text is what a model needs to fix its SQL. Truncate it.""" try: @@ -246,3 +268,132 @@ class RestClient: query_id=payload.get("queryId"), truncated=max_rows > 0 and len(rows) >= max_rows, ) + + # -- metadata ---------------------------------------------------------- + + def schemas(self) -> list[dict[str, Any]]: + result = self.query( + "SELECT SCHEMA_NAME, TYPE FROM INFORMATION_SCHEMA.SCHEMATA " + "ORDER BY SCHEMA_NAME", + max_rows=10_000, + ) + return [ + {"name": row.get("SCHEMA_NAME"), "type": row.get("TYPE")} + for row in result.rows + ] + + def plugin_type(self, schema: str) -> str | None: + """Return the storage plugin TYPE backing `schema`, or None if unknown. + + File-based plugins (`dfs`, `s3`) do not register their contents in + INFORMATION_SCHEMA, so `tables` and `columns` must branch on this. + """ + result = self.query( + "SELECT SCHEMA_NAME, TYPE FROM INFORMATION_SCHEMA.`SCHEMATA` " + f"WHERE SCHEMA_NAME = {quote_literal_path(schema)}", + max_rows=1, + ) + return result.rows[0].get("TYPE") if result.rows else None + + def tables(self, schema: str) -> list[dict[str, Any]]: + # File plugins are absent from INFORMATION_SCHEMA.`TABLES`; querying it + # for `dfs.tmp` returns an empty list that looks like an empty workspace. + # `SHOW FILES` is the only way to enumerate them. sqlalchemy-drill's + # get_table_names branches the same way. + if self.plugin_type(schema) == "file": + result = self.query( + f"SHOW FILES FROM {quote_identifier_path(schema)}", max_rows=10_000 + ) + tables: list[dict[str, Any]] = [] + for row in result.rows: + name = row.get("name") + if not name: + continue + # Drill stores a view as a `<name>.view.drill` file in the workspace. + if name.endswith(".view.drill"): + tables.append({"name": name[: -len(".view.drill")], "type": "VIEW"}) + else: + is_dir = str(row.get("isDirectory", "")).lower() == "true" + tables.append({"name": name, "type": "DIRECTORY" if is_dir else "TABLE"}) + return sorted(tables, key=lambda t: t["name"]) + + result = self.query( + "SELECT TABLE_NAME, TABLE_TYPE FROM INFORMATION_SCHEMA.`TABLES` " + f"WHERE TABLE_SCHEMA = {quote_literal_path(schema)} ORDER BY TABLE_NAME", + max_rows=10_000, + ) + return [ + {"name": row.get("TABLE_NAME"), "type": row.get("TABLE_TYPE")} + for row in result.rows + ] + + def columns(self, schema: str, table: str) -> list[dict[str, Any]]: + # Validate the table name up front, before the plugin_type lookup fires + # a query: an invalid table name should never make it to the network. + # File-plugin table names are filenames and may contain a literal "." + # (e.g. "sales.csv"), so validate segment-by-segment like a dotted path. + if any(not _IDENTIFIER.fullmatch(part) for part in table.split(".")): + raise DrillError(f"invalid identifier: {table!r}") + + # Same split: file plugins have dynamic schemas and no + # INFORMATION_SCHEMA.`COLUMNS` rows. DESCRIBE is metadata-only -- + # deliberately NOT a `SELECT * ... LIMIT 1` probe, which would read user + # data to answer a metadata question. + if self.plugin_type(schema) == "file": + result = self.query( + f"DESCRIBE {quote_identifier_path(schema + '.' + table)}", + max_rows=10_000, + ) + return [ + { + "name": row.get("COLUMN_NAME"), + "data_type": row.get("DATA_TYPE"), + "nullable": str(row.get("IS_NULLABLE", "")).upper() == "YES", + } + for row in result.rows + ] + + result = self.query( + "SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE FROM INFORMATION_SCHEMA.`COLUMNS` " + f"WHERE TABLE_SCHEMA = {quote_literal_path(schema)} " + f"AND TABLE_NAME = {quote_literal(table)} ORDER BY ORDINAL_POSITION", + max_rows=10_000, + ) + return [ + { + "name": row.get("COLUMN_NAME"), + "data_type": row.get("DATA_TYPE"), + "nullable": str(row.get("IS_NULLABLE", "")).upper() == "YES", + } + for row in result.rows + ] + + # -- management -------------------------------------------------------- + + def storage_plugins(self) -> list[dict[str, Any]]: + payload = self._request("GET", "/storage.json").json() + return redact(payload) + + def cluster_status(self) -> dict[str, Any]: + cluster = self._request("GET", "/cluster.json").json() + status = self._request("GET", "/status.json").json() + merged = dict(cluster) if isinstance(cluster, dict) else {"cluster": cluster} + if isinstance(status, dict): + merged.update(status) + else: + merged["status"] = status + return merged + + def profiles(self, limit: int) -> list[dict[str, Any]]: + payload = self._request("GET", "/profiles.json").json() + running = payload.get("runningQueries") or [] + finished = payload.get("finishedQueries") or [] + return (list(running) + list(finished))[:limit] + + def profile(self, query_id: str) -> dict[str, Any]: + _check_query_id(query_id) + return self._request("GET", f"/profiles/{query_id}.json").json() + + def cancel_query(self, query_id: str) -> str: + _check_query_id(query_id) + return self._request("GET", f"/profiles/cancel/{query_id}").text diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py index 87a0adb..fb804d0 100644 --- a/tests/test_client_rest.py +++ b/tests/test_client_rest.py @@ -376,3 +376,234 @@ class TestKerberosAuth: monkeypatch.setitem(sys.modules, "httpx_gssapi", stub) client = make_client(auth="kerberos") assert client._http.auth is sentinel + + +def query_response(columns, rows): + return httpx.Response(200, json={"columns": columns, "rows": rows, "queryId": "q"}) + + +class TestMetadata: + @respx.mock + def test_schemas_queries_information_schema(self): + route = respx.post(f"{BASE}/query.json").mock( + return_value=query_response( + ["SCHEMA_NAME", "TYPE"], [{"SCHEMA_NAME": "dfs.tmp", "TYPE": "file"}] + ) + ) + assert make_client().schemas() == [{"name": "dfs.tmp", "type": "file"}] + assert b"INFORMATION_SCHEMA" in route.calls.last.request.read() + + @respx.mock + def test_tables_filters_by_schema(self): + route = respx.post(f"{BASE}/query.json").mock( + return_value=query_response( + ["TABLE_NAME", "TABLE_TYPE"], [{"TABLE_NAME": "t", "TABLE_TYPE": "TABLE"}] + ) + ) + assert make_client().tables("dfs.tmp") == [{"name": "t", "type": "TABLE"}] + assert b"'dfs.tmp'" in route.calls.last.request.read() + + @respx.mock + def test_columns_returns_name_type_nullable(self): + respx.post(f"{BASE}/query.json").mock( + return_value=query_response( + ["COLUMN_NAME", "DATA_TYPE", "IS_NULLABLE"], + [{"COLUMN_NAME": "id", "DATA_TYPE": "INTEGER", "IS_NULLABLE": "YES"}], + ) + ) + assert make_client().columns("dfs.tmp", "t") == [ + {"name": "id", "data_type": "INTEGER", "nullable": True} + ] + + @respx.mock + def test_metadata_rejects_injection_in_schema_name(self): + with pytest.raises(DrillError, match="invalid identifier"): + make_client().tables("dfs'; DROP TABLE x --") + + +class TestFilePluginMetadata: + """File plugins are absent from INFORMATION_SCHEMA; they need SHOW FILES.""" + + @staticmethod + def _schemata(plugin_type): + return query_response(["SCHEMA_NAME", "TYPE"], [{"SCHEMA_NAME": "dfs.tmp", "TYPE": plugin_type}]) + + @respx.mock + def test_tables_uses_show_files_for_a_file_plugin(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response(["name", "isDirectory"], [{"name": "sales.csv", "isDirectory": "false"}]), + ] + ) + assert make_client().tables("dfs.tmp") == [{"name": "sales.csv", "type": "TABLE"}] + assert b"SHOW FILES FROM" in route.calls[1].request.read() + + @respx.mock + def test_show_files_marks_directories(self): + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response(["name", "isDirectory"], [{"name": "year=2024", "isDirectory": "true"}]), + ] + ) + assert make_client().tables("dfs.tmp")[0]["type"] == "DIRECTORY" + + @respx.mock + def test_show_files_strips_the_view_drill_suffix(self): + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response(["name", "isDirectory"], [{"name": "top_sales.view.drill", "isDirectory": "false"}]), + ] + ) + assert make_client().tables("dfs.tmp") == [{"name": "top_sales", "type": "VIEW"}] + + @respx.mock + def test_tables_uses_information_schema_for_a_non_file_plugin(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("jdbc"), + query_response(["TABLE_NAME", "TABLE_TYPE"], [{"TABLE_NAME": "t", "TABLE_TYPE": "TABLE"}]), + ] + ) + assert make_client().tables("mysql.app") == [{"name": "t", "type": "TABLE"}] + assert b"INFORMATION_SCHEMA" in route.calls[1].request.read() + + @respx.mock + def test_columns_uses_describe_for_a_file_plugin(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("file"), + query_response( + ["COLUMN_NAME", "DATA_TYPE", "IS_NULLABLE"], + [{"COLUMN_NAME": "id", "DATA_TYPE": "BIGINT", "IS_NULLABLE": "YES"}], + ), + ] + ) + assert make_client().columns("dfs.tmp", "sales.csv") == [ + {"name": "id", "data_type": "BIGINT", "nullable": True} + ] + body = route.calls[1].request.read() + assert b"DESCRIBE" in body + # Metadata-only: never read user rows to answer a metadata question. + assert b"LIMIT 1" not in body + assert b"SELECT *" not in body + + @respx.mock + def test_columns_uses_information_schema_for_a_non_file_plugin(self): + route = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + self._schemata("jdbc"), + query_response( + ["COLUMN_NAME", "DATA_TYPE", "IS_NULLABLE"], + [{"COLUMN_NAME": "id", "DATA_TYPE": "INTEGER", "IS_NULLABLE": "NO"}], + ), + ] + ) + result = make_client().columns("mysql.app", "t") + assert result == [{"name": "id", "data_type": "INTEGER", "nullable": False}] + assert b"INFORMATION_SCHEMA" in route.calls[1].request.read() + + @respx.mock + def test_unknown_plugin_type_falls_back_to_information_schema(self): + respx.post(f"{BASE}/query.json").mock( + side_effect=[ + query_response(["SCHEMA_NAME", "TYPE"], []), + query_response(["TABLE_NAME", "TABLE_TYPE"], []), + ] + ) + assert make_client().tables("nope") == [] + + @respx.mock + def test_show_files_path_still_rejects_injection(self): + respx.post(f"{BASE}/query.json").mock(return_value=self._schemata("file")) + with pytest.raises(DrillError, match="invalid identifier"): + make_client().tables("dfs`; DROP TABLE x --") + + @respx.mock + def test_metadata_rejects_injection_in_table_name(self): + with pytest.raises(DrillError, match="invalid identifier"): + make_client().columns("dfs.tmp", "t' OR '1'='1") + + +class TestManagement: + @respx.mock + def test_storage_plugins_are_redacted(self): + respx.get(f"{BASE}/storage.json").mock( + return_value=httpx.Response( + 200, + json=[ + { + "name": "s3", + "config": {"type": "file", "fs.s3a.secret.key": "verysecret"}, + } + ], + ) + ) + plugins = make_client().storage_plugins() + assert plugins[0]["config"]["fs.s3a.secret.key"] == "***REDACTED***" + assert plugins[0]["name"] == "s3" + + @respx.mock + def test_cluster_status_merges_cluster_and_status(self): + respx.get(f"{BASE}/cluster.json").mock( + return_value=httpx.Response(200, json={"drillbits": [{"address": "n1"}]}) + ) + respx.get(f"{BASE}/status.json").mock( + return_value=httpx.Response(200, json={"status": "Running!"}) + ) + result = make_client().cluster_status() + assert result["drillbits"] == [{"address": "n1"}] + assert result["status"] == "Running!" + + @respx.mock + def test_profiles_are_limited(self): + respx.get(f"{BASE}/profiles.json").mock( + return_value=httpx.Response( + 200, + json={ + "finishedQueries": [{"queryId": f"q{i}"} for i in range(10)], + "runningQueries": [], + }, + ) + ) + assert len(make_client().profiles(limit=3)) == 3 + + @respx.mock + def test_profiles_include_running_queries_first(self): + respx.get(f"{BASE}/profiles.json").mock( + return_value=httpx.Response( + 200, + json={ + "runningQueries": [{"queryId": "live"}], + "finishedQueries": [{"queryId": "done"}], + }, + ) + ) + assert make_client().profiles(limit=5)[0]["queryId"] == "live" + + @respx.mock + def test_profile_fetches_one_query(self): + respx.get(f"{BASE}/profiles/abc.json").mock( + return_value=httpx.Response(200, json={"queryId": "abc", "state": "COMPLETED"}) + ) + assert make_client().profile("abc")["state"] == "COMPLETED" + + @respx.mock + def test_profile_rejects_a_malformed_query_id(self): + with pytest.raises(DrillError, match="invalid"): + make_client().profile("../../etc/passwd") + + @respx.mock + def test_cancel_query(self): + route = respx.get(f"{BASE}/profiles/cancel/abc").mock( + return_value=httpx.Response(200, text="Cancelled query abc") + ) + assert "abc" in make_client().cancel_query("abc") + assert route.called + + @respx.mock + def test_cancel_rejects_a_malformed_query_id(self): + with pytest.raises(DrillError, match="invalid"): + make_client().cancel_query("abc; rm -rf /")
