Fixed config override not playing nicely with pool executors

mcp_clickhouse/mcp_server.py

@@ -1,6 +1,7 @@
 import asyncio
 import atexit
 import concurrent.futures
+import copy
 import json
 import logging
 import os
@@ -478,8 +479,26 @@ def _validate_query_for_destructive_ops(query: str) -> None:
         )
 
 
-def execute_query(query: str) -> str:
-    client = create_clickhouse_client()
+def get_session_client_config_overrides() -> dict | None:
+    """Return request-scoped ClickHouse client config overrides, if any."""
+    try:
+        ctx = get_context()
+    except RuntimeError:
+        return None
+
+    session_config_overrides = ctx.get_state(CLIENT_CONFIG_OVERRIDES_KEY)
+    if session_config_overrides and not isinstance(session_config_overrides, dict):
+        logger.warning(
+            "%s must be a dict, got %s. Ignoring.",
+            CLIENT_CONFIG_OVERRIDES_KEY,
+            type(session_config_overrides).__name__,
+        )
+        return None
+    return copy.deepcopy(session_config_overrides) if session_config_overrides else None
+
+
+def execute_query(query: str, session_config_overrides: dict | None = None) -> str:
+    client = create_clickhouse_client(session_config_overrides)
     try:
         _validate_query_for_destructive_ops(query)
 
@@ -502,7 +521,8 @@ def run_query(query: str) -> str:
     """
     logger.info(f"Executing query: {query}")
     try:
-        future = QUERY_EXECUTOR.submit(execute_query, query)
+        session_config_overrides = get_session_client_config_overrides()
+        future = QUERY_EXECUTOR.submit(execute_query, query, session_config_overrides)
         timeout_secs = get_mcp_config().query_timeout
         try:
             return future.result(timeout=timeout_secs)
@@ -521,7 +541,8 @@ async def run_query_async(query: str) -> str:
     """Async MCP-facing wrapper for ClickHouse queries."""
     logger.info(f"Executing query: {query}")
     try:
-        future = QUERY_EXECUTOR.submit(execute_query, query)
+        session_config_overrides = get_session_client_config_overrides()
+        future = QUERY_EXECUTOR.submit(execute_query, query, session_config_overrides)
         timeout_secs = get_mcp_config().query_timeout
         try:
             return await asyncio.wait_for(
@@ -538,24 +559,17 @@ async def run_query_async(query: str) -> str:
         raise RuntimeError(f"Unexpected error during query execution: {str(e)}")
 
 
-def create_clickhouse_client():
+def create_clickhouse_client(session_config_overrides: dict | None = None):
     client_config = get_config().get_client_config()
 
-    try:
-        ctx = get_context()
-        session_config_overrides = ctx.get_state(CLIENT_CONFIG_OVERRIDES_KEY)
-        if session_config_overrides and not isinstance(session_config_overrides, dict):
-            logger.warning(
-                f"{CLIENT_CONFIG_OVERRIDES_KEY} must be a dict, got {type(session_config_overrides).__name__}. Ignoring."
-            )
-        elif session_config_overrides:
-            logger.debug(
-                f"Applying session-specific ClickHouse client config overrides: {list(session_config_overrides.keys())}"
-            )
-            client_config.update(session_config_overrides)
-    except RuntimeError:
-        # If we're outside a request context, just proceed with the default config
-        pass
+    if session_config_overrides is None:
+        session_config_overrides = get_session_client_config_overrides()
+    if session_config_overrides:
+        logger.debug(
+            "Applying session-specific ClickHouse client config overrides: %s",
+            list(session_config_overrides.keys()),
+        )
+        client_config.update(session_config_overrides)
 
     config_fields = [
         f"secure={client_config['secure']}",

tests/test_context_config_override.py

@@ -1,5 +1,8 @@
 """Tests for context state-based ClickHouse client configuration overrides."""
 
+import threading
+from types import SimpleNamespace
+
 import pytest
 from unittest.mock import patch, MagicMock
 
@@ -10,10 +13,43 @@
 from mcp_clickhouse.mcp_server import (
     mcp,
     create_clickhouse_client,
+    run_query,
     CLIENT_CONFIG_OVERRIDES_KEY,
 )
 
 
+def _base_client_config() -> dict:
+    return {
+        "host": "clickhouse.example.test",
+        "port": 8443,
+        "username": "huginn",
+        "password": "secret",
+        "interface": "https",
+        "secure": True,
+        "verify": True,
+        "connect_timeout": 30,
+        "send_receive_timeout": 300,
+        "client_name": "mcp_clickhouse",
+    }
+
+
+def _mock_clickhouse_config():
+    config = MagicMock()
+    config.get_client_config.return_value = _base_client_config()
+    config.allow_write_access = False
+    config.allow_drop = False
+    return config
+
+
+def _mock_clickhouse_query_client():
+    client = MagicMock(server_version="24.1", server_settings={})
+    result = MagicMock()
+    result.column_names = ["one"]
+    result.result_rows = [[1]]
+    client.query.return_value = result
+    return client
+
+
 class ConfigOverrideMiddleware(Middleware):
     """Test middleware that sets ClickHouse client config overrides."""
 
@@ -29,62 +65,125 @@ 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."""
 
