mcp-hydrolix 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

mcp_hydrolix/mcp_env.py

@@ -4,10 +4,12 @@ This module handles all environment variable configuration with sensible default
 and type conversion.
 """
 
-from dataclasses import dataclass
 import os
-from typing import Optional
+from dataclasses import dataclass
 from enum import Enum
+from typing import Optional
+
+from .auth.credentials import HydrolixCredential, ServiceAccountToken, UsernamePassword
 
 
 class TransportType(str, Enum):
@@ -46,15 +48,49 @@ class HydrolixConfig:
         HYDROLIX_MCP_SERVER_TRANSPORT: MCP server transport method - "stdio", "http", or "sse" (default: stdio)
         HYDROLIX_MCP_BIND_HOST: Host to bind the MCP server to when using HTTP or SSE transport (default: 127.0.0.1)
         HYDROLIX_MCP_BIND_PORT: Port to bind the MCP server to when using HTTP or SSE transport (default: 8000)
+        HYDROLIX_QUERIES_POOL_SIZE 100
+        HYDROLIX_MCP_REQUEST_TIMEOUT 120
+        HYDROLIX_MCP_WORKERS 3
+        HYDROLIX_MCP_WORKER_CONNECTIONS 200
+        HYDROLIX_MCP_MAX_REQUESTS 10000
+        HYDROLIX_MCP_MAX_REQUESTS_JITTER 1000
+        HYDROLIX_MCP_MAX_KEEPALIVE 10
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the configuration from environment variables."""
         self._validate_required_vars()
+        # Credential to use for clickhouse connections when no per-request credential is provided
+        self._default_credential: Optional[HydrolixCredential] = None
+
+        # Set the default credential to the service account from the environment, if available
+        if (global_service_account := os.environ.get("HYDROLIX_TOKEN")) is not None:
+            self._default_credential = ServiceAccountToken(
+                global_service_account, f"https://{self.host}/config"
+            )
+        elif (global_username := os.environ.get("HYDROLIX_USER")) is not None and (
+            global_password := os.environ.get("HYDROLIX_PASSWORD")
+        ) is not None:
+            # No global service account available. Set the default credential to the username/password
+            # from the environment, if available
+            self._default_credential = UsernamePassword(global_username, global_password)
+
+    def creds_with(self, request_credential: Optional[HydrolixCredential]) -> HydrolixCredential:
+        if request_credential is not None:
+            return request_credential
+        elif self._default_credential is not None:
+            return self._default_credential
+        else:
+            raise ValueError(
+                "No credentials available for Hydrolix connection. "
+                "Please provide credentials either through HYDROLIX_TOKEN or "
+                "HYDROLIX_USER/HYDROLIX_PASSWORD environment variables, "
+                "or pass credentials explicitly via the creds parameter."
+            )
 
     @property
     def host(self) -> str:
-        """Get the Hydrolix host."""
+        """Get the Hydrolix host. Called during __init__"""
         return os.environ["HYDROLIX_HOST"]
 
     @property
@@ -68,40 +104,6 @@ class HydrolixConfig:
             return int(os.environ["HYDROLIX_PORT"])
         return 8088
 
-    @property
-    def service_account(self) -> bool:
-        """Determine if service account is enabled
-
-        Defaults to false.
-        Can be overridden if HYDROLIX_TOKEN environment variable.
-        """
-        return "HYDROLIX_TOKEN" in os.environ
-
-    @property
-    def service_account_token(self) -> str:
-        """Get the service account token
-
-        Defaults to None.
-        Can be overridden if HYDROLIX_TOKEN environment variable.
-        """
-        if "HYDROLIX_TOKEN" in os.environ:
-            return os.environ["HYDROLIX_TOKEN"]
-        return None
-
-    @property
-    def username(self) -> str:
-        """Get the Hydrolix username."""
-        if "HYDROLIX_USER" in os.environ:
-            return os.environ["HYDROLIX_USER"]
-        return None
-
-    @property
-    def password(self) -> str:
-        """Get the Hydrolix password."""
-        if "HYDROLIX_PASSWORD" in os.environ:
-            return os.environ["HYDROLIX_PASSWORD"]
-        return None
-
     @property
     def database(self) -> Optional[str]:
         """Get the default database name if set."""
