diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index 6758b020..c283f745 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/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']}", diff --git a/tests/test_context_config_override.py b/tests/test_context_config_override.py index 789d131d..836f71ac 100644 --- a/tests/test_context_config_override.py +++ b/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 diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index da4ba857..6ca65389 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, _session_config_overrides=None): time.sleep(0.75) return json.dumps({"columns": ["value"], "rows": [[1]]})