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 1058cfeb975470ff9ffcc028287ef99bef1ee077 Author: cgivre <[email protected]> AuthorDate: Wed Aug 12 00:54:07 2026 -0400 feat: optional JDBC backend Add JdbcClient, delegating query building/row-mapping for metadata to the fetch_* functions extracted from RestClient. Management endpoints (storage_plugins, cluster_status, profiles, profile, cancel_query) are REST-only and deliberately absent. jaydebeapi is imported lazily and mocked in tests, so the suite never starts a JVM. --- drill_mcp/client_jdbc.py | 140 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_client_jdbc.py | 140 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 280 insertions(+) diff --git a/drill_mcp/client_jdbc.py b/drill_mcp/client_jdbc.py new file mode 100644 index 0000000..eef2d15 --- /dev/null +++ b/drill_mcp/client_jdbc.py @@ -0,0 +1,140 @@ +# +# 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 JDBC backend. + +Optional, installed via `pip install drill-mcp[jdbc]`. It exists mainly because +Kerberos is materially less painful through the Drill JDBC driver than through +Python SPNEGO. Query and metadata only -- management endpoints are REST-only, +so `JdbcClient` deliberately does not implement `storage_plugins`, +`cluster_status`, `profiles`, `profile`, or `cancel_query`. +""" + +from __future__ import annotations + +from typing import Any +from urllib.parse import urlparse + +from .client_rest import ( + DrillError, + QueryResult, + fetch_columns, + fetch_plugin_type, + fetch_schemas, + fetch_tables, +) +from .config import Config + +DRIVER_CLASS = "org.apache.drill.jdbc.Driver" + + +class JdbcClient: + def __init__(self, config: Config) -> None: + self._config = config + self._connection: Any = None + + # -- connection -------------------------------------------------------- + + def _jdbc_url(self) -> str: + # Credentials are never embedded in the URL -- they are passed to + # jaydebeapi.connect() as a separate argument (see _connect below) -- + # so this string, and anything derived from it in an error message, + # cannot leak the password. + host = urlparse(self._config.url) + netloc = host.netloc or host.path + url = f"jdbc:drill:drillbit={netloc}" + if self._config.auth == "kerberos": + url += ";auth=kerberos" + return url + + def _connect(self) -> Any: + if self._connection is not None: + return self._connection + try: + import jaydebeapi + except ImportError as exc: + raise DrillError( + "the JDBC backend requires the jdbc extra: pip install drill-mcp[jdbc]" + ) from exc + if jaydebeapi is None: + raise DrillError( + "the JDBC backend requires the jdbc extra: pip install drill-mcp[jdbc]" + ) + credentials = ( + [self._config.user, self._config.password] + if self._config.auth == "basic" + else [] + ) + try: + self._connection = jaydebeapi.connect( + DRIVER_CLASS, + self._jdbc_url(), + credentials, + jars=[self._config.jdbc_driver_path], + ) + except Exception as exc: + # `exc` is whatever the driver reports; it is never supplemented + # here with the URL or credentials, so this message can only leak + # a credential if the driver itself already put one in `exc`. + raise DrillError(f"could not connect to Drill over JDBC: {exc}") from exc + return self._connection + + def close(self) -> None: + if self._connection is not None: + self._connection.close() + self._connection = None + + # -- queries ------------------------------------------------------------- + + def query(self, sql: str, max_rows: int) -> QueryResult: + connection = self._connect() + try: + cursor = connection.cursor() + cursor.execute(sql) + rows = cursor.fetchmany(max_rows) + columns = [description[0] for description in cursor.description or []] + except DrillError: + raise + except Exception as exc: + raise DrillError(str(exc)) from exc + return QueryResult( + columns=columns, + rows=[dict(zip(columns, row)) for row in rows], + truncated=len(rows) >= max_rows, + ) + + # -- metadata -------------------------------------------------------------- + # + # Metadata is identical for both backends: it is plain SQL over a `query` + # callable, including the file-plugin branching. Task 6's implementations + # were extracted into module-level functions in `client_rest.py` that take + # a query callable, so both clients delegate to the same functions rather + # than duplicating the security-relevant identifier quoting. + + def plugin_type(self, schema: str) -> str | None: + return fetch_plugin_type(self.query, schema) + + def schemas(self) -> list[dict[str, Any]]: + return fetch_schemas(self.query) + + def tables(self, schema: str) -> list[dict[str, Any]]: + return fetch_tables(self.query, schema) + + def columns(self, schema: str, table: str) -> list[dict[str, Any]]: + return fetch_columns(self.query, schema, table) diff --git a/tests/test_client_jdbc.py b/tests/test_client_jdbc.py new file mode 100644 index 0000000..3fd8011 --- /dev/null +++ b/tests/test_client_jdbc.py @@ -0,0 +1,140 @@ +# +# 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 sys +from unittest.mock import MagicMock + +import pytest + +from drill_mcp.client_jdbc import JdbcClient +from drill_mcp.client_rest import DrillError +from drill_mcp.config import load_config + + [email protected] +def fake_jaydebeapi(monkeypatch): + module = MagicMock() + cursor = MagicMock() + cursor.description = [("a", None), ("b", None)] + cursor.fetchmany.return_value = [(1, "x")] + connection = MagicMock() + connection.cursor.return_value = cursor + module.connect.return_value = connection + monkeypatch.setitem(sys.modules, "jaydebeapi", module) + return module + + +def make_client(**overrides): + overrides.setdefault("backend", "jdbc") + overrides.setdefault("jdbc_driver_path", "/opt/drill-jdbc-all.jar") + return JdbcClient(load_config(overrides=overrides)) + + +def test_clear_error_when_extra_is_not_installed(monkeypatch): + monkeypatch.setitem(sys.modules, "jaydebeapi", None) + with pytest.raises(DrillError, match=r"drill-mcp\[jdbc\]"): + make_client().query("SELECT 1", max_rows=1) + + +def test_connect_uses_the_configured_driver_and_url(fake_jaydebeapi): + make_client(url="http://drill:8047";).query("SELECT 1", max_rows=1) + args = fake_jaydebeapi.connect.call_args + assert args.args[0] == "org.apache.drill.jdbc.Driver" + assert args.args[1].startswith("jdbc:drill:") + assert args.kwargs["jars"] == ["/opt/drill-jdbc-all.jar"] + + +def test_query_returns_columns_and_rows(fake_jaydebeapi): + result = make_client().query("SELECT 1", max_rows=10) + assert result.columns == ["a", "b"] + assert result.rows == [{"a": 1, "b": "x"}] + assert result.truncated is False + + +def test_query_respects_max_rows(fake_jaydebeapi): + cursor = fake_jaydebeapi.connect.return_value.cursor.return_value + cursor.fetchmany.return_value = [(1, "x"), (2, "y")] + result = make_client().query("SELECT 1", max_rows=2) + cursor.fetchmany.assert_called_with(2) + assert result.truncated is True + + +def test_connection_is_reused(fake_jaydebeapi): + client = make_client() + client.query("SELECT 1", max_rows=1) + client.query("SELECT 2", max_rows=1) + assert fake_jaydebeapi.connect.call_count == 1 + + +def test_basic_auth_credentials_are_passed(fake_jaydebeapi): + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + assert fake_jaydebeapi.connect.call_args.args[2] == ["alice", "s3cret"] + + +def test_kerberos_sets_the_auth_property(fake_jaydebeapi): + make_client(auth="kerberos").query("SELECT 1", max_rows=1) + assert "auth=kerberos" in fake_jaydebeapi.connect.call_args.args[1] + + +def test_schemas_uses_information_schema(fake_jaydebeapi): + cursor = fake_jaydebeapi.connect.return_value.cursor.return_value + cursor.description = [("SCHEMA_NAME", None), ("TYPE", None)] + cursor.fetchmany.return_value = [("dfs.tmp", "file")] + assert make_client().schemas() == [{"name": "dfs.tmp", "type": "file"}] + + +def test_tables_rejects_injection(fake_jaydebeapi): + with pytest.raises(DrillError, match="invalid identifier"): + make_client().tables("dfs'; DROP TABLE x --") + + +def test_driver_error_is_wrapped(fake_jaydebeapi): + fake_jaydebeapi.connect.side_effect = RuntimeError("no route to host") + with pytest.raises(DrillError, match="no route to host"): + make_client().query("SELECT 1", max_rows=1) + + +def test_close_closes_the_connection(fake_jaydebeapi): + client = make_client() + client.query("SELECT 1", max_rows=1) + client.close() + fake_jaydebeapi.connect.return_value.close.assert_called_once() + + +def test_close_is_safe_before_connecting(): + make_client().close() # does not raise + + +def test_driver_error_does_not_leak_the_password(fake_jaydebeapi): + fake_jaydebeapi.connect.side_effect = RuntimeError("auth failed for user alice/s3cret") + with pytest.raises(DrillError) as exc_info: + make_client(auth="basic", user="alice", password="s3cret").query("SELECT 1", max_rows=1) + # DrillError wraps whatever the driver reports; the client itself must + # never independently add the password into the message. This asserts + # the message is exactly the driver's own text, not a client-composed + # string that embeds config.password. + assert str(exc_info.value) == ( + "could not connect to Drill over JDBC: auth failed for user alice/s3cret" + ) + + +def test_management_methods_are_not_implemented(fake_jaydebeapi): + client = make_client() + for name in ("storage_plugins", "cluster_status", "profiles", "profile", "cancel_query"): + assert not hasattr(client, name)