@@ -115,6 +117,14 @@ class HydrolixConfig:
         """
         return os.getenv("HYDROLIX_VERIFY", "true").lower() == "true"
 
+    @property
+    def secure(self) -> bool:
+        """Get whether use secured connection.
+
+        Default: True
+        """
+        return os.getenv("HYDROLIX_SECURE", "true").lower() == "true"
+
     @property
     def connect_timeout(self) -> int:
         """Get the connection timeout in seconds.
@@ -132,7 +142,23 @@ class HydrolixConfig:
         return int(os.getenv("HYDROLIX_SEND_RECEIVE_TIMEOUT", "300"))
 
     @property
-    def proxy_path(self) -> str:
+    def query_pool_size(self) -> int:
+        """Get the send/receive timeout in seconds.
+
+        Default: 300 (Hydrolix default)
+        """
+        return int(os.getenv("HYDROLIX_QUERIES_POOL_SIZE", 100))
+
+    @property
+    def query_timeout_sec(self) -> int:
+        """Get the send/receive timeout in seconds.
+
+        Default: 300 (Hydrolix default)
+        """
+        return int(os.getenv("HYDROLIX_QUERY_TIMEOUT_SECS", 30))
+
+    @property
+    def proxy_path(self) -> Optional[str]:
         return os.getenv("HYDROLIX_PROXY_PATH")
 
     @property
@@ -168,19 +194,83 @@ class HydrolixConfig:
         """
         return int(os.getenv("HYDROLIX_MCP_BIND_PORT", "8000"))
 
-    def get_client_config(self) -> dict:
-        """Get the configuration dictionary for clickhouse_connect client.
+    @property
+    def mcp_timeout(self) -> int:
+        """Get the request timeout secunds.
+
+        Only used when transport is "http" or "sse".
+        Default: 120
+        """
+        return int(os.getenv("HYDROLIX_MCP_REQUEST_TIMEOUT", 120))
+
+    @property
+    def mcp_workers(self) -> int:
+        """Get the number of worker processes.
+
+        Only used when transport is "http" or "sse".
+        Default: 1
+        """
+        return int(os.getenv("HYDROLIX_MCP_WORKERS", 1))
+
+    @property
+    def mcp_worker_connections(self) -> int:
+        """Get the max number of concurrent requests per worker.
+
+        Only used when transport is "http" or "sse".
+        Default: 200
+        """
+        return int(os.getenv("HYDROLIX_MCP_WORKER_CONNECTIONS", 100))
+
+    @property
+    def mcp_max_requests_jitter(self) -> int:
+        """Get the random parameter to randomize time process is reloaded after max_requests.
+
+        Only used when transport is "http" or "sse".
+        Default: 10000
+        """
+        return int(os.getenv("HYDROLIX_MCP_MAX_REQUESTS_JITTER", 1000))
+
+    @property
+    def mcp_max_requests(self) -> int:
+        """Get the max number of requests handled by worker before it is restarted.
+
+        Only used when transport is "http" or "sse".
+        Default: 1000
+        """
+        return int(os.getenv("HYDROLIX_MCP_MAX_REQUESTS", 10000))
+
+    @property
+    def mcp_keepalive(self) -> int:
+        """Get a seconds of idle keepalive connections are kept alive.
+
+        Only used when transport is "http" or "sse".
+        Default: 10
+        """
+        return int(os.getenv("HYDROLIX_MCP_MAX_KEEPALIVE", 10))
+
+    def get_client_config(self, request_credential: Optional[HydrolixCredential]) -> dict:
+        """
+        Get the configuration dictionary for clickhouse_connect client.
+
+        Args:
+            request_credential: Optional credentials to use for this request. If not provided,
+                   falls back to the default credential for this HydrolixConfig
 
         Returns:
             dict: Configuration ready to be passed to clickhouse_connect.get_client()
+
+        Raises:
+            ValueError: If no credentials could be inferred for the request (either from
+                       the startup environment or provided in the request)
         """
         config = {
             "host": self.host,
             "port": self.port,
-            "secure": True,
+            "secure": self.secure,
             "verify": self.verify,
             "connect_timeout": self.connect_timeout,
             "send_receive_timeout": self.send_receive_timeout,
+            "executor_threads": self.query_pool_size,
             "client_name": "mcp_hydrolix",
         }
 