-    @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):
+    def test_overrides_merged_into_client_config(self):
         """Verify overrides from context state are merged into the client config."""
         mock_ctx = MagicMock()
         mock_ctx.get_state.return_value = {"connect_timeout": 99, "send_receive_timeout": 199}
-        mock_get_context.return_value = mock_ctx
-        mock_cc.get_client.return_value = MagicMock(server_version="24.1")
 
-        create_clickhouse_client()
+        with (
+            patch("mcp_clickhouse.mcp_server.get_config", return_value=_mock_clickhouse_config()),
+            patch("mcp_clickhouse.mcp_server.clickhouse_connect") as mock_cc,
+            patch("mcp_clickhouse.mcp_server.get_context", return_value=mock_ctx),
+        ):
+            mock_cc.get_client.return_value = _mock_clickhouse_query_client()
+            create_clickhouse_client()
 
         call_kwargs = mock_cc.get_client.call_args[1]
         assert call_kwargs["connect_timeout"] == 99
         assert call_kwargs["send_receive_timeout"] == 199
 
-    @patch("mcp_clickhouse.mcp_server.clickhouse_connect")
-    @patch("mcp_clickhouse.mcp_server.get_context")
-    def test_empty_overrides_no_change(self, mock_get_context, mock_cc):
+    def test_empty_overrides_no_change(self):
         """Empty overrides dict should not alter the base config."""
         mock_ctx = MagicMock()
         mock_ctx.get_state.return_value = {}
-        mock_get_context.return_value = mock_ctx
-        mock_cc.get_client.return_value = MagicMock(server_version="24.1")
 
-        create_clickhouse_client()
+        with (
+            patch("mcp_clickhouse.mcp_server.get_config", return_value=_mock_clickhouse_config()),
+            patch("mcp_clickhouse.mcp_server.clickhouse_connect") as mock_cc,
+            patch("mcp_clickhouse.mcp_server.get_context", return_value=mock_ctx),
+        ):
+            mock_cc.get_client.return_value = _mock_clickhouse_query_client()
+            create_clickhouse_client()
 
         call_kwargs = mock_cc.get_client.call_args[1]
         # Base config values from env should pass through unchanged
         assert "host" in call_kwargs
         assert "username" in call_kwargs
 
-    @patch("mcp_clickhouse.mcp_server.clickhouse_connect")
-    @patch("mcp_clickhouse.mcp_server.get_context")
-    def test_no_overrides_in_context(self, mock_get_context, mock_cc):
+    def test_no_overrides_in_context(self):
         """When context state has no overrides, base config is used as-is."""
         mock_ctx = MagicMock()
         mock_ctx.get_state.return_value = None
-        mock_get_context.return_value = mock_ctx
-        mock_cc.get_client.return_value = MagicMock(server_version="24.1")
 
-        create_clickhouse_client()
+        with (
+            patch("mcp_clickhouse.mcp_server.get_config", return_value=_mock_clickhouse_config()),
+            patch("mcp_clickhouse.mcp_server.clickhouse_connect") as mock_cc,
+            patch("mcp_clickhouse.mcp_server.get_context", return_value=mock_ctx),
+        ):
+            mock_cc.get_client.return_value = _mock_clickhouse_query_client()
+            create_clickhouse_client()
 
         call_kwargs = mock_cc.get_client.call_args[1]
         assert "host" in call_kwargs
 
-    @patch("mcp_clickhouse.mcp_server.clickhouse_connect")
-    def test_no_request_context_falls_back_to_defaults(self, mock_cc):
+    def test_run_query_applies_role_override_across_executor_thread(self):
+        """run_query must preserve request overrides when work moves to QUERY_EXECUTOR."""
+        mock_ctx = MagicMock()
+        mock_ctx.get_state.return_value = {
+            "settings": {
+                "role": "merchant_role_123",
+            }
+        }
+        main_thread = threading.current_thread()
+
+        def get_context_only_in_request_thread():
+            if threading.current_thread() is main_thread:
+                return mock_ctx
+            raise RuntimeError("No FastMCP context in executor thread")
+
+        with (
+            patch("mcp_clickhouse.mcp_server.get_config", return_value=_mock_clickhouse_config()),
+            patch(
+                "mcp_clickhouse.mcp_server.get_mcp_config",
+                return_value=SimpleNamespace(query_timeout=5),
+            ),
+            patch("mcp_clickhouse.mcp_server.clickhouse_connect") as mock_cc,
+            patch(
+                "mcp_clickhouse.mcp_server.get_context",
+                side_effect=get_context_only_in_request_thread,
+            ),
+        ):
+            mock_cc.get_client.return_value = _mock_clickhouse_query_client()
+            run_query("SELECT merchant FROM analytics.bundle_product_sales")
+
+        call_kwargs = mock_cc.get_client.call_args.kwargs
+        assert call_kwargs["settings"] == {"role": "merchant_role_123"}
+
+    def test_invalid_override_type_ignored_with_warning(self, caplog):
+        """Invalid override values should not reach clickhouse_connect."""
+        mock_ctx = MagicMock()
+        mock_ctx.get_state.return_value = ["not", "a", "dict"]
+
+        with (
+            patch("mcp_clickhouse.mcp_server.get_config", return_value=_mock_clickhouse_config()),
+            patch("mcp_clickhouse.mcp_server.clickhouse_connect") as mock_cc,
+            patch("mcp_clickhouse.mcp_server.get_context", return_value=mock_ctx),
+            caplog.at_level("WARNING", logger="mcp-clickhouse"),
+        ):
+            mock_cc.get_client.return_value = _mock_clickhouse_query_client()
+            create_clickhouse_client()
+
+        call_kwargs = mock_cc.get_client.call_args.kwargs
+        assert "settings" not in call_kwargs
+        assert (
+            f"{CLIENT_CONFIG_OVERRIDES_KEY} must be a dict, got list. Ignoring."
+            in caplog.text
+        )
+
+    def test_no_request_context_falls_back_to_defaults(self):
         """Outside a request context (RuntimeError), base config is used."""
-        mock_cc.get_client.return_value = MagicMock(server_version="24.1")
-
-        # get_context is NOT mocked, so it will raise RuntimeError
-        # since there's no active FastMCP request context
-        create_clickhouse_client()
+        with (
+            patch("mcp_clickhouse.mcp_server.get_config", return_value=_mock_clickhouse_config()),
+            patch("mcp_clickhouse.mcp_server.clickhouse_connect") as mock_cc,
+        ):
+            mock_cc.get_client.return_value = _mock_clickhouse_query_client()
+            # get_context is NOT mocked, so it will raise RuntimeError
+            # since there's no active FastMCP request context
+            create_clickhouse_client()
 
         call_kwargs = mock_cc.get_client.call_args[1]
         assert "host" in call_kwargs
+        assert "settings" not in call_kwargs
 
 
 @pytest.fixture

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, _session_config_overrides=None):
         time.sleep(0.75)
         return json.dumps({"columns": ["value"], "rows": [[1]]})