diff --git a/CHANGELOG.md b/CHANGELOG.md index a8b84a44..2547264a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,12 +5,17 @@ All notable changes to this project will be documented in this file. ## Unreleased ### Added +- Client connection reuse across tool calls via a config-keyed cache, eliminating per-call connection overhead. ([#152](https://github.com/ClickHouse/mcp-clickhouse/pull/152)) +- Server-side query cancellation: timed-out queries now issue `KILL QUERY` on the ClickHouse server instead of leaving zombie workers consuming threads and server resources. ([#152](https://github.com/ClickHouse/mcp-clickhouse/pull/152)) +- `CLICKHOUSE_MCP_MAX_WORKERS` environment variable to configure the query worker thread pool size (default: `10`). ([#152](https://github.com/ClickHouse/mcp-clickhouse/pull/152)) - Support for FastMCP OAuth/OIDC auth providers on HTTP/SSE transports via the `FASTMCP_SERVER_AUTH` environment variable (e.g. Azure Entra, Google, GitHub, WorkOS). Static token, FastMCP OAuth, and disabled mode are now mutually exclusive; configure exactly one. ([#171](https://github.com/ClickHouse/mcp-clickhouse/issues/171)) ### Changed +- `CLICKHOUSE_SEND_RECEIVE_TIMEOUT` is now auto-capped to `CLICKHOUSE_MCP_QUERY_TIMEOUT + 5` unless explicitly set, so HTTP reads unblock shortly after an MCP timeout fires. ([#152](https://github.com/ClickHouse/mcp-clickhouse/pull/152)) - `/health` endpoint is now unauthenticated across all auth modes (previously gated only under static-token mode, which was asymmetric and incompatible with redirect-based OAuth providers). Response bodies trimmed to `OK` / generic error strings to avoid leaking ClickHouse version information or connection exception details; underlying errors are logged server-side. ### Fixed +- Session config overrides from PR #115 are now resolved on the request thread where the FastMCP context is available, so overrides are correctly applied to queries dispatched to the worker pool. ([#152](https://github.com/ClickHouse/mcp-clickhouse/pull/152)) - Tool responses now return JSON-encoded strings, avoiding MCP protocol validation errors on successful queries. ([#154](https://github.com/ClickHouse/mcp-clickhouse/pull/154)) - Long-running queries no longer block other tool calls. The MCP-facing `run_query` and `run_chdb_select_query` tools now await their thread-pool futures asynchronously, so concurrent tool calls are served while a slow query is in flight. ([#128](https://github.com/ClickHouse/mcp-clickhouse/issues/128)) diff --git a/README.md b/README.md index 2358b4a3..83ab6b74 100644 --- a/README.md +++ b/README.md @@ -516,9 +516,9 @@ The following environment variables are used to configure the ClickHouse and chD * `CLICKHOUSE_CONNECT_TIMEOUT`: Connection timeout in seconds * Default: `"30"` * Increase this value if you experience connection timeouts -* `CLICKHOUSE_SEND_RECEIVE_TIMEOUT`: Send/receive timeout in seconds - * Default: `"300"` - * Increase this value for long-running queries +* `CLICKHOUSE_SEND_RECEIVE_TIMEOUT`: Send/receive timeout in seconds for the underlying HTTP connection + * Default: automatically set to `CLICKHOUSE_MCP_QUERY_TIMEOUT + 5` so worker threads unblock shortly after a query timeout + * If explicitly set, the value is used as-is (e.g. `"300"` for long-running queries) * `CLICKHOUSE_DATABASE`: Default database to use * Default: None (uses server default) * Set this to automatically connect to a specific database @@ -532,9 +532,14 @@ The following environment variables are used to configure the ClickHouse and chD * `CLICKHOUSE_MCP_BIND_PORT`: Port to bind the MCP server to when using HTTP or SSE transport * Default: `"8000"` * Only used when transport is `"http"` or `"sse"` -* `CLICKHOUSE_MCP_QUERY_TIMEOUT`: Timeout in seconds for SELECT tools +* `CLICKHOUSE_MCP_QUERY_TIMEOUT`: Timeout in seconds for query tool calls * Default: `"30"` * Increase this if you see `Query timed out after ...` errors for heavy queries + * When a query times out, the server issues a `KILL QUERY` on the ClickHouse server to cancel it + * Unless `CLICKHOUSE_SEND_RECEIVE_TIMEOUT` is explicitly set, the HTTP read timeout is automatically aligned to this value plus a small buffer, so worker threads unblock shortly after a timeout +* `CLICKHOUSE_MCP_MAX_WORKERS`: Maximum number of concurrent query worker threads + * Default: `"10"` + * Increase if your workload requires many concurrent tool calls * `CLICKHOUSE_MCP_AUTH_TOKEN`: Static bearer token for HTTP/SSE transports * Default: None * One of `CLICKHOUSE_MCP_AUTH_TOKEN`, `FASTMCP_SERVER_AUTH`, or `CLICKHOUSE_MCP_AUTH_DISABLED=true` is **required** for HTTP/SSE transports diff --git a/mcp_clickhouse/mcp_env.py b/mcp_clickhouse/mcp_env.py index 0492ca6a..07f0fd3a 100644 --- a/mcp_clickhouse/mcp_env.py +++ b/mcp_clickhouse/mcp_env.py @@ -300,6 +300,7 @@ class MCPServerConfig: CLICKHOUSE_MCP_BIND_HOST: Bind host for HTTP/SSE (default: 127.0.0.1) CLICKHOUSE_MCP_BIND_PORT: Bind port for HTTP/SSE (default: 8000) CLICKHOUSE_MCP_QUERY_TIMEOUT: SELECT tool timeout in seconds (default: 30) + CLICKHOUSE_MCP_MAX_WORKERS: Maximum thread pool workers for query execution (default: 10) CLICKHOUSE_MCP_AUTH_TOKEN: Static bearer token for HTTP/SSE transports. One authentication mode must be configured for HTTP/SSE; the other two options are FASTMCP_SERVER_AUTH (FastMCP OAuth/OIDC providers) and @@ -328,6 +329,14 @@ def bind_port(self) -> int: def query_timeout(self) -> int: return int(os.getenv("CLICKHOUSE_MCP_QUERY_TIMEOUT", "30")) + @property + def max_workers(self) -> int: + """Maximum thread pool workers for query execution. + + Default: 10 + """ + return int(os.getenv("CLICKHOUSE_MCP_MAX_WORKERS", "10")) + @property def auth_token(self) -> Optional[str]: """Get the authentication token for HTTP/SSE transports.""" diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index 6758b020..18f4c804 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -5,9 +5,11 @@ import logging import os import re +import threading +import time import uuid from dataclasses import asdict, dataclass, field -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import clickhouse_connect from cachetools import TTLCache @@ -67,11 +69,24 @@ class Table: ) logger = logging.getLogger(MCP_SERVER_NAME) -QUERY_EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=10) -atexit.register(lambda: QUERY_EXECUTOR.shutdown(wait=True)) - load_dotenv() +_max_workers = get_mcp_config().max_workers +QUERY_EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=_max_workers) + +# --- Client cache --- +# Cache of ClickHouse clients keyed by frozen config, enabling client reuse +# across tool calls. Each entry is (client, last_used_timestamp). +_client_cache: Dict[Tuple, Tuple] = {} +_client_cache_lock = threading.Lock() +_CLIENT_IDLE_PING_THRESHOLD = 60 # seconds before we ping to verify liveness + +# --- Active query tracker --- +# Maps query_id -> (cache_key, query_text) so we can KILL QUERY on the +# correct server when a timeout fires. +_active_queries: Dict[str, Tuple] = {} +_active_queries_lock = threading.Lock() + _HTTP_TRANSPORTS = (TransportType.HTTP.value, TransportType.SSE.value) @@ -190,8 +205,19 @@ def _serialize_tool_result(obj: Any) -> str: def list_databases() -> str: """List available ClickHouse databases""" logger.info("Listing all databases") - client = create_clickhouse_client() - result = client.command("SHOW DATABASES") + config = _resolve_client_config() + + for attempt in range(2): + try: + client = create_clickhouse_client(config=config) + result = client.command("SHOW DATABASES") + break + except Exception as err: + if attempt == 0 and _is_connection_error(err): + logger.warning("list_databases connection error, retrying: %s", err) + _evict_cached_client(config) + continue + raise # Convert newline-separated string to list and trim whitespace if isinstance(result, str): @@ -366,8 +392,33 @@ def list_tables( page_size, include_detailed_columns, ) - client = create_clickhouse_client() + config = _resolve_client_config() + for attempt in range(2): + try: + client = create_clickhouse_client(config=config) + return _list_tables_impl( + client, database, like, not_like, page_token, + page_size, include_detailed_columns, + ) + except Exception as err: + if attempt == 0 and _is_connection_error(err): + logger.warning("list_tables connection error, retrying: %s", err) + _evict_cached_client(config) + continue + raise + + +def _list_tables_impl( + client, + database: str, + like: Optional[str], + not_like: Optional[str], + page_token: Optional[str], + page_size: int, + include_detailed_columns: bool, +) -> Dict[str, Any]: + """Inner implementation of list_tables, separated for retry logic.""" if page_token and page_token in table_pagination_cache: cached_state = table_pagination_cache[page_token] cached_include_detailed = cached_state.get("include_detailed_columns", True) @@ -478,20 +529,88 @@ def _validate_query_for_destructive_ops(query: str) -> None: ) -def execute_query(query: str) -> str: - client = create_clickhouse_client() +def _is_connection_error(err: Exception) -> bool: + """Check if an exception indicates a broken connection rather than a query error.""" + from clickhouse_connect.driver.exceptions import OperationalError + if isinstance(err, (OSError, ConnectionError, OperationalError)): + return True + err_str = str(err).lower() + return any(s in err_str for s in ("connection", "timed out", "reset by peer", "eof")) + + +def execute_query(query: str, query_id: str, client_config: dict) -> str: + """Execute a query in a worker thread. + + Args: + query: SQL to execute. + query_id: Unique identifier for server-side tracking / cancellation. + client_config: Pre-resolved config dict (resolved on the request thread). + """ + cache_key = _config_to_cache_key(client_config) + with _active_queries_lock: + _active_queries[query_id] = (cache_key, query) + try: + client = create_clickhouse_client(config=client_config) _validate_query_for_destructive_ops(query) query_settings = build_query_settings(client) + query_settings["query_id"] = query_id res = client.query(query, settings=query_settings) - logger.info(f"Query returned {len(res.result_rows)} rows") + logger.info(f"Query {query_id} returned {len(res.result_rows)} rows") return _serialize_tool_result({"columns": res.column_names, "rows": res.result_rows}) except ToolError: raise except Exception as err: - logger.error(f"Error executing query: {err}") + # Evict the cached client on connection errors so the next call + # creates a fresh one. We do NOT retry here because the query may + # involve writes and retrying could duplicate side effects. + if _is_connection_error(err): + _evict_cached_client(client_config) + logger.error(f"Error executing query {query_id}: {err}") raise ToolError(f"Query execution failed: {str(err)}") + finally: + with _active_queries_lock: + _active_queries.pop(query_id, None) + + +def _cancel_query(query_id: str): + """Issue KILL QUERY on the ClickHouse server for a timed-out query. + + Uses the same cached client (same server/credentials) that originated + the query. Failures are logged but never raised — cancellation errors + must not mask the original timeout. + """ + with _active_queries_lock: + entry = _active_queries.pop(query_id, None) + + if entry is None: + logger.debug("Query %s already completed, nothing to cancel", query_id) + return + + cache_key, _query_text = entry + try: + safe_id = str(uuid.UUID(query_id)) + except ValueError: + logger.warning("Refusing to KILL QUERY with non-UUID query_id: %r", query_id) + return + + try: + with _client_cache_lock: + cached = _client_cache.get(cache_key) + if cached is None: + logger.warning( + "No cached client for query %s cancel — server-side query may still run", + safe_id, + ) + return + + client, _ = cached + logger.info("Cancelling query %s via KILL QUERY", safe_id) + client.command(f"KILL QUERY WHERE query_id = '{safe_id}'") + logger.info("Successfully cancelled query %s", safe_id) + except Exception as e: + logger.warning("Failed to cancel query %s: %s", safe_id, e) def run_query(query: str) -> str: @@ -501,14 +620,29 @@ def run_query(query: str) -> str: to allow DDL and DML statements when your ClickHouse server permits them. """ logger.info(f"Executing query: {query}") + + # Resolve config on the request thread where FastMCP Context is available + client_config = _resolve_client_config() + query_id = str(uuid.uuid4()) + try: - future = QUERY_EXECUTOR.submit(execute_query, query) + with _active_queries_lock: + in_flight = len(_active_queries) + if in_flight >= _max_workers: + logger.warning( + "Thread pool saturated: %d in-flight vs %d workers", + in_flight, _max_workers, + ) + + future = QUERY_EXECUTOR.submit(execute_query, query, query_id, client_config) timeout_secs = get_mcp_config().query_timeout try: return future.result(timeout=timeout_secs) except concurrent.futures.TimeoutError: - logger.warning(f"Query timed out after {timeout_secs} seconds: {query}") - future.cancel() + logger.warning( + "Query %s timed out after %s seconds: %s", query_id, timeout_secs, query + ) + _cancel_query(query_id) raise ToolError(f"Query timed out after {timeout_secs} seconds") except ToolError: raise @@ -517,19 +651,50 @@ def run_query(query: str) -> str: raise RuntimeError(f"Unexpected error during query execution: {str(e)}") +def _config_to_cache_key(config: dict) -> tuple: + """Convert a client config dict into a hashable cache key. + + Handles nested dicts (e.g. 'settings') by recursively sorting items. + """ + items = [] + for k, v in sorted(config.items()): + if isinstance(v, dict): + v = _config_to_cache_key(v) + items.append((k, v)) + return tuple(items) + + async def run_query_async(query: str) -> str: - """Async MCP-facing wrapper for ClickHouse queries.""" + """Async MCP-facing wrapper for ClickHouse queries. + + Awaits the worker-pool future asynchronously so concurrent tool calls are + served while a slow query is in flight. + """ logger.info(f"Executing query: {query}") + + client_config = _resolve_client_config() + query_id = str(uuid.uuid4()) + try: - future = QUERY_EXECUTOR.submit(execute_query, query) + with _active_queries_lock: + in_flight = len(_active_queries) + if in_flight >= _max_workers: + logger.warning( + "Thread pool saturated: %d in-flight vs %d workers", + in_flight, _max_workers, + ) + + future = QUERY_EXECUTOR.submit(execute_query, query, query_id, client_config) timeout_secs = get_mcp_config().query_timeout try: return await asyncio.wait_for( asyncio.wrap_future(future), timeout=timeout_secs ) except asyncio.TimeoutError: - logger.warning(f"Query timed out after {timeout_secs} seconds: {query}") - future.cancel() + logger.warning( + "Query %s timed out after %s seconds: %s", query_id, timeout_secs, query + ) + _cancel_query(query_id) raise ToolError(f"Query timed out after {timeout_secs} seconds") except ToolError: raise @@ -538,8 +703,15 @@ async def run_query_async(query: str) -> str: raise RuntimeError(f"Unexpected error during query execution: {str(e)}") -def create_clickhouse_client(): +def _resolve_client_config() -> dict: + """Build the merged client config on the request thread. + + Must be called from the request thread where FastMCP Context is available. + Merges base config with any per-session overrides, then aligns + send_receive_timeout with the MCP query timeout. + """ client_config = get_config().get_client_config() + srt_explicitly_set = "CLICKHOUSE_SEND_RECEIVE_TIMEOUT" in os.environ try: ctx = get_context() @@ -552,36 +724,174 @@ def create_clickhouse_client(): logger.debug( f"Applying session-specific ClickHouse client config overrides: {list(session_config_overrides.keys())}" ) + if "send_receive_timeout" in session_config_overrides: + srt_explicitly_set = True client_config.update(session_config_overrides) except RuntimeError: - # If we're outside a request context, just proceed with the default config + # Outside a request context — proceed with base config pass + # Align send_receive_timeout with MCP query timeout so worker threads + # unblock shortly after the MCP-level timeout fires, preventing zombie threads. + # Only auto-cap when neither env var nor session override explicitly set it. + if not srt_explicitly_set: + query_timeout = get_mcp_config().query_timeout + effective_srt = client_config.get("send_receive_timeout", 300) + if effective_srt > query_timeout + 5: + client_config["send_receive_timeout"] = query_timeout + 5 + + return client_config + + +def _evict_cached_client(config: dict) -> None: + """Evict a cached client for the given config, closing it. + + Call this when a query or command fails with a connection error so the + next call creates a fresh client instead of reusing the broken one. + """ + cache_key = _config_to_cache_key(config) + with _client_cache_lock: + entry = _client_cache.pop(cache_key, None) + if entry is not None: + client, _ = entry + logger.info("Evicted stale cached client for %s", config.get("host", "?")) + try: + client.close() + except Exception: + pass + + +def create_clickhouse_client(config: Optional[dict] = None): + """Get or create a cached ClickHouse client for the given config. + + Args: + config: Pre-resolved client config dict. When None the config is + resolved from env + session overrides (requires request context). + Pass an explicit config when calling from a worker thread. + """ + if config is None: + config = _resolve_client_config() + + cache_key = _config_to_cache_key(config) + + # Check cache — extract candidate without holding the lock during ping + candidate = None + with _client_cache_lock: + if cache_key in _client_cache: + client, last_used = _client_cache[cache_key] + if time.time() - last_used > _CLIENT_IDLE_PING_THRESHOLD: + candidate = client + else: + _client_cache[cache_key] = (client, time.time()) + logger.debug("Reusing cached client") + return client + + # Ping outside the lock so we don't serialize unrelated configs + if candidate is not None: + try: + alive = candidate.ping() + except Exception: + alive = False + + # Identity-check under the lock: another thread may have replaced + # the cached client while we were pinging — never overwrite or evict + # a newer entry based on stale candidate state. + if alive: + with _client_cache_lock: + entry = _client_cache.get(cache_key) + if entry is not None and entry[0] is candidate: + _client_cache[cache_key] = (candidate, time.time()) + logger.debug("Reusing cached client (ping OK after idle)") + return candidate + if entry is not None: + current_client = entry[0] + _client_cache[cache_key] = (current_client, time.time()) + else: + current_client = None + if current_client is not None: + try: + candidate.close() + except Exception: + pass + logger.debug("Reusing cached client (another thread replaced idle client)") + return current_client + else: + logger.warning("Cached client failed ping, creating new client") + with _client_cache_lock: + entry = _client_cache.get(cache_key) + if entry is not None and entry[0] is candidate: + _client_cache.pop(cache_key, None) + try: + candidate.close() + except Exception: + pass + + # Create new client outside the lock (client creation is slow) config_fields = [ - f"secure={client_config['secure']}", - f"verify={client_config['verify']}", - f"connect_timeout={client_config['connect_timeout']}s", - f"send_receive_timeout={client_config['send_receive_timeout']}s", + f"secure={config['secure']}", + f"verify={config['verify']}", + f"connect_timeout={config['connect_timeout']}s", + f"send_receive_timeout={config['send_receive_timeout']}s", ] - if "server_host_name" in client_config: - config_fields.append(f"server_host_name={client_config['server_host_name']}") - log_msg = ( - f"Creating ClickHouse client connection to {client_config['host']}:{client_config['port']} " - f"as {client_config['username']} " + if "server_host_name" in config: + config_fields.append(f"server_host_name={config['server_host_name']}") + logger.info( + f"Creating ClickHouse client connection to {config['host']}:{config['port']} " + f"as {config['username']} " f"({', '.join(config_fields)})" ) - logger.info(log_msg) try: - client = clickhouse_connect.get_client(**client_config) - # Test the connection + # Disable autogenerate_session_id so the client is safe for concurrent + # use from the thread pool. clickhouse_connect rejects concurrent queries + # on the same session_id, but with this disabled each query runs + # without session affinity. + client = clickhouse_connect.get_client( + **config, autogenerate_session_id=False + ) version = client.server_version logger.info(f"Successfully connected to ClickHouse server version {version}") - return client except Exception as e: logger.error(f"Failed to connect to ClickHouse: {str(e)}") raise + with _client_cache_lock: + # Another thread may have raced and cached a client for this key + if cache_key in _client_cache: + try: + client.close() + except Exception: + pass + client, _ = _client_cache[cache_key] + _client_cache[cache_key] = (client, time.time()) + return client + _client_cache[cache_key] = (client, time.time()) + + return client + + +def _clear_client_cache(): + """Clear the client cache, closing all cached clients. + + Used during shutdown and for testing. + """ + with _client_cache_lock: + for _, (client, _) in list(_client_cache.items()): + try: + client.close() + except Exception: + pass + _client_cache.clear() + + +def _shutdown(): + # Order matters: drain workers before closing the clients they hold. + QUERY_EXECUTOR.shutdown(wait=True) + _clear_client_cache() + + +atexit.register(_shutdown) + def build_query_settings(client) -> dict[str, str]: """Build query settings dict for ClickHouse queries. diff --git a/tests/test_client_cache.py b/tests/test_client_cache.py new file mode 100644 index 00000000..b52e0adf --- /dev/null +++ b/tests/test_client_cache.py @@ -0,0 +1,365 @@ +"""Tests for ClickHouse client caching and reuse.""" + +import time +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp.exceptions import ToolError + +from mcp_clickhouse.mcp_env import get_mcp_config +from mcp_clickhouse.mcp_server import ( + _active_queries, + _active_queries_lock, + _clear_client_cache, + _client_cache, + _client_cache_lock, + _config_to_cache_key, + _resolve_client_config, + _shutdown, + create_clickhouse_client, + execute_query, +) + + +class TestConfigToCacheKey: + """Tests for the _config_to_cache_key helper.""" + + def test_deterministic_key(self): + config = {"host": "localhost", "port": 8443, "username": "default"} + assert _config_to_cache_key(config) == _config_to_cache_key(config) + + def test_order_independent(self): + config_a = {"host": "localhost", "port": 8443} + config_b = {"port": 8443, "host": "localhost"} + assert _config_to_cache_key(config_a) == _config_to_cache_key(config_b) + + def test_nested_dict(self): + config = {"host": "localhost", "settings": {"role": "admin", "readonly": "1"}} + key = _config_to_cache_key(config) + assert isinstance(key, tuple) + # Nested dict should also be a tuple + for k, v in key: + if k == "settings": + assert isinstance(v, tuple) + + def test_different_configs_different_keys(self): + config_a = {"host": "host1", "port": 8443} + config_b = {"host": "host2", "port": 8443} + assert _config_to_cache_key(config_a) != _config_to_cache_key(config_b) + + +class TestClientCaching: + """Tests for client cache behavior.""" + + def setup_method(self): + _clear_client_cache() + + def teardown_method(self): + _clear_client_cache() + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_same_config_returns_cached_client(self, _mock_ctx, mock_cc): + """Same config should return the cached client without creating a new one.""" + mock_client = MagicMock(server_version="24.1") + mock_cc.get_client.return_value = mock_client + + client1 = create_clickhouse_client() + client2 = create_clickhouse_client() + + assert client1 is client2 + # get_client should only be called once + assert mock_cc.get_client.call_count == 1 + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context") + def test_different_config_creates_new_client(self, mock_get_context, mock_cc): + """Different session configs should produce different cached clients.""" + mock_client_a = MagicMock(server_version="24.1") + mock_client_b = MagicMock(server_version="24.1") + mock_cc.get_client.side_effect = [mock_client_a, mock_client_b] + + # First call: no overrides + mock_ctx = MagicMock() + mock_ctx.get_state.return_value = None + mock_get_context.return_value = mock_ctx + client1 = create_clickhouse_client() + + _clear_client_cache() + + # Second call: with override that changes the config key + mock_ctx2 = MagicMock() + mock_ctx2.get_state.return_value = {"connect_timeout": 99} + mock_get_context.return_value = mock_ctx2 + client2 = create_clickhouse_client() + + assert client1 is not client2 + assert mock_cc.get_client.call_count == 2 + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_stale_client_evicted_on_ping_failure(self, _mock_ctx, mock_cc): + """Client that fails ping after idle should be evicted and recreated.""" + mock_client_old = MagicMock(server_version="24.1") + mock_client_old.ping.return_value = False + mock_client_new = MagicMock(server_version="24.2") + mock_cc.get_client.side_effect = [mock_client_old, mock_client_new] + + client1 = create_clickhouse_client() + assert client1 is mock_client_old + + # Simulate idle time exceeding threshold + with _client_cache_lock: + for key, val in _client_cache.items(): + client, _ = val + _client_cache[key] = (client, time.time() - 120) + + client2 = create_clickhouse_client() + assert client2 is mock_client_new + assert mock_cc.get_client.call_count == 2 + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_autogenerate_session_id_disabled(self, _mock_ctx, mock_cc): + """Cached clients should be created with autogenerate_session_id=False.""" + mock_cc.get_client.return_value = MagicMock(server_version="24.1") + + create_clickhouse_client() + + call_kwargs = mock_cc.get_client.call_args[1] + assert call_kwargs["autogenerate_session_id"] is False + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_clear_cache_closes_clients(self, _mock_ctx, mock_cc): + """_clear_client_cache should close all cached clients.""" + mock_client = MagicMock(server_version="24.1") + mock_cc.get_client.return_value = mock_client + + create_clickhouse_client() + _clear_client_cache() + + mock_client.close.assert_called_once() + + +class TestResolveClientConfig: + """Tests for _resolve_client_config.""" + + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_send_receive_timeout_capped_when_not_explicit(self, _mock_ctx): + """send_receive_timeout should be capped to query_timeout + 5 by default.""" + config = _resolve_client_config() + + expected = get_mcp_config().query_timeout + 5 + assert config["send_receive_timeout"] == expected + + @patch.dict("os.environ", {"CLICKHOUSE_SEND_RECEIVE_TIMEOUT": "200"}) + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_send_receive_timeout_not_capped_when_explicit(self, _mock_ctx): + """Explicit env var should bypass the auto-cap.""" + config = _resolve_client_config() + assert config["send_receive_timeout"] == 200 + + @patch("mcp_clickhouse.mcp_server.get_context") + def test_session_override_timeout_not_capped(self, mock_get_context): + """Session override of send_receive_timeout should bypass the auto-cap.""" + mock_ctx = MagicMock() + mock_ctx.get_state.return_value = {"send_receive_timeout": 300} + mock_get_context.return_value = mock_ctx + + config = _resolve_client_config() + assert config["send_receive_timeout"] == 300 + + +class TestEvictionOnError: + """Tests for client eviction on connection errors.""" + + def setup_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + def teardown_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_execute_query_evicts_on_connection_error(self, _mock_ctx, mock_cc): + """execute_query should evict the cached client on connection errors.""" + mock_client = MagicMock(server_version="24.1") + mock_client.server_settings = {} + mock_client.query.side_effect = ConnectionError("connection reset") + mock_cc.get_client.return_value = mock_client + + config = _resolve_client_config() + + with pytest.raises(ToolError, match="connection reset"): + execute_query("SELECT 1", "evict-test", config) + + # Client should have been evicted — next call creates a new one + mock_client_new = MagicMock(server_version="24.2") + mock_client_new.server_settings = {} + mock_result = MagicMock() + mock_result.result_rows = [] + mock_result.column_names = [] + mock_client_new.query.return_value = mock_result + mock_cc.get_client.return_value = mock_client_new + + execute_query("SELECT 1", "evict-test-2", config) + assert mock_cc.get_client.call_count == 2 + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_execute_query_no_evict_on_sql_error(self, _mock_ctx, mock_cc): + """execute_query should NOT evict on normal SQL errors (not connection).""" + mock_client = MagicMock(server_version="24.1") + mock_client.server_settings = {} + mock_client.query.side_effect = Exception("Unknown column 'x'") + mock_cc.get_client.return_value = mock_client + + config = _resolve_client_config() + + with pytest.raises(ToolError): + execute_query("SELECT x", "no-evict-test", config) + + # Client should still be cached, second call reuses it + mock_client.query.side_effect = None + mock_result = MagicMock() + mock_result.result_rows = [] + mock_result.column_names = [] + mock_client.query.return_value = mock_result + execute_query("SELECT 1", "no-evict-test-2", config) + + # get_client only called once, reused from cache + assert mock_cc.get_client.call_count == 1 + + +class TestPingExceptionHandling: + """Tests for ping exception handling in create_clickhouse_client.""" + + def setup_method(self): + _clear_client_cache() + + def teardown_method(self): + _clear_client_cache() + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_ping_exception_evicts_and_recreates(self, _mock_ctx, mock_cc): + """A ping() that raises should evict the client and create a new one.""" + mock_client_old = MagicMock(server_version="24.1") + mock_client_old.ping.side_effect = Exception("boom") + mock_client_new = MagicMock(server_version="24.2") + mock_cc.get_client.side_effect = [mock_client_old, mock_client_new] + + client1 = create_clickhouse_client() + assert client1 is mock_client_old + + # Simulate idle time exceeding threshold + with _client_cache_lock: + for key, val in _client_cache.items(): + client, _ = val + _client_cache[key] = (client, time.time() - 120) + + client2 = create_clickhouse_client() + assert client2 is mock_client_new + assert mock_cc.get_client.call_count == 2 + + +class TestCacheRaceHandling: + """Tests for identity-checked cache updates around the idle-ping path.""" + + def setup_method(self): + _clear_client_cache() + + def teardown_method(self): + _clear_client_cache() + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_ping_ok_yields_to_newer_cached_client(self, _mock_ctx, mock_cc): + """If another thread replaced the cached client during ping, + a successful ping must not resurrect our stale candidate.""" + stale = MagicMock(server_version="24.1", name="stale") + replacement = MagicMock(server_version="24.2", name="replacement") + mock_cc.get_client.return_value = stale + + # Seed the cache with stale and backdate so create_clickhouse_client + # takes the idle-ping path on the next call. + create_clickhouse_client() + with _client_cache_lock: + (key,) = list(_client_cache.keys()) + _client_cache[key] = (stale, time.time() - 120) + + # While pinging, simulate another thread replacing the entry. + def ping_and_replace(): + with _client_cache_lock: + _client_cache[key] = (replacement, time.time()) + return True + + stale.ping.side_effect = ping_and_replace + + result = create_clickhouse_client() + + assert result is replacement + # Stale candidate must be closed; replacement must still be cached. + stale.close.assert_called_once() + with _client_cache_lock: + cached_client, _ = _client_cache[key] + assert cached_client is replacement + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_ping_fail_preserves_newer_cached_client(self, _mock_ctx, mock_cc): + """A failed ping must not drop a replacement installed by another thread.""" + stale = MagicMock(server_version="24.1", name="stale") + replacement = MagicMock(server_version="24.2", name="replacement") + # Second get_client call returns the newly created client in the + # fall-through path — it will get closed because the replacement wins. + fresh_create = MagicMock(server_version="24.3", name="fresh_create") + mock_cc.get_client.side_effect = [stale, fresh_create] + + create_clickhouse_client() + with _client_cache_lock: + (key,) = list(_client_cache.keys()) + _client_cache[key] = (stale, time.time() - 120) + + def ping_and_replace(): + with _client_cache_lock: + _client_cache[key] = (replacement, time.time()) + return False # ping fails + + stale.ping.side_effect = ping_and_replace + + result = create_clickhouse_client() + + # Replacement wins because it's still cached when the tail block runs. + assert result is replacement + stale.close.assert_called_once() + # The freshly created client was closed by the post-create race check. + fresh_create.close.assert_called_once() + with _client_cache_lock: + cached_client, _ = _client_cache[key] + assert cached_client is replacement + + +class TestShutdownOrdering: + """Tests that atexit shutdown closes the executor before the cache.""" + + @patch("mcp_clickhouse.mcp_server._clear_client_cache") + @patch("mcp_clickhouse.mcp_server.QUERY_EXECUTOR") + def test_executor_shutdown_runs_before_cache_clear( + self, mock_executor, mock_clear + ): + """The consolidated _shutdown callback must drain the executor first.""" + call_order = [] + mock_executor.shutdown.side_effect = lambda wait: call_order.append("executor") + mock_clear.side_effect = lambda: call_order.append("cache") + + _shutdown() + + assert call_order == ["executor", "cache"] + mock_executor.shutdown.assert_called_once_with(wait=True) + mock_clear.assert_called_once_with() diff --git a/tests/test_context_config_override.py b/tests/test_context_config_override.py index 789d131d..0248ae91 100644 --- a/tests/test_context_config_override.py +++ b/tests/test_context_config_override.py @@ -10,6 +10,7 @@ from mcp_clickhouse.mcp_server import ( mcp, create_clickhouse_client, + _clear_client_cache, CLIENT_CONFIG_OVERRIDES_KEY, ) @@ -29,6 +30,14 @@ async def on_call_tool(self, context: MiddlewareContext, call_next: CallNext): class TestConfigOverrideUnit: """Unit tests for the config override merge logic in create_clickhouse_client.""" + def setup_method(self): + """Clear the client cache before each test.""" + _clear_client_cache() + + def teardown_method(self): + """Clear the client cache after each test.""" + _clear_client_cache() + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") @patch("mcp_clickhouse.mcp_server.get_context") def test_overrides_merged_into_client_config(self, mock_get_context, mock_cc): diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index da4ba857..5e9903ed 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -383,7 +383,7 @@ async def test_concurrent_queries(mcp_server, setup_test_database): async def test_run_query_does_not_block_other_mcp_requests(mcp_server): """list_tools should complete while a query is in flight.""" - def slow_execute_query(_query: str): + def slow_execute_query(_query: str, _query_id: str, _client_config: dict): time.sleep(0.75) return json.dumps({"columns": ["value"], "rows": [[1]]}) diff --git a/tests/test_query_cancellation.py b/tests/test_query_cancellation.py new file mode 100644 index 00000000..f4b4483e --- /dev/null +++ b/tests/test_query_cancellation.py @@ -0,0 +1,212 @@ +"""Tests for query ID tracking and server-side cancellation.""" + +import concurrent.futures +import uuid +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp.exceptions import ToolError + +from mcp_clickhouse.mcp_server import ( + _active_queries, + _active_queries_lock, + _cancel_query, + _clear_client_cache, + _client_cache, + _client_cache_lock, + _resolve_client_config, + execute_query, + run_query, +) + + +class TestQueryIdTracking: + """Tests for query_id propagation through execute_query.""" + + def setup_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + def teardown_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_query_id_passed_in_settings(self, _mock_ctx, mock_cc): + """query_id should be included in the settings dict passed to client.query().""" + mock_client = MagicMock(server_version="24.1") + mock_client.server_settings = {} + mock_result = MagicMock() + mock_result.result_rows = [("row1",)] + mock_result.column_names = ["col1"] + mock_client.query.return_value = mock_result + mock_cc.get_client.return_value = mock_client + + config = _resolve_client_config() + execute_query("SELECT 1", "test-query-id-123", config) + + # Verify query_id was passed in settings + call_args = mock_client.query.call_args + settings = call_args[1].get("settings") or call_args.kwargs.get("settings") + assert settings["query_id"] == "test-query-id-123" + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_active_queries_tracked_and_cleaned(self, _mock_ctx, mock_cc): + """execute_query should register in _active_queries and clean up on completion.""" + mock_client = MagicMock(server_version="24.1") + mock_client.server_settings = {} + mock_result = MagicMock() + mock_result.result_rows = [] + mock_result.column_names = [] + mock_client.query.return_value = mock_result + mock_cc.get_client.return_value = mock_client + + config = _resolve_client_config() + + # Before execution + with _active_queries_lock: + assert "tracking-test-id" not in _active_queries + + execute_query("SELECT 1", "tracking-test-id", config) + + # After completion, should be cleaned up + with _active_queries_lock: + assert "tracking-test-id" not in _active_queries + + @patch("mcp_clickhouse.mcp_server.clickhouse_connect") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_active_queries_cleaned_on_error(self, _mock_ctx, mock_cc): + """execute_query should clean up _active_queries even on error.""" + mock_client = MagicMock(server_version="24.1") + mock_client.server_settings = {} + mock_client.query.side_effect = Exception("DB error") + mock_cc.get_client.return_value = mock_client + + config = _resolve_client_config() + + with pytest.raises(ToolError): + execute_query("SELECT bad", "error-test-id", config) + + with _active_queries_lock: + assert "error-test-id" not in _active_queries + + +class TestCancelQuery: + """Tests for _cancel_query server-side cancellation.""" + + def setup_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + def teardown_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + def test_cancel_issues_kill_query(self): + """_cancel_query should issue KILL QUERY via the cached client.""" + mock_client = MagicMock() + cache_key = (("host", "localhost"), ("port", 8443)) + query_id = str(uuid.uuid4()) + + # Set up cached client and active query + with _client_cache_lock: + _client_cache[cache_key] = (mock_client, 0) + with _active_queries_lock: + _active_queries[query_id] = (cache_key, "SELECT sleep(60)") + + _cancel_query(query_id) + + mock_client.command.assert_called_once_with( + f"KILL QUERY WHERE query_id = '{query_id}'" + ) + # Should be removed from active queries + with _active_queries_lock: + assert query_id not in _active_queries + + def test_cancel_noop_for_completed_query(self): + """_cancel_query should be a no-op if the query already completed.""" + # No entry in _active_queries + _cancel_query(str(uuid.uuid4())) # Should not raise + + def test_cancel_warns_without_cached_client(self): + """_cancel_query should log warning if no cached client is available.""" + cache_key = (("host", "gone"),) + query_id = str(uuid.uuid4()) + with _active_queries_lock: + _active_queries[query_id] = (cache_key, "SELECT 1") + + # No client in cache for this key + _cancel_query(query_id) # Should not raise + + with _active_queries_lock: + assert query_id not in _active_queries + + def test_cancel_failure_does_not_raise(self): + """_cancel_query should swallow exceptions from KILL QUERY.""" + mock_client = MagicMock() + mock_client.command.side_effect = Exception("Permission denied") + cache_key = (("host", "localhost"),) + query_id = str(uuid.uuid4()) + + with _client_cache_lock: + _client_cache[cache_key] = (mock_client, 0) + with _active_queries_lock: + _active_queries[query_id] = (cache_key, "SELECT 1") + + _cancel_query(query_id) # Should not raise + + def test_cancel_rejects_non_uuid_query_id(self): + """A non-UUID query_id must be refused before any KILL QUERY is issued.""" + mock_client = MagicMock() + cache_key = (("host", "localhost"),) + hostile = "foo'; DROP TABLE x; --" + + with _client_cache_lock: + _client_cache[cache_key] = (mock_client, 0) + with _active_queries_lock: + _active_queries[hostile] = (cache_key, "SELECT 1") + + _cancel_query(hostile) + + mock_client.command.assert_not_called() + # The active-query entry is always popped first, so it's gone either way. + with _active_queries_lock: + assert hostile not in _active_queries + + +class TestRunQueryTimeout: + """Tests for run_query timeout triggering _cancel_query.""" + + def setup_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + def teardown_method(self): + _clear_client_cache() + with _active_queries_lock: + _active_queries.clear() + + @patch("mcp_clickhouse.mcp_server._cancel_query") + @patch("mcp_clickhouse.mcp_server.QUERY_EXECUTOR") + @patch("mcp_clickhouse.mcp_server.get_context", side_effect=RuntimeError) + def test_timeout_triggers_cancel(self, _mock_ctx, mock_executor, mock_cancel): + """When run_query times out, it should call _cancel_query with the query_id.""" + mock_future = MagicMock() + mock_future.result.side_effect = concurrent.futures.TimeoutError() + mock_executor.submit.return_value = mock_future + + with pytest.raises(ToolError, match="timed out"): + run_query("SELECT sleep(999)") + + # _cancel_query should have been called with the generated query_id + mock_cancel.assert_called_once() + query_id = mock_cancel.call_args[0][0] + assert isinstance(query_id, str) + assert len(query_id) > 0