@@ -191,29 +281,29 @@ class HydrolixConfig:
         if self.proxy_path:
             config["proxy_path"] = self.proxy_path
 
-        if self.service_account:
-            config["access_token"] = self.service_account_token
-        else:
-            config["username"] = self.username
-            config["password"] = self.password
+        # Add credentials
+        config |= self.creds_with(request_credential).clickhouse_config_entries()
 
         return config
 
     def _validate_required_vars(self) -> None:
-        """Validate that all required environment variables are set.
+        """Validate that all required environment variables are set. Called during __init__.
 
         Raises:
             ValueError: If any required environment variable is missing.
         """
         missing_vars = []
-        if self.service_account:
-            required_vars = ["HYDROLIX_HOST", "HYDROLIX_TOKEN"]
-        else:
-            required_vars = ["HYDROLIX_HOST", "HYDROLIX_USER", "HYDROLIX_PASSWORD"]
+        required_vars = ["HYDROLIX_HOST"]
         for var in required_vars:
             if var not in os.environ:
                 missing_vars.append(var)
 
+        # HYDROLIX_USER and HYDROLIX_PASSWORD must either be both present or both absent
+        if ("HYDROLIX_USER" in os.environ) != ("HYDROLIX_PASSWORD" in os.environ):
+            raise ValueError(
+                "User/password authentication is only partially configured: set both HYDROLIX_USER and HYDROLIX_PASSWORD"
+            )
+
         if missing_vars:
             raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
 

mcp_hydrolix/mcp_server.py

@@ -1,19 +1,32 @@
-import logging
 import json
-from typing import Optional, List, Any
-import concurrent.futures
-import atexit
+import logging
+import signal
+from collections.abc import Sequence
+from dataclasses import asdict, is_dataclass
+from typing import Any, Final, Optional, List, cast, TypedDict
 
 import clickhouse_connect
+from clickhouse_connect import common
+from clickhouse_connect.driver import httputil
 from clickhouse_connect.driver.binding import format_query_value
 from dotenv import load_dotenv
 from fastmcp import FastMCP
 from fastmcp.exceptions import ToolError
-from dataclasses import dataclass, field, asdict, is_dataclass
+from fastmcp.server.dependencies import get_access_token
+from pydantic import Field
+from pydantic.dataclasses import dataclass
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse
 
-from mcp_hydrolix.mcp_env import get_config
+from .auth import (
+    AccessToken,
+    HydrolixCredential,
+    HydrolixCredentialChain,
+    ServiceAccountToken,
+    UsernamePassword,
+)
+from .mcp_env import HydrolixConfig, get_config
+from .utils import with_serializer
 
 
 @dataclass
@@ -33,35 +46,34 @@ class Table:
     name: str
     engine: str
     create_table_query: str
-    dependencies_database: str
-    dependencies_table: str
+    dependencies_database: List[str]
+    dependencies_table: List[str]
     engine_full: str
     sorting_key: str
     primary_key: str
-    total_rows: int
-    total_bytes: int
-    total_bytes_uncompressed: int
-    parts: int
-    active_parts: int
-    total_marks: int
+    total_rows: Optional[int]
+    total_bytes: Optional[int]
+    total_bytes_uncompressed: Optional[int]
+    parts: Optional[int]
+    active_parts: Optional[int]
+    total_marks: Optional[int]
+    columns: Optional[List[Column]] = Field([])
     comment: Optional[str] = None
-    columns: List[Column] = field(default_factory=list)
 
 
-MCP_SERVER_NAME = "mcp-hydrolix"
+@dataclass
+class HdxQueryResult(TypedDict):
+    columns: List[str]
+    rows: List[List[Any]]
 
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-)
-logger = logging.getLogger(MCP_SERVER_NAME)
 
-QUERY_EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=10)
-atexit.register(lambda: QUERY_EXECUTOR.shutdown(wait=True))
-SELECT_QUERY_TIMEOUT_SECS = 30
+MCP_SERVER_NAME = "mcp-hydrolix"
+logger = logging.getLogger(MCP_SERVER_NAME)
 
 load_dotenv()
 
+HYDROLIX_CONFIG: Final[HydrolixConfig] = get_config()
+
 mcp = FastMCP(
     name=MCP_SERVER_NAME,
     dependencies=[
@@ -69,9 +81,105 @@ mcp = FastMCP(
         "python-dotenv",
         "pip-system-certs",
     ],
+    auth=HydrolixCredentialChain(f"https://{HYDROLIX_CONFIG.host}/config"),
 )
 
 
+def get_request_credential() -> Optional[HydrolixCredential]:
+    if (token := get_access_token()) is not None:
+        if isinstance(token, AccessToken):
+            return token.as_credential()
+        else:
+            raise ValueError(
+                "Found non-hydrolix access token on request -- this should be impossible!"
+            )
+    return None
+
+
+async def create_hydrolix_client(pool_mgr, request_credential: Optional[HydrolixCredential]):
+    """
+    Create a client for operations against query-head. Note that this eagerly issues requests for initialization
+    of properties like `server_version`, and so may throw exceptions.
+    INV: clients returned by this method MUST NOT be reused across sessions, because they can close over per-session
+    credentials.
+    """
+    creds = HYDROLIX_CONFIG.creds_with(request_credential)
+    auth_info = (
+        f"as {creds.username}"
+        if isinstance(creds, UsernamePassword)
+        else f"using service account {cast(ServiceAccountToken, creds).service_account_id}"
+    )
+    logger.info(
+        f"Creating Hydrolix client connection to {HYDROLIX_CONFIG.host}:{HYDROLIX_CONFIG.port} "
+        f"{auth_info} "
+        f"(connect_timeout={HYDROLIX_CONFIG.connect_timeout}s, "
+        f"send_receive_timeout={HYDROLIX_CONFIG.send_receive_timeout}s)"
+    )
+
+    try:
+        client = await clickhouse_connect.get_async_client(
+            pool_mgr=pool_mgr, **HYDROLIX_CONFIG.get_client_config(request_credential)
+        )
+        # Test the connection
+        version = client.client.server_version
+        logger.info(f"Successfully connected to Hydrolix compatible with ClickHouse {version}")
+        return client
+    except Exception as e:
+        logger.error(f"Failed to connect to Hydrolix: {str(e)}")
+        raise
+
+
+# allow custom hydrolix settings in CH client
+common.set_setting("invalid_setting_action", "send")
+common.set_setting("autogenerate_session_id", False)
+client_shared_pool = httputil.get_pool_manager(maxsize=HYDROLIX_CONFIG.query_pool_size, num_pools=1)
+
+
+def term(*args, **kwargs):
+    client_shared_pool.clear()
+
+
+signal.signal(signal.SIGTERM, term)
+signal.signal(signal.SIGINT, term)
+signal.signal(signal.SIGQUIT, term)
+
+
+async def execute_query(query: str) -> HdxQueryResult:
+    try:
+        async with await create_hydrolix_client(
+            client_shared_pool, get_request_credential()
+        ) as client:
+            res = await client.query(
+                query,
+                settings={
+                    "readonly": 1,
+                    "hdx_query_max_execution_time": HYDROLIX_CONFIG.query_timeout_sec,
+                    "hdx_query_max_attempts": 1,
+                    "hdx_query_max_result_rows": 100_000,
+                    "hdx_query_max_memory_usage": 2 * 1024 * 1024 * 1024,  # 2GiB
+                    "hdx_query_admin_comment": f"User: {MCP_SERVER_NAME}",
+                },
+            )
+            logger.info(f"Query returned {len(res.result_rows)} rows")
+            return HdxQueryResult(columns=res.column_names, rows=res.result_rows)
+    except Exception as err:
+        logger.error(f"Error executing query: {err}")
+        raise ToolError(f"Query execution failed: {str(err)}")
+
+
+async def execute_cmd(query: str):
+    try:
+        async with await create_hydrolix_client(
+            client_shared_pool, get_request_credential()
+        ) as client:
+            res = await client.command(query)
+            logger.info("Command returned executed.")
+            return res
+    except Exception as err:
+        logger.error(f"Error executing command: {err}")
+        raise ToolError(f"Command execution failed: {str(err)}")
+
+
 @mcp.custom_route("/health", methods=["GET"])
 async def health_check(request: Request) -> PlainTextResponse:
     """Health check endpoint for monitoring server status.
