Merge pull request #11 from keboola/KAB-817-statefull-sessions

KAB-817: add the concept of "session state factory"

Author: jordanrburger

src/keboola_mcp_server/cli.py

@@ -1,7 +1,6 @@
 """Command-line interface for the Keboola MCP server."""
 
 import argparse
-import asyncio
 import logging
 import sys
 from typing import List, Optional
@@ -24,7 +23,7 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Keboola MCP Server")
     parser.add_argument(
         "--transport",
-        choices=["stdio"],
+        choices=["stdio", "sse"],
         default="stdio",
         help="Transport to use for MCP communication",
     )
@@ -35,7 +34,7 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
         help="Logging level",
     )
     parser.add_argument(
-        "--api-url", help="Keboola Storage API URL (defaults to https://connection.keboola.com)"
+        "--api-url", default="https://connection.keboola.com", help="Keboola Storage API URL"
     )
 
     return parser.parse_args(args)
@@ -49,11 +48,10 @@ def main(args: Optional[List[str]] = None) -> None:
     """
     parsed_args = parse_args(args)
 
-    # Create config from environment, but override with CLI args
-    config = Config.from_env()
-    if parsed_args.api_url:
-        config.storage_api_url = parsed_args.api_url
-    config.log_level = parsed_args.log_level
+    # Create config from the CLI arguments
+    config = Config.from_dict(
+        {"storage_api_url": parsed_args.api_url, "log_level": parsed_args.log_level}
+    )
 
     try:
         # Create and run server

src/keboola_mcp_server/config.py

@@ -1,18 +1,18 @@
 """Configuration handling for the Keboola MCP server."""
 
+import dataclasses
 import logging
-import os
 from dataclasses import dataclass
-from typing import Optional
+from typing import Mapping, Optional
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
+@dataclass(frozen=True)
 class Config:
     """Server configuration."""
 
-    storage_token: str
+    storage_token: Optional[str] = None
     storage_api_url: str = "https://connection.keboola.com"
     log_level: str = "INFO"
     # Add Snowflake credentials
@@ -24,71 +24,42 @@ class Config:
     snowflake_schema: Optional[str] = None
     snowflake_role: Optional[str] = None
 
-    def __init__(
-        self,
-        storage_token: str,
-        storage_api_url: str = "https://connection.keboola.com",
-        snowflake_account: Optional[str] = None,
-        snowflake_user: Optional[str] = None,
-        snowflake_password: Optional[str] = None,
-        snowflake_warehouse: Optional[str] = None,
-        snowflake_database: Optional[str] = None,
-        snowflake_role: Optional[str] = None,
-        snowflake_schema: Optional[str] = None,
-        log_level: str = "INFO",
-    ):
-        self.storage_token = storage_token
-        self.storage_api_url = storage_api_url
-        self.snowflake_account = snowflake_account
-        self.snowflake_user = snowflake_user
-        self.snowflake_password = snowflake_password
-        self.snowflake_warehouse = snowflake_warehouse
-        self.snowflake_database = snowflake_database
-        self.snowflake_role = snowflake_role
-        self.snowflake_schema = snowflake_schema
-        self.log_level = log_level
+    @classmethod
+    def _read_options(cls, d: Mapping[str, str]) -> Mapping[str, str]:
+        options: dict[str, str] = {}
+        for f in dataclasses.fields(cls):
+            if f.name in d:
+                options[f.name] = d.get(f.name)
+            elif (dict_name := f"KBC_{f.name.upper()}") in d:
+                options[f.name] = d.get(dict_name)
+        return options
 
     @classmethod
-    def from_env(cls) -> "Config":
-        """Create config from environment variables."""
-        # Add debug logging using logger instead of print
-        for env_var in [
-            "KBC_SNOWFLAKE_ACCOUNT",
-            "KBC_SNOWFLAKE_USER",
-            "KBC_SNOWFLAKE_PASSWORD",
-            "KBC_SNOWFLAKE_WAREHOUSE",
-            "KBC_SNOWFLAKE_DATABASE",
-            "KBC_SNOWFLAKE_ROLE",
-            "KBC_SNOWFLAKE_SCHEMA",
-        ]:
-            logger.debug(f"Reading {env_var}: {'set' if os.getenv(env_var) else 'not set'}")
+    def from_dict(cls, d: Mapping[str, str]) -> "Config":
+        """
+        Creates new `Config` instance with values read from the input mapping.
+        The keys in the input mapping can either be the names of the fields in `Config` class
+        or their uppercase variant prefixed with 'KBC_'.
+        """
+        return cls(**cls._read_options(d))
 
-        storage_token = os.getenv("KBC_STORAGE_TOKEN")
-        if not storage_token:
-            raise ValueError("KBC_STORAGE_TOKEN environment variable is required")
+    def replace_by(self, d: Mapping[str, str]) -> "Config":
+        """
+        Creates new `Config` instance from the existing one by replacing the values from the input mapping.
+        The keys in the input mapping can either be the names of the fields in `Config` class
+        or their uppercase variant prefixed with 'KBC_'.
+        """
+        return dataclasses.replace(self, **self._read_options(d))
 
-        return cls(
-            storage_token=storage_token,
-            storage_api_url=os.getenv("KBC_STORAGE_API_URL", "https://connection.keboola.com"),
-            snowflake_account=os.getenv("KBC_SNOWFLAKE_ACCOUNT"),
-            snowflake_user=os.getenv("KBC_SNOWFLAKE_USER"),
-            snowflake_password=os.getenv("KBC_SNOWFLAKE_PASSWORD"),
-            snowflake_warehouse=os.getenv("KBC_SNOWFLAKE_WAREHOUSE"),
-            snowflake_database=os.getenv("KBC_SNOWFLAKE_DATABASE"),
-            snowflake_role=os.getenv("KBC_SNOWFLAKE_ROLE"),
-            snowflake_schema=os.getenv("KBC_SNOWFLAKE_SCHEMA"),
-            log_level=os.getenv("KBC_LOG_LEVEL", "INFO"),
+    def has_storage_config(self) -> bool:
+        """Check if Storage API configuration is complete."""
+        return all(
+            [
+                self.storage_token,
+                self.storage_api_url,
+            ]
         )
 
-    def validate(self) -> None:
-        """Validate the configuration."""
-        if not self.storage_token:
-            raise ValueError("Storage token not configured")
-        if not self.storage_api_url:
-            raise ValueError("Storage API URL is required")
-        if self.log_level not in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
-            raise ValueError(f"Invalid log level: {self.log_level}")
-
     def has_snowflake_config(self) -> bool:
         """Check if Snowflake configuration is complete."""
         return all(
@@ -100,3 +71,16 @@ def has_snowflake_config(self) -> bool:
                 self.snowflake_database,
             ]
         )
+
+    def __repr__(self):
+        params: list[str] = []
+        for f in dataclasses.fields(self):
+            value = getattr(self, f.name)
+            if value:
+                if "token" in f.name or "password" in f.name:
+                    params.append(f"{f.name}='****'")
+                else:
+                    params.append(f"{f.name}='{value}'")
+            else:
+                params.append(f"{f.name}=None")
+        return f'Config({", ".join(params)})'

src/keboola_mcp_server/database.py

@@ -2,12 +2,10 @@
 
 import logging
 from contextlib import contextmanager
-
-import snowflake.connector
 from dataclasses import dataclass
-from typing import Optional, List, Tuple, Dict, Any
+from typing import Any, Dict, List
 
-from .config import Config
+import snowflake.connector
 
 logger = logging.getLogger(__name__)
 
@@ -179,38 +177,25 @@ def find_working_connection(self) -> str:
             error_msg += f" - {pattern} ({db}): {error}\n"
         raise DatabaseConnectionError(error_msg)
 
+    def create_snowflake_connection(self) -> snowflake.connector.connection:
+        """Create and return a Snowflake connection using configured credentials.
+        Raises:
+            ValueError: If credentials are not fully configured or connection fails
+        """
+        try:
+            database = self.find_working_connection()
+            conn = snowflake.connector.connect(
+                account=self.config.snowflake_account,
+                user=self.config.snowflake_user,
+                password=self.config.snowflake_password,
+                warehouse=self.config.snowflake_warehouse,
+                database=database,
+                schema=self.config.snowflake_schema,
+                role=self.config.snowflake_role,
+            )
 
-def create_snowflake_connection(config: Config) -> snowflake.connector.connection:
-    """Create a return a Snowflake connection using configured credentials.
-
-    Args:
-        config: Configuration object containing Snowflake credentials
-
-    Returns:
-        snowflake.connector.connection: established Snowflake connection
-
-    Raises:
-        ValueError: If credentials are not fully configured or connection fails
-    """
-    if not config.has_snowflake_config():
-        raise ValueError("Snowflake credentials are not fully configured")
-
-    try:
-        connection_manager = ConnectionManager(config)
-        database = connection_manager.find_working_connection()
-
-        conn = snowflake.connector.connect(
-            account=config.snowflake_account,
-            user=config.snowflake_user,
-            password=config.snowflake_password,
-            warehouse=config.snowflake_warehouse,
-            database=database,
-            schema=config.snowflake_schema,
-            role=config.snowflake_role,
-        )
-
-        return conn
-    except DatabaseConnectionError as e:
-        raise ValueError(f"Failed to find working database connection: {str(e)}")
-    except Exception as e:
-        raise ValueError(f"Failed to create Snowflake connection: {str(e)}")
+            return conn
+        except DatabaseConnectionError as e:
+            raise ValueError(f"Failed to find working database connection: {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Failed to create Snowflake connection: {str(e)}")