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 ee18ced2453e2e6f8b4db2c447d7b756eb9abe6b Author: cgivre <[email protected]> AuthorDate: Tue Aug 11 16:43:31 2026 -0400 feat: Drill REST client with auth and query execution Fixes two defects found in the task brief's reference implementation: - _login's URL-based check ("j_security_check" in str(response.url)) false-positives on a successful, unredirected login, since the response URL still points at the login endpoint. Status code alone is the only reliable signal here. - test_connection_failure_message_names_the_url_not_the_password only mocked /query.json; since auth=basic logs in first, the unmocked /j_security_check route raised respx's AllMockedAssertionError instead of exercising the intended DrillError path. Now mocks both endpoints, matching how a real outage would behave. --- drill_mcp/client_rest.py | 202 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_client_rest.py | 174 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 376 insertions(+) diff --git a/drill_mcp/client_rest.py b/drill_mcp/client_rest.py new file mode 100644 index 0000000..4781462 --- /dev/null +++ b/drill_mcp/client_rest.py @@ -0,0 +1,202 @@ +# +# 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. +# + +"""Drill REST backend. + +Talks to a Drillbit's HTTP endpoints. Metadata methods (Task 6) issue +INFORMATION_SCHEMA queries directly and deliberately bypass `guard.py` — the +guard governs SQL that originated from the model, not SQL this module composes +itself. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any + +import httpx + +from .config import Config + +# -- quoting ----------------------------------------------------------------- +# +# Trust boundary: schema and table names arrive from the model and are +# interpolated into query strings (Drill's REST API has no bind parameters). +# `_IDENTIFIER` must reject anything that could break out of the surrounding +# single quotes or change SQL structure: no quotes, backslashes, semicolons, +# whitespace, or empty segments. `+` requires at least one character, so the +# empty string never matches, and `quote_literal_path` splits on "." and +# validates every segment, so a lone "." (which splits into two empty +# segments) is rejected too. +_IDENTIFIER = re.compile(r"^[A-Za-z0-9_$-]+$") + + +class DrillError(Exception): + """Any failure talking to Drill: connection, auth, or query error.""" + + +@dataclass +class QueryResult: + columns: list[str] = field(default_factory=list) + rows: list[dict[str, Any]] = field(default_factory=list) + query_id: str | None = None + truncated: bool = False + + +def quote_literal(value: str) -> str: + """Return one identifier as a single-quoted SQL literal, rejecting unsafe input. + + Trust boundary: schema and table names arrive from the model and are + interpolated into INFORMATION_SCHEMA queries. Drill's REST API has no bind + parameters, so anything outside the safe character set is rejected rather + than escaped. + """ + if not _IDENTIFIER.match(value): + raise DrillError(f"invalid identifier: {value!r}") + return f"'{value}'" + + +def quote_literal_path(value: str) -> str: + """Same as `quote_literal`, but permits a dotted schema path like `dfs.tmp`.""" + parts = value.split(".") + if not parts or any(not _IDENTIFIER.match(part) for part in parts): + raise DrillError(f"invalid identifier: {value!r}") + return f"'{value}'" + + +# -- client -------------------------------------------------------------- + + +class RestClient: + def __init__(self, config: Config) -> None: + self._config = config + self._authenticated = False + auth = None + if config.auth == "kerberos": + try: + from httpx_gssapi import HTTPSPNEGOAuth + except ImportError as exc: # pragma: no cover - exercised in Task 7 style + raise DrillError( + "auth: kerberos requires the kerberos extra: pip install drill-mcp[kerberos]" + ) from exc + auth = HTTPSPNEGOAuth() + self._http = httpx.Client( + base_url=config.url, + timeout=config.timeout_seconds, + auth=auth, + follow_redirects=True, + ) + + def close(self) -> None: + self._http.close() + + # -- transport --------------------------------------------------------- + + def _login(self) -> None: + if self._config.auth != "basic": + self._authenticated = True + return + try: + response = self._http.post( + "/j_security_check", + data={ + "j_username": self._config.user, + "j_password": self._config.password, + }, + ) + except httpx.HTTPError as exc: + raise self._transport_error(exc) from exc + # Drill's j_security_check endpoint (standard Java EE FORM auth) returns + # 200 for both outcomes when no redirect target is configured; failure + # is not reliably distinguishable by URL shape (an unredirected success + # response's URL also still points at "/j_security_check", so that + # check is vacuous at best and a false positive at worst). Status code + # is the only signal we can trust here. + if response.status_code >= 400: + raise DrillError( + f"authentication failed for user {self._config.user!r} at {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"{self._config.timeout_seconds}s" + ) + return DrillError( + f"could not reach Drill at {self._config.url} " + f"(auth mode: {self._config.auth}): {type(exc).__name__}" + ) + + def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response: + if not self._authenticated: + self._login() + try: + response = self._http.request(method, path, **kwargs) + if response.status_code == 401 and self._config.auth == "basic": + self._authenticated = False + self._login() + response = self._http.request(method, path, **kwargs) + except httpx.HTTPError as exc: + raise self._transport_error(exc) from exc + + if response.status_code == 401: + raise DrillError( + f"authentication rejected by Drill at {self._config.url} " + f"for user {self._config.user!r}" + ) + if response.status_code >= 400: + raise DrillError(_error_text(response)) + return response + + # -- queries ----------------------------------------------------------- + + def query(self, sql: str, max_rows: int) -> QueryResult: + response = self._request( + "POST", + "/query.json", + json={"queryType": "SQL", "query": sql, "autoLimit": max_rows}, + ) + payload = response.json() + rows = payload.get("rows") or [] + return QueryResult( + columns=payload.get("columns") or [], + rows=rows, + query_id=payload.get("queryId"), + truncated=len(rows) >= max_rows, + ) + + +def _error_text(response: httpx.Response) -> str: + """Drill's own error text is what a model needs to fix its SQL. Truncate it.""" + try: + payload = response.json() + except ValueError: + payload = None + message = "" + if isinstance(payload, dict): + message = payload.get("errorMessage") or payload.get("message") or "" + if not message: + message = response.text + message = " ".join(message.split()) + if len(message) > 2000: + message = message[:2000] + " ... [truncated]" + return message or f"Drill returned HTTP {response.status_code}" diff --git a/tests/test_client_rest.py b/tests/test_client_rest.py new file mode 100644 index 0000000..affb997 --- /dev/null +++ b/tests/test_client_rest.py @@ -0,0 +1,174 @@ +# +# 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. +# + +import httpx +import pytest +import respx + +from drill_mcp.client_rest import DrillError, RestClient, quote_literal, quote_literal_path +from drill_mcp.config import load_config + +BASE = "http://drill:8047"; + + +def make_client(**overrides): + overrides.setdefault("url", BASE) + return RestClient(load_config(overrides=overrides)) + + +class TestQuoting: + """Trust boundary: schema and table names arrive from the model.""" + + def test_literal_quotes_a_single_identifier(self): + assert quote_literal("foo") == "'foo'" + + def test_literal_path_allows_dots(self): + assert quote_literal_path("dfs.tmp") == "'dfs.tmp'" + + def test_literal_path_allows_ordinary_drill_names(self): + assert quote_literal_path("my_ws.data-2024") == "'my_ws.data-2024'" + + @pytest.mark.parametrize( + "bad", + ["foo'bar", "foo;DROP", "foo bar", "foo\nbar", "", "foo\\bar", "foo`bar", "dfs.tmp"], + ) + def test_literal_rejects_dangerous_input(self, bad): + with pytest.raises(DrillError, match="invalid identifier"): + quote_literal(bad) + + @pytest.mark.parametrize( + "bad", + ["foo'bar", "foo;DROP", "foo bar", "foo\nbar", "", "..", "foo\\bar", "foo`bar"], + ) + def test_literal_path_rejects_dangerous_input(self, bad): + with pytest.raises(DrillError, match="invalid identifier"): + quote_literal_path(bad) + + +class TestQuery: + @respx.mock + def test_posts_sql_with_autolimit(self): + route = respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response( + 200, json={"columns": ["a"], "rows": [{"a": "1"}], "queryId": "q1"} + ) + ) + result = make_client().query("SELECT 1", max_rows=10) + assert route.called + body = route.calls.last.request.read() + assert b'"queryType": "SQL"' in body or b'"queryType":"SQL"' in body + assert b"autoLimit" in body + assert result.columns == ["a"] + assert result.rows == [{"a": "1"}] + assert result.query_id == "q1" + assert result.truncated is False + + @respx.mock + def test_marks_result_truncated_at_the_cap(self): + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response( + 200, json={"columns": ["a"], "rows": [{"a": "1"}, {"a": "2"}]} + ) + ) + assert make_client().query("SELECT 1", max_rows=2).truncated is True + + @respx.mock + def test_drill_error_text_is_surfaced(self): + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response(500, json={"errorMessage": "VALIDATION ERROR: no such table"}) + ) + with pytest.raises(DrillError, match="no such table"): + make_client().query("SELECT * FROM nope", max_rows=10) + + @respx.mock + def test_connection_failure_message_names_the_url_not_the_password(self): + # A real connection failure affects every request to the host, including + # the basic-auth login that precedes the first query, so both endpoints + # must fail the same way for this to simulate a real outage. + respx.post(f"{BASE}/j_security_check").mock(side_effect=httpx.ConnectError("refused")) + respx.post(f"{BASE}/query.json").mock(side_effect=httpx.ConnectError("refused")) + with pytest.raises(DrillError) as exc: + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + assert BASE in str(exc.value) + assert "s3cret" not in str(exc.value) + + @respx.mock + def test_timeout_is_reported_clearly(self): + respx.post(f"{BASE}/query.json").mock(side_effect=httpx.ReadTimeout("slow")) + with pytest.raises(DrillError, match="timed out"): + make_client().query("SELECT 1", max_rows=1) + + +class TestBasicAuth: + @respx.mock + def test_logs_in_before_the_first_query(self): + login = respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response(200, json={"columns": [], "rows": []}) + ) + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + assert login.called + assert b"j_username=alice" in login.calls.last.request.read() + + @respx.mock + def test_session_is_reused_across_queries(self): + login = respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response(200, json={"columns": [], "rows": []}) + ) + client = make_client(auth="basic", user="alice", password="s3cret") + client.query("SELECT 1", max_rows=1) + client.query("SELECT 2", max_rows=1) + assert login.call_count == 1 + + @respx.mock + def test_reauthenticates_once_on_401(self): + login = respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) + query = respx.post(f"{BASE}/query.json").mock( + side_effect=[ + httpx.Response(401), + httpx.Response(200, json={"columns": ["a"], "rows": []}), + ] + ) + result = make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + assert result.columns == ["a"] + assert login.call_count == 2 + assert query.call_count == 2 + + @respx.mock + def test_gives_up_after_one_retry(self): + respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) + respx.post(f"{BASE}/query.json").mock(return_value=httpx.Response(401)) + with pytest.raises(DrillError, match="authentication"): + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + + @respx.mock + def test_login_failure_is_reported(self): + respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(401)) + with pytest.raises(DrillError, match="authentication"): + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + + @respx.mock + def test_no_login_when_auth_is_none(self): + login = respx.post(f"{BASE}/j_security_check").mock(return_value=httpx.Response(200)) + respx.post(f"{BASE}/query.json").mock( + return_value=httpx.Response(200, json={"columns": [], "rows": []}) + ) + make_client().query("SELECT 1", max_rows=1) + assert not login.called