@@ -80,8 +188,10 @@ async def health_check(request: Request) -> PlainTextResponse:
     """
     try:
         # Try to create a client connection to verify query-head connectivity
-        client = create_hydrolix_client()
-        version = client.server_version
+        async with await create_hydrolix_client(
+            client_shared_pool, get_request_credential()
+        ) as client:
+            version = client.client.server_version
         return PlainTextResponse(f"OK - Connected to Hydrolix compatible with ClickHouse {version}")
     except Exception as e:
         # Return 503 Service Unavailable if we can't connect to Hydrolix
@@ -97,21 +207,24 @@ def result_to_column(query_columns, result) -> List[Column]:
 
 
 def to_json(obj: Any) -> str:
+    # This function technically returns different types:
+    # - str for dataclasses (the primary use case)
+    # - list/dict/Any for recursive processing during serialization
+    # Type checking is suppressed for non-str returns as they're only used internally by json.dumps
     if is_dataclass(obj):
         return json.dumps(asdict(obj), default=to_json)
     elif isinstance(obj, list):
-        return [to_json(item) for item in obj]
+        return [to_json(item) for item in obj]  # type: ignore[return-value]
     elif isinstance(obj, dict):
-        return {key: to_json(value) for key, value in obj.items()}
-    return obj
+        return {key: to_json(value) for key, value in obj.items()}  # type: ignore[return-value]
+    return obj  # type: ignore[return-value]
 
 
 @mcp.tool()
-def list_databases():
+async def list_databases() -> List[str]:
     """List available Hydrolix databases"""
     logger.info("Listing all databases")
-    client = create_hydrolix_client()
-    result = client.command("SHOW DATABASES")
+    result = await execute_cmd("SHOW DATABASES")
 
     # Convert newline-separated string to list and trim whitespace
     if isinstance(result, str):
@@ -120,65 +233,53 @@ def list_databases():
         databases = [result]
 
     logger.info(f"Found {len(databases)} databases")
-    return json.dumps(databases)
+    return databases
 
 
 @mcp.tool()
-def list_tables(database: str, like: Optional[str] = None, not_like: Optional[str] = None):
+async def list_tables(
+    database: str, like: Optional[str] = None, not_like: Optional[str] = None
+) -> List[Table]:
     """List available Hydrolix tables in a database, including schema, comment,
     row count, and column count."""
     logger.info(f"Listing tables in database '{database}'")
-    client = create_hydrolix_client()
-    query = f"SELECT database, name, engine, create_table_query, dependencies_database, dependencies_table, engine_full, sorting_key, primary_key, total_rows, total_bytes, total_bytes_uncompressed, parts, active_parts, total_marks, comment FROM system.tables WHERE database = {format_query_value(database)}"
+    query = f"""
+        SELECT database, name, engine, create_table_query, dependencies_database,
+            dependencies_table, engine_full, sorting_key, primary_key, total_rows, total_bytes,
+            total_bytes_uncompressed, parts, active_parts, total_marks, comment
+        FROM system.tables WHERE database = {format_query_value(database)}"""
     if like:
         query += f" AND name LIKE {format_query_value(like)}"
 
     if not_like:
         query += f" AND name NOT LIKE {format_query_value(not_like)}"
 
-    result = client.query(query)
+    result = await execute_query(query)
 
     # Deserialize result as Table dataclass instances
-    tables = result_to_table(result.column_names, result.result_rows)
+    tables = result_to_table(result["columns"], result["rows"])
 
     for table in tables:
-        column_data_query = f"SELECT database, table, name, type AS column_type, default_kind, default_expression, comment FROM system.columns WHERE database = {format_query_value(database)} AND table = {format_query_value(table.name)}"
-        column_data_query_result = client.query(column_data_query)
+        column_data_query = f"""
+            SELECT database, table, name, type AS column_type, default_kind, default_expression, comment
+            FROM system.columns
+            WHERE database = {format_query_value(database)} AND table = {format_query_value(table.name)}"""
+        column_data_query_result = await execute_query(column_data_query)
         table.columns = [
             c
             for c in result_to_column(
-                column_data_query_result.column_names,
-                column_data_query_result.result_rows,
+                column_data_query_result["columns"],
+                column_data_query_result["rows"],
             )
         ]
 
     logger.info(f"Found {len(tables)} tables")
-    return [asdict(table) for table in tables]
-
-
-def execute_query(query: str):
-    client = create_hydrolix_client()
-    try:
-        res = client.query(
-            query,
-            settings={
-                "readonly": 1,
-                "hdx_query_max_execution_time": SELECT_QUERY_TIMEOUT_SECS,
-                "hdx_query_max_attempts": 1,
-                "hdx_query_max_result_rows": 100_000,
-                "hdx_query_max_memory_usage": 2 * 1024 * 1024 * 1024,  # 2GiB
-                "hdx_query_admin_comment": f"User: {MCP_SERVER_NAME}",
-            },
-        )
-        logger.info(f"Query returned {len(res.result_rows)} rows")
-        return {"columns": res.column_names, "rows": res.result_rows}
-    except Exception as err:
-        logger.error(f"Error executing query: {err}")
-        raise ToolError(f"Query execution failed: {str(err)}")
+    return tables
 
 
 @mcp.tool()
-def run_select_query(query: str):
+@with_serializer
+async def run_select_query(query: str) -> dict[str, tuple | Sequence[str | Sequence[Any]]]:
     """Run a SELECT query in a Hydrolix time-series database using the Clickhouse SQL dialect.
     Queries run using this tool will timeout after 30 seconds.
 
