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

mcp_hydrolix/auth/__init__.py

@@ -0,0 +1,29 @@
+"""Authentication package for MCP Hydrolix.
+
+This package contains authentication-related types used to define hydrolix auth
+in terms of FastMCP infrastructure
+"""
+
+from .credentials import (
+    HydrolixCredential,
+    ServiceAccountToken,
+    UsernamePassword,
+)
+from .mcp_providers import (
+    TOKEN_PARAM,
+    AccessToken,
+    ChainedAuthBackend,
+    GetParamAuthBackend,
+    HydrolixCredentialChain,
+)
+
+__all__ = [
+    "HydrolixCredential",
+    "ServiceAccountToken",
+    "UsernamePassword",
+    "AccessToken",
+    "ChainedAuthBackend",
+    "GetParamAuthBackend",
+    "HydrolixCredentialChain",
+    "TOKEN_PARAM",
+]

mcp_hydrolix/auth/credentials.py

@@ -0,0 +1,63 @@
+"""Hydrolix credential types for authentication."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+import jwt
+
+
+class HydrolixCredential(ABC):
+    @abstractmethod
+    def clickhouse_config_entries(self) -> dict:
+        """
+        Returns the entries needed for a ClickHouse client config to use this credential.
+        This will typically add `access_token` or (`username` and `password`)
+        """
+        ...
+
+
+@dataclass
+class ServiceAccountToken(HydrolixCredential):
+    """Hydrolix credentials using a service account token."""
+
+    def __init__(self, token: str, expected_iss: Optional[str]):
+        """
+        Initialize a ServiceAccountToken from a token JWT (or raise an error if the claims are invalid).
+        NB the claims' signatures are NOT checked by this function -- these validations MUST NOT be considered
+        authoritative.
+        """
+
+        claims = jwt.decode(
+            token,
+            key="",  # NB service account signing key is not publicly-hosted, so we can't verify the signature
+            options={
+                "verify_signature": False,
+                "verify_iss": True,
+                "verify_iat": True,
+                "verify_exp": True,
+            },
+            issuer=expected_iss,
+        )
+        self.token = token
+        self.service_account_id = claims["sub"]
+        self.issued_at = claims["iss"]
+        self.expires_at = claims["exp"]
+
+    def clickhouse_config_entries(self) -> dict:
+        return {"access_token": self.token}
+
+    token: str
+    service_account_id: str
+    issued_at: int
+    expires_at: int
+
+
+@dataclass
+class UsernamePassword(HydrolixCredential):
+    """Hydrolix credentials using username and password."""
+
+    def clickhouse_config_entries(self) -> dict:
+        return {"username": self.username, "password": self.password}
+
+    username: str
+    password: str

mcp_hydrolix/auth/mcp_providers.py

@@ -0,0 +1,137 @@
+"""Authentication backends and providers for MCP Hydrolix server."""
+
+import time
+from abc import abstractmethod, ABC
+from typing import List, ClassVar, Final, Optional
+
+from fastmcp.server.auth import AccessToken as FastMCPAccessToken, AuthProvider
+from mcp.server.auth.middleware.auth_context import (
+    AuthContextMiddleware as McpAuthContextMiddleware,
+)
+from mcp.server.auth.middleware.bearer_auth import (
+    AuthenticatedUser as McpAuthenticatedUser,
+    BearerAuthBackend,
+)
+from mcp.server.auth.provider import TokenVerifier as McpTokenVerifier
+from starlette.authentication import AuthCredentials, AuthenticationBackend
+from starlette.middleware import Middleware
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.requests import HTTPConnection, Request
+
+from .credentials import HydrolixCredential, ServiceAccountToken
+
+TOKEN_PARAM: Final[str] = "token"
+
+
+class ChainedAuthBackend(AuthenticationBackend):
+    """
+    Generic authentication backend that tries multiple backends in order. Returns the first successful
+    authentication result. Only tries an auth method once all previous auth methods have failed.
+    """
+
+    def __init__(self, backends: List[AuthenticationBackend]):
+        self.backends = backends
+
+    async def authenticate(self, conn: HTTPConnection):
+        # due to a very strange quirk of python syntax, this CANNOT be an anonymous async generator. The quirk is
+        # that async generator expressions aren't allowed to have `await` in their if conditions (though async
+        # generators have no such restriction on their if statements)
+        async def successful_results():
+            for backend in self.backends:
+                if (result := await backend.authenticate(conn)) is not None:
+                    yield result
+
+        return await anext(successful_results(), None)
+
+
+class GetParamAuthBackend(AuthenticationBackend):
+    """
+    Authentication backend that validates tokens from a query parameter
+    """
+
+    def __init__(self, token_verifier: McpTokenVerifier, token_get_param: str):
+        self.token_verifier = token_verifier
+        self.token_get_param = token_get_param
+
+    async def authenticate(self, conn: HTTPConnection):
+        token = Request(conn.scope).query_params.get(self.token_get_param)
+
+        if token is None:
+            return None
+
+        # Validate the token with the verifier
+        auth_info = await self.token_verifier.verify_token(token)
+
+        if not auth_info:
+            return None
+
+        if auth_info.expires_at and auth_info.expires_at < int(time.time()):
+            return None
+
+        return AuthCredentials(auth_info.scopes), McpAuthenticatedUser(auth_info)
+
+
+class AccessToken(FastMCPAccessToken, ABC):
+    @abstractmethod
+    def as_credential(self) -> HydrolixCredential: ...
+
+
+class HydrolixCredentialChain(AuthProvider):
+    """
+    AuthProvider that authenticates with the following precedence (highest to lowest):
+
+    1. Per-request Bearer token: Service account token via Authorization: Bearer <token> header
+    2. Per-request GET parameter: Service account token via ?token=<token> query parameter
+
+    NB MCP-standard oAuth is not currently implemented
+    NB all per-request credentials take precedence over all environment-variable-supplied credentials
+    """
+
+    class ServiceAccountAccess(AccessToken):
+        FAKE_CLIENT_ID: ClassVar[Final[str]] = "MCP_CLIENT_VIA_SERVICE_ACCOUNT"
+        FAKE_SCOPE: ClassVar[Final[str]] = "MCP_SERVICE_ACCOUNT_SCOPE"
+
+        expected_issuer: Optional[str] = None
+
+        def as_credential(self) -> ServiceAccountToken:
+            return ServiceAccountToken(self.token, self.expected_issuer)
+
+    def __init__(self, expected_issuer: Optional[str]):
+        """
+        Initialize HydrolixCredentialChain.
+
+        Args:
+            expected_issuer: The issuer URL that must be used (mitigates credential-stuffing)
+        """
+        super().__init__()
+        self.expected_issuer = expected_issuer
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        """
+        This is responsible for validating and authenticating the `token`.
+        See ChainedAuthBackend for how the token is obtained in the first place.
+        Authorization is performed by individual endpoints via `fastmcp.server.dependencies.get_access_token`
+        """
+        return HydrolixCredentialChain.ServiceAccountAccess(
+            token=token,
+            client_id=HydrolixCredentialChain.ServiceAccountAccess.FAKE_CLIENT_ID,
+            scopes=[HydrolixCredentialChain.ServiceAccountAccess.FAKE_SCOPE],
+            expires_at=None,
+            resource=None,
+            claims={},
+            expected_issuer=self.expected_issuer,
+        )
+
+    def get_middleware(self) -> list:
+        return [
+            Middleware(
+                AuthenticationMiddleware,
+                backend=ChainedAuthBackend(
+                    [
+                        BearerAuthBackend(self),
+                        GetParamAuthBackend(self, TOKEN_PARAM),
+                    ]
+                ),
+            ),
+            Middleware(McpAuthContextMiddleware),
+        ]

mcp_hydrolix/log/__init__.py

@@ -0,0 +1,6 @@
+from .log import JsonFormatter, setup_logging
+
+__all__ = [
+    "setup_logging",
+    "JsonFormatter",
+]

mcp_hydrolix/log/log.py

@@ -0,0 +1,60 @@
+import json
+import logging
+import logging.config
+import os
+from pathlib import Path
+
+import yaml
+
+
+class JsonFormatter(logging.Formatter):
+    """
+    Custom formatter to output logs in JSON format.
+    """
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Convert log record to JSON format."""
+        log_record = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "component": "mcp-hydrolix",
+            "logger": record.name,
+        }
+
+        if isinstance(record.msg, dict):
+            log_record["message"] = json.dumps(record.msg)
+        else:
+            log_record["message"] = record.getMessage()
+
+        if record.exc_info:
+            log_record["exception"] = self.formatException(record.exc_info)
+
+        return json.dumps(log_record)
+
+
+def setup_logging(config_path: str | None, log_level: str, log_format: str) -> dict | None:
+    """
+    Configures logging from a YAML file and overrides level/format.
+    """
+    if config_path is None:
+        # print(f"Warning: Logging config file not provided at '{config_path}'. Using basic config.")
+        config_path = f"{os.path.dirname(__file__)}/log.yaml"
+
+    config_file = Path(config_path)
+    if not config_file.is_file():
+        logging.basicConfig(level=log_level.upper())
+        return None
+
+    with open(config_file) as f:
+        config = yaml.safe_load(f)
+
+    # Override level and formatter based on function arguments
+    config["root"]["level"] = log_level.upper()
+    if "loggers" in config and isinstance(config["loggers"], dict):
+        for logger in config["loggers"].keys():
+            config["loggers"][logger]["level"] = log_level.upper()
+
+    config["handlers"]["default"]["formatter"] = log_format
+
+    # logging.config.dictConfig(config)
+    return config

