This is an automated email from the ASF dual-hosted git repository.
cgivre pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/drill-mcp.git
The following commit(s) were added to refs/heads/main by this push:
new 73f84ad Open a JDBC connection per query
73f84ad is described below
commit 73f84adcb6cededbd9219a26ba701df5b80ba4c6
Author: cgivre <[email protected]>
AuthorDate: Sun Oct 11 15:47:52 2026 -0400
Open a JDBC connection per query
The MCP SDK runs sync tools on worker threads, so parallel tool calls
raced on the lazily created shared connection and could leak one. Each
query now opens and closes its own connection, matching the Quarkus JDBC
MCP server.
---
README.md | 3 +--
drill_mcp/client_jdbc.py | 25 +++++-------------
tests/test_client_jdbc.py | 65 ++++++++++++++++++++++++++++++++++-------------
3 files changed, 55 insertions(+), 38 deletions(-)
diff --git a/README.md b/README.md
index e2f8fed..17f95f3 100644
--- a/README.md
+++ b/README.md
@@ -137,8 +137,7 @@ that would share one Drill identity, one policy, and
unrestricted
An HTTP transport is a possible future enhancement. Shipping one safely
needs client authentication, a Drill identity (or impersonation) per
-caller, authorization for `cancel_query`, and confirmation that the JDBC
-backend tolerates concurrent calls.
+caller, and authorization for `cancel_query`.
## Backends
diff --git a/drill_mcp/client_jdbc.py b/drill_mcp/client_jdbc.py
index 8d3ee68..9adc04b 100644
--- a/drill_mcp/client_jdbc.py
+++ b/drill_mcp/client_jdbc.py
@@ -70,7 +70,6 @@ def _type_name(description_entry: tuple) -> str | None:
class JdbcClient:
def __init__(self, config: Config) -> None:
self._config = config
- self._connection: Any = None
# -- connection --------------------------------------------------------
@@ -108,8 +107,10 @@ class JdbcClient:
return text
def _connect(self) -> Any:
- if self._connection is not None:
- return self._connection
+ # A fresh connection per call, like the Quarkus JDBC MCP server. The
+ # MCP SDK runs sync tools on worker threads, so a shared connection
+ # would be raced on; nothing shared means nothing to lock.
+ # ponytail: pays a connect per query; pool if connect latency matters.
try:
import jaydebeapi
except ImportError as exc:
@@ -122,7 +123,7 @@ class JdbcClient:
else []
)
try:
- self._connection = jaydebeapi.connect(
+ return jaydebeapi.connect(
DRIVER_CLASS,
self._jdbc_url(),
credentials,
@@ -132,26 +133,14 @@ class JdbcClient:
raise DrillError(
f"could not connect to Drill over JDBC:
{self._scrub(str(exc))}"
) from exc
- return self._connection
-
- def close(self) -> None:
- if self._connection is None:
- return
- try:
- self._connection.close()
- except Exception as exc:
- raise DrillError(
- f"could not close the JDBC connection: {self._scrub(str(exc))}"
- ) from exc
- finally:
- self._connection = None
# -- queries -------------------------------------------------------------
def query(self, sql: str, max_rows: int) -> QueryResult:
connection = self._connect()
try:
- with closing(connection.cursor()) as cursor:
+ # Closing inside the try routes a close failure through _scrub too.
+ with closing(connection), closing(connection.cursor()) as cursor:
cursor.execute(sql)
rows = cursor.fetchmany(max_rows)
description = cursor.description or []
diff --git a/tests/test_client_jdbc.py b/tests/test_client_jdbc.py
index 452b743..bd19cad 100644
--- a/tests/test_client_jdbc.py
+++ b/tests/test_client_jdbc.py
@@ -18,6 +18,8 @@
#
import sys
+import threading
+import time
from unittest.mock import MagicMock
import pytest
@@ -146,11 +148,51 @@ def test_query_respects_max_rows(fake_jaydebeapi):
assert result.truncated is True
-def test_connection_is_reused(fake_jaydebeapi):
+def test_each_query_opens_and_closes_its_own_connection(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
+ assert fake_jaydebeapi.connect.call_count == 2
+ assert fake_jaydebeapi.connect.return_value.close.call_count == 2
+
+
+def test_connection_is_closed_when_the_query_fails(fake_jaydebeapi):
+ cursor = fake_jaydebeapi.connect.return_value.cursor.return_value
+ cursor.execute.side_effect = RuntimeError("boom")
+ with pytest.raises(DrillError):
+ make_client().query("SELECT 1", max_rows=1)
+ fake_jaydebeapi.connect.return_value.close.assert_called_once()
+
+
+def test_concurrent_queries_do_not_share_a_connection(fake_jaydebeapi):
+ # The MCP SDK runs sync tools on worker threads, so parallel tool calls
+ # reach `query` concurrently. A slow connect makes the calls overlap.
+ connections = []
+
+ def slow_connect(*args, **kwargs):
+ time.sleep(0.05)
+ connection = MagicMock()
+ connection.cursor.return_value.description = [("a", None)]
+ connection.cursor.return_value.fetchmany.return_value = [(1,)]
+ connections.append(connection)
+ return connection
+
+ fake_jaydebeapi.connect.side_effect = slow_connect
+ client = make_client()
+ barrier = threading.Barrier(2)
+
+ def run():
+ barrier.wait()
+ client.query("SELECT 1", max_rows=1)
+
+ threads = [threading.Thread(target=run) for _ in range(2)]
+ for thread in threads:
+ thread.start()
+ for thread in threads:
+ thread.join()
+ assert len(connections) == 2
+ for connection in connections:
+ connection.close.assert_called_once()
def test_basic_auth_credentials_are_passed(fake_jaydebeapi):
@@ -185,17 +227,6 @@ def test_driver_error_is_wrapped(fake_jaydebeapi):
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"
@@ -226,16 +257,14 @@ def
test_query_error_does_not_leak_the_password(fake_jaydebeapi):
def test_close_error_does_not_leak_the_password(fake_jaydebeapi):
- client = make_client(auth="basic", user="alice", password="s3cret")
- client.query("SELECT 1", max_rows=1)
fake_jaydebeapi.connect.return_value.close.side_effect = RuntimeError(
"close failed: password was s3cret"
)
with pytest.raises(DrillError) as exc_info:
- client.close()
+ make_client(auth="basic", user="alice", password="s3cret").query(
+ "SELECT 1", max_rows=1
+ )
assert "s3cret" not in str(exc_info.value)
- # The stale handle must not linger after a failed close.
- assert client._connection is None
def test_management_methods_are_not_implemented(fake_jaydebeapi):