@@ -218,51 +319,8 @@ def run_select_query(query: str):
     """
     logger.info(f"Executing SELECT query: {query}")
     try:
-        future = QUERY_EXECUTOR.submit(execute_query, query)
-        try:
-            result = future.result(timeout=SELECT_QUERY_TIMEOUT_SECS)
-            # Check if we received an error structure from execute_query
-            if isinstance(result, dict) and "error" in result:
-                logger.warning(f"Query failed: {result['error']}")
-                # MCP requires structured responses; string error messages can cause
-                # serialization issues leading to BrokenResourceError
-                return {
-                    "status": "error",
-                    "message": f"Query failed: {result['error']}",
-                }
-            return result
-        except concurrent.futures.TimeoutError:
-            logger.warning(f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds: {query}")
-            future.cancel()
-            raise ToolError(f"Query timed out after {SELECT_QUERY_TIMEOUT_SECS} seconds")
-    except ToolError:
-        raise
+        result = await execute_query(query=query)
+        return result
     except Exception as e:
         logger.error(f"Unexpected error in run_select_query: {str(e)}")
-        raise RuntimeError(f"Unexpected error during query execution: {str(e)}")
-
-
-def create_hydrolix_client():
-    client_config = get_config().get_client_config()
-    auth_info = (
-        f"as {client_config['username']}"
-        if "username" in client_config
-        else "using service account token"
-    )
-    logger.info(
-        f"Creating Hydrolix client connection to {client_config['host']}:{client_config['port']} "
-        f"{auth_info} "
-        f"(secure={client_config['secure']}, verify={client_config['verify']}, "
-        f"connect_timeout={client_config['connect_timeout']}s, "
-        f"send_receive_timeout={client_config['send_receive_timeout']}s)"
-    )
-
-    try:
-        client = clickhouse_connect.get_client(**client_config)
-        # Test the connection
-        version = client.server_version
-        logger.info(f"Successfully connected to Hydrolix compatible with ClickHouse {version}")
-        return client
-    except Exception as e:
-        logger.error(f"Failed to connect to Hydrolix: {str(e)}")
-        raise
+        raise ToolError(f"Unexpected error during query execution: {str(e)}")