mcp_hydrolix/log/log.yaml

@@ -0,0 +1,40 @@
+version: 1
+disable_existing_loggers: false
+
+filters:
+    token_filter:
+      '()': mcp_hydrolix.log.utils.AccessLogTokenRedactingFilter
+
+# This section defines the format of your log messages
+formatters:
+  default:
+    format: "%(levelname)s:     %(name)s - %(message)s"
+  json:
+    # This special '()' key points to the custom JsonFormatter class in main.py
+    '()': mcp_hydrolix.log.JsonFormatter
+
+# This section defines where the logs are sent (e.g., console)
+handlers:
+  default:
+    # The formatter used by this handler will be set dynamically in the code
+    formatter: default
+    class: logging.StreamHandler
+    stream: ext://sys.stderr
+    filters:
+      - token_filter
+
+root:
+  level: INFO
+  handlers: [ default ]
+  propagate: false
+
+# This section defines the loggers for different parts of the application
+loggers:
+  uvicorn:
+    handlers: [ default ]
+    level: INFO
+    propagate: false
+  gunicorn:
+    handlers: [ default ]
+    level: INFO
+    propagate: false

mcp_hydrolix/log/utils.py

@@ -0,0 +1,56 @@
+"""Logging utilities for redacting sensitive information from logs."""
+
+import logging
+import re
+
+from ..auth import TOKEN_PARAM
+
+
+class AccessLogTokenRedactingFilter(logging.Filter):
+    """
+    Filter that redacts token query parameters from uvicorn access logs.
+
+    This filter is specifically designed to intercept log messages that contain
+    request URLs with query parameters and replace token values with [REDACTED].
+    """
+
+    # Regex pattern to match token=<value> in query strings
+    # Matches: token=<anything except & or whitespace>
+    TOKEN_PATTERN = re.compile(rf"{TOKEN_PARAM}=[^&\s]+")
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        """
+        Filter method that redacts tokens from the log message.
+
+        Args:
+            record: The log record to filter
+
+        Returns:
+            True (side-effect only)
+        """
+        if hasattr(record, "msg") and isinstance(record.msg, str):
+            record.msg = self.TOKEN_PATTERN.sub(rf"{TOKEN_PARAM}=[REDACTED]", record.msg)
+
+        # Also check args if they exist (for formatted log messages)
+        if hasattr(record, "args") and record.args:
+            # Convert args to list for modification
+            if isinstance(record.args, tuple):
+                modified_args: list = []
+                for arg in record.args:
+                    if isinstance(arg, str):
+                        # Redact tokens from string arguments
+                        modified_args.append(
+                            self.TOKEN_PATTERN.sub(rf"{TOKEN_PARAM}=[REDACTED]", arg)
+                        )
+                    elif isinstance(arg, bytes):
+                        # Redact tokens from string arguments
+                        modified_args.append(
+                            self.TOKEN_PATTERN.sub(
+                                rf"{TOKEN_PARAM}=[REDACTED]", arg.decode("utf-8")
+                            )
+                        )
+                    else:
+                        modified_args.append(arg)
+                record.args = tuple(modified_args)
+
+        return True

mcp_hydrolix/main.py

@@ -1,8 +1,76 @@
+import logging.config as lconfig
+
+from fastmcp.server.http import StarletteWithLifespan
+from gunicorn.app.base import BaseApplication
+
+from .log import setup_logging
+from .mcp_env import TransportType, get_config
 from .mcp_server import mcp
 
 
+class CoreApplication(BaseApplication):
+    """Gunicorn Core Application"""
+
+    def __init__(self, app: StarletteWithLifespan, options: dict = None) -> None:
+        """Initialize the core application."""
+        self.options = options or {}
+        self.app = app
+        super().__init__()
+
+    def load_config(self) -> None:
+        """Load the options specific to this application."""
+        config = {
+            key: value
+            for key, value in self.options.items()
+            if key in self.cfg.settings and value is not None
+        }
+        for key, value in config.items():
+            self.cfg.set(key.lower(), value)
+
+    def load(self) -> BaseApplication:
+        """Load the application."""
+        return self.app
+
+
 def main():
-    mcp.run()
+    config = get_config()
+    transport = config.mcp_server_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)
+        workers = config.mcp_workers
+        if workers == 1:
+            log_dict_config = setup_logging(None, "INFO", "json")
+            lconfig.dictConfig(log_dict_config)
+            mcp.run(
+                transport=transport,
+                host=config.mcp_bind_host,
+                port=config.mcp_bind_port,
+                uvicorn_config={"log_config": log_dict_config},
+            )
+        else:
+            log_dict_config = setup_logging(None, "INFO", "json")
+            lconfig.dictConfig(log_dict_config)
+            options = {
+                "bind": f"{config.mcp_bind_host}:{config.mcp_bind_port}",
+                "timeout": config.mcp_timeout,
+                "workers": config.mcp_workers,
+                "worker_class": "uvicorn.workers.UvicornWorker",
+                "worker_connections": config.mcp_worker_connections,
+                "max_requests": config.mcp_max_requests,
+                "max_requests_jitter": config.mcp_max_requests_jitter,
+                "keepalive": config.mcp_keepalive,
+                "logconfig_dict": log_dict_config,
+            }
+            CoreApplication(
+                mcp.http_app(path="/mcp", stateless_http=True, transport=transport), options
+            ).run()
+    else:
+        # For stdio transport, no host or port is needed
+        mcp.run(transport=transport)
 
 
 if __name__ == "__main__":