diff --git a/README.md b/README.md index cad481b..cb8991a 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,18 @@ An MCP server for ClickHouse. * Input: `sql` (string): The SQL query to execute. * Query data directly from various sources (files, URLs, databases) without ETL processes. +### Health Check Endpoint + +When running with HTTP or SSE transport, a health check endpoint is available at `/health`. This endpoint: +- Returns `200 OK` with the ClickHouse version if the server is healthy and can connect to ClickHouse +- Returns `503 Service Unavailable` if the server cannot connect to ClickHouse + +Example: +```bash +curl http://localhost:8000/health +# Response: OK - Connected to ClickHouse 24.3.1 +``` + ## Configuration This MCP server supports both ClickHouse and chDB. You can enable either or both depending on your needs. @@ -179,6 +191,18 @@ CLICKHOUSE_PASSWORD=clickhouse 4. For easy testing with the MCP Inspector, run `fastmcp dev mcp_clickhouse/mcp_server.py` to start the MCP server. +5. To test with HTTP transport and the health check endpoint: + ```bash + # Using default port 8000 + CLICKHOUSE_MCP_SERVER_TRANSPORT=http python -m mcp_clickhouse.main + + # Or with a custom port + CLICKHOUSE_MCP_SERVER_TRANSPORT=http CLICKHOUSE_MCP_BIND_PORT=4200 python -m mcp_clickhouse.main + + # Then in another terminal: + curl http://localhost:8000/health # or http://localhost:4200/health for custom port + ``` + ### Environment Variables The following environment variables are used to configure the ClickHouse and chDB connections: @@ -216,7 +240,14 @@ The following environment variables are used to configure the ClickHouse and chD * Set this to automatically connect to a specific database * `CLICKHOUSE_MCP_SERVER_TRANSPORT`: Sets the transport method for the MCP server. * Default: `"stdio"` - * Valid options: `"stdio"`, `"http"`, `"streamable-http"`, `"sse"`. This is useful for local development with tools like MCP Inspector. + * Valid options: `"stdio"`, `"http"`, `"sse"`. This is useful for local development with tools like MCP Inspector. +* `CLICKHOUSE_MCP_BIND_HOST`: Host to bind the MCP server to when using HTTP or SSE transport + * Default: `"127.0.0.1"` + * Set to `"0.0.0.0"` to bind to all network interfaces (useful for Docker or remote access) + * Only used when transport is `"http"` or `"sse"` +* `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_ENABLED`: Enable/disable ClickHouse functionality * Default: `"true"` * Set to `"false"` to disable ClickHouse tools when using chDB only @@ -227,7 +258,7 @@ The following environment variables are used to configure the ClickHouse and chD * Default: `"false"` * Set to `"true"` to enable chDB tools * `CHDB_DATA_PATH`: The path to the chDB data directory - * Default: `":memory:"` (in-memory database) + * Default: `":memory:"` (in-memory database) * Use `:memory:` for in-memory database * Use a file path for persistent storage (e.g., `/path/to/chdb/data`) @@ -286,6 +317,21 @@ CLICKHOUSE_ENABLED=false CHDB_DATA_PATH=/path/to/chdb/data ``` +For MCP Inspector or remote access with HTTP transport: + +```env +CLICKHOUSE_HOST=localhost +CLICKHOUSE_USER=default +CLICKHOUSE_PASSWORD=clickhouse +CLICKHOUSE_MCP_SERVER_TRANSPORT=http +CLICKHOUSE_MCP_BIND_HOST=0.0.0.0 # Bind to all interfaces +CLICKHOUSE_MCP_BIND_PORT=4200 # Custom port (default: 8000) +``` + +When using HTTP transport, the server will run on the configured port (default 8000). For example, with the above configuration: +- MCP endpoint: `http://localhost:4200/mcp` +- Health check: `http://localhost:4200/health` + You can set these variables in your environment, in a `.env` file, or in the Claude Desktop configuration: ```json @@ -305,13 +351,18 @@ You can set these variables in your environment, in a `.env` file, or in the Cla "CLICKHOUSE_HOST": "", "CLICKHOUSE_USER": "", "CLICKHOUSE_PASSWORD": "", - "CLICKHOUSE_DATABASE": "" + "CLICKHOUSE_DATABASE": "", + "CLICKHOUSE_MCP_SERVER_TRANSPORT": "stdio", + "CLICKHOUSE_MCP_BIND_HOST": "127.0.0.1", + "CLICKHOUSE_MCP_BIND_PORT": "8000" } } } } ``` +Note: The bind host and port settings are only used when transport is set to "http" or "sse". + ### Running tests ```bash diff --git a/mcp_clickhouse/main.py b/mcp_clickhouse/main.py index 002971e..97599a4 100644 --- a/mcp_clickhouse/main.py +++ b/mcp_clickhouse/main.py @@ -1,11 +1,20 @@ from .mcp_server import mcp -from .mcp_env import get_config +from .mcp_env import get_config, TransportType def main(): config = get_config() transport = config.mcp_server_transport - mcp.run(transport=transport) + + # For HTTP and SSE transports, we need to specify host and port + http_transports = [TransportType.HTTP.value, TransportType.SSE.value] + if transport in http_transports: + # Use the configured bind host (defaults to 127.0.0.1, can be set to 0.0.0.0) + # and bind port (defaults to 8000) + mcp.run(transport=transport, host=config.mcp_bind_host, port=config.mcp_bind_port) + else: + # For stdio transport, no host or port is needed + mcp.run(transport=transport) if __name__ == "__main__": diff --git a/mcp_clickhouse/mcp_env.py b/mcp_clickhouse/mcp_env.py index 6b271ac..c201a20 100644 --- a/mcp_clickhouse/mcp_env.py +++ b/mcp_clickhouse/mcp_env.py @@ -7,6 +7,20 @@ from dataclasses import dataclass import os from typing import Optional +from enum import Enum + + +class TransportType(str, Enum): + """Supported MCP server transport types.""" + + STDIO = "stdio" + HTTP = "http" + SSE = "sse" + + @classmethod + def values(cls) -> list[str]: + """Get all valid transport values.""" + return [transport.value for transport in cls] @dataclass @@ -30,6 +44,8 @@ class ClickHouseConfig: CLICKHOUSE_DATABASE: Default database to use (default: None) CLICKHOUSE_PROXY_PATH: Path to be added to the host URL. For instance, for servers behind an HTTP proxy (default: None) CLICKHOUSE_MCP_SERVER_TRANSPORT: MCP server transport method - "stdio", "http", or "sse" (default: stdio) + CLICKHOUSE_MCP_BIND_HOST: Host to bind the MCP server to when using HTTP or SSE transport (default: 127.0.0.1) + CLICKHOUSE_MCP_BIND_PORT: Port to bind the MCP server to when using HTTP or SSE transport (default: 8000) CLICKHOUSE_ENABLED: Enable ClickHouse server (default: true) """ @@ -120,14 +136,32 @@ def mcp_server_transport(self) -> str: Valid options: "stdio", "http", "sse" Default: "stdio" """ - transport = os.getenv("CLICKHOUSE_MCP_SERVER_TRANSPORT", "stdio").lower() - valid_transports = "stdio", "http", "streamable-http", "sse" - if transport not in valid_transports: - raise ValueError( - f"Invalid transport '{transport}'. Valid options: {', '.join(valid_transports)}" - ) + transport = os.getenv("CLICKHOUSE_MCP_SERVER_TRANSPORT", TransportType.STDIO.value).lower() + + # Validate transport type + if transport not in TransportType.values(): + valid_options = ", ".join(f'"{t}"' for t in TransportType.values()) + raise ValueError(f"Invalid transport '{transport}'. Valid options: {valid_options}") return transport + @property + def mcp_bind_host(self) -> str: + """Get the host to bind the MCP server to. + + Only used when transport is "http" or "sse". + Default: "127.0.0.1" + """ + return os.getenv("CLICKHOUSE_MCP_BIND_HOST", "127.0.0.1") + + @property + def mcp_bind_port(self) -> int: + """Get the port to bind the MCP server to. + + Only used when transport is "http" or "sse". + Default: 8000 + """ + return int(os.getenv("CLICKHOUSE_MCP_BIND_PORT", "8000")) + def get_client_config(self) -> dict: """Get the configuration dictionary for clickhouse_connect client. @@ -239,7 +273,7 @@ def get_chdb_config() -> ChDBConfig: """ Gets the singleton instance of ChDBConfig. Instantiates it on the first call. - + Returns: ChDBConfig: The chDB configuration instance """ diff --git a/mcp_clickhouse/mcp_server.py b/mcp_clickhouse/mcp_server.py index 6473685..e9b9a09 100644 --- a/mcp_clickhouse/mcp_server.py +++ b/mcp_clickhouse/mcp_server.py @@ -12,7 +12,10 @@ from fastmcp import FastMCP from fastmcp.tools import Tool from fastmcp.prompts import Prompt +from fastmcp.exceptions import ToolError from dataclasses import dataclass, field, asdict, is_dataclass +from starlette.requests import Request +from starlette.responses import PlainTextResponse from mcp_clickhouse.mcp_env import get_config, get_chdb_config from mcp_clickhouse.chdb_prompt import CHDB_PROMPT @@ -75,6 +78,22 @@ class Table: ) +@mcp.custom_route("/health", methods=["GET"]) +async def health_check(request: Request) -> PlainTextResponse: + """Health check endpoint for monitoring server status. + + Returns OK if the server is running and can connect to ClickHouse. + """ + try: + # Try to create a client connection to verify ClickHouse connectivity + client = create_clickhouse_client() + version = client.server_version + return PlainTextResponse(f"OK - Connected to ClickHouse {version}") + except Exception as e: + # Return 503 Service Unavailable if we can't connect to ClickHouse + return PlainTextResponse(f"ERROR - Cannot connect to ClickHouse: {str(e)}", status_code=503) + + def result_to_table(query_columns, result) -> List[Table]: return [Table(**dict(zip(query_columns, row))) for row in result] @@ -150,9 +169,7 @@ def execute_query(query: str): return {"columns": res.column_names, "rows": res.result_rows} except Exception as err: logger.error(f"Error executing query: {err}") - # Return a structured dictionary rather than a string to ensure proper serialization - # by the MCP protocol. String responses for errors can cause BrokenResourceError. - return {"error": str(err)} + raise ToolError(f"Query execution failed: {str(err)}") def run_select_query(query: str): @@ -175,16 +192,12 @@ def run_select_query(query: str): except concurrent.futures.TimeoutError: logger.warning(f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds: {query}") future.cancel() - # Return a properly structured response for timeout errors - return { - "status": "error", - "message": f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds", - } + raise ToolError(f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds") + except ToolError: + raise except Exception as e: logger.error(f"Unexpected error in run_select_query: {str(e)}") - # Catch all other exceptions and return them in a structured format - # to prevent MCP serialization failures - return {"status": "error", "message": f"Unexpected error: {str(e)}"} + raise RuntimeError(f"Unexpected error during query execution: {str(e)}") def create_clickhouse_client(): @@ -249,7 +262,7 @@ def execute_chdb_query(query: str): """Execute a query using chDB client.""" client = create_chdb_client() try: - res = client.query(query, 'JSON') + res = client.query(query, "JSON") if res.has_error(): error_msg = res.error_message() logger.error(f"Error executing chDB query: {error_msg}") @@ -310,7 +323,7 @@ def _init_chdb_client(): return None client_config = get_chdb_config().get_client_config() - data_path = client_config['data_path'] + data_path = client_config["data_path"] logger.info(f"Creating chDB client with data_path={data_path}") client = chs.Session(path=data_path) logger.info(f"Successfully connected to chDB with data_path={data_path}") @@ -337,7 +350,7 @@ def _init_chdb_client(): chdb_prompt = Prompt.from_function( chdb_initial_prompt, name="chdb_initial_prompt", - description="This prompt helps users understand how to interact and perform common operations in chDB" + description="This prompt helps users understand how to interact and perform common operations in chDB", ) mcp.add_prompt(chdb_prompt) logger.info("chDB tools and prompts registered") diff --git a/pyproject.toml b/pyproject.toml index df7e35d..f3fed54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mcp-clickhouse" -version = "0.1.9" +version = "0.1.10" description = "An MCP server for ClickHouse." readme = "README.md" license = "Apache-2.0" diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index f969207..0119790 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,6 +1,7 @@ import pytest import pytest_asyncio from fastmcp import Client +from fastmcp.exceptions import ToolError import asyncio from mcp_clickhouse.mcp_server import mcp, create_clickhouse_client from dotenv import load_dotenv @@ -256,15 +257,12 @@ async def test_run_select_query_error(mcp_server, setup_test_database): async with Client(mcp_server) as client: # Query non-existent table query = f"SELECT * FROM {test_db}.non_existent_table" - result = await client.call_tool("run_select_query", {"query": query}) - query_result = json.loads(result[0].text) + # Should raise ToolError + with pytest.raises(ToolError) as exc_info: + await client.call_tool("run_select_query", {"query": query}) - # Should return error structure - assert "status" in query_result - assert query_result["status"] == "error" - assert "message" in query_result - assert "Query failed" in query_result["message"] + assert "Query execution failed" in str(exc_info.value) @pytest.mark.asyncio @@ -273,13 +271,12 @@ async def test_run_select_query_syntax_error(mcp_server): async with Client(mcp_server) as client: # Invalid SQL syntax query = "SELECT FROM WHERE" - result = await client.call_tool("run_select_query", {"query": query}) - query_result = json.loads(result[0].text) + # Should raise ToolError + with pytest.raises(ToolError) as exc_info: + await client.call_tool("run_select_query", {"query": query}) - # Should return error structure - assert query_result["status"] == "error" - assert "Query failed" in query_result["message"] + assert "Query execution failed" in str(exc_info.value) @pytest.mark.asyncio diff --git a/tests/test_tool.py b/tests/test_tool.py index d038f03..50878c4 100644 --- a/tests/test_tool.py +++ b/tests/test_tool.py @@ -2,6 +2,7 @@ import json from dotenv import load_dotenv +from fastmcp.exceptions import ToolError from mcp_clickhouse import create_clickhouse_client, list_databases, list_tables, run_select_query @@ -73,10 +74,12 @@ def test_run_select_query_success(self): def test_run_select_query_failure(self): """Test running a SELECT query with an error.""" query = f"SELECT * FROM {self.test_db}.non_existent_table" - result = run_select_query(query) - self.assertIsInstance(result, dict) - self.assertEqual(result["status"], "error") - self.assertIn("Query failed", result["message"]) + + # Should raise ToolError + with self.assertRaises(ToolError) as context: + run_select_query(query) + + self.assertIn("Query execution failed", str(context.exception)) def test_table_and_column_comments(self): """Test that table and column comments are correctly retrieved.""" diff --git a/uv.lock b/uv.lock index 16c0a98..8222315 100644 --- a/uv.lock +++ b/uv.lock @@ -434,7 +434,7 @@ wheels = [ [[package]] name = "mcp-clickhouse" -version = "0.1.9" +version = "0.1.10" source = { editable = "." } dependencies = [ { name = "chdb" },