Support dot env file (#72)

Signed-off-by: Manabu McCloskey <[email protected]>

Author: nabuskey

.env.example

@@ -1,15 +1,15 @@
 # Spark History Server MCP Configuration
 
 # MCP Server Settings
-MCP_PORT=18888
-MCP_DEBUG=false
+SHS_MCP_PORT=18888 # Port for MCP server (default: 18888)
+SHS_MCP_DEBUG=true # Enable debug mode (default: false)
+SHS_MCP_ADDRESS=0.0.0.0 # Address for MCP server (default: localhost)
+SHS_MCP_TRANSPORT=treamable-http
 
-# Spark Authentication (Optional)
-# SPARK_USERNAME=your_spark_username
-# SPARK_PASSWORD=your_spark_password
-# SPARK_TOKEN=your_spark_token
-
-# Example for production:
-# SPARK_USERNAME=prod_user
-# SPARK_PASSWORD=secure_password_here
-# SPARK_TOKEN=jwt_token_here
+# Spark History Server Settings
+# SHS_SERVERS_*_URL - URL for a specific server
+# SHS_SERVERS_*_AUTH_USERNAME - Username for a specific server
+# SHS_SERVERS_*_AUTH_PASSWORD - Password for a specific server
+# SHS_SERVERS_*_AUTH_TOKEN - Token for a specific server
+# SHS_SERVERS_*_VERIFY_SSL - Whether to verify SSL for a specific server (true/false)
+# SHS_SERVERS_*_EMR_CLUSTER_ARN - EMR cluster ARN for a specific server

README.md

@@ -262,13 +262,17 @@ servers:
 ```
 
 ### 🔐 Environment Variables
-```bash
-SHS_SPARK_USERNAME=your_username
-SHS_SPARK_PASSWORD=your_password
-SHS_SPARK_TOKEN=your_jwt_token
-SHS_MCP_PORT=18888
-SHS_MCP_DEBUG=false
-SHS_MCP_ADDRESS=0.0.0.0
+```
+SHS_MCP_PORT - Port for MCP server (default: 18888)
+SHS_MCP_DEBUG - Enable debug mode (default: false)
+SHS_MCP_ADDRESS - Address for MCP server (default: localhost)
+SHS_MCP_TRANSPORT - MCP transport mode (default: streamable-http)
+SHS_SERVERS_*_URL - URL for a specific server
+SHS_SERVERS_*_AUTH_USERNAME - Username for a specific server
+SHS_SERVERS_*_AUTH_PASSWORD - Password for a specific server
+SHS_SERVERS_*_AUTH_TOKEN - Token for a specific server
+SHS_SERVERS_*_VERIFY_SSL - Whether to verify SSL for a specific server (true/false)
+SHS_SERVERS_*_EMR_CLUSTER_ARN - EMR cluster ARN for a specific server
 ```
 
 ## 🤖 AI Agent Integration

config.yaml

@@ -5,8 +5,8 @@ servers:
     url: "http://localhost:18080"
     # Optional authentication (can also use environment variables).
     # auth:
-    #   username: ${SHS_SPARK_USERNAME}
-    #   password: ${SHS_SPARK_PASSWORD}
+    #   username: ${SHS_SERVERS_LOCAL_AUTH_USERNAME}
+    #   password: ${SHS_SERVERS_LOCAL_AUTH_PASSWORD}
     #   token: ${SHS_SPARK_TOKEN}
 
   # Production server example
@@ -15,17 +15,17 @@ servers:
   #   verify_ssl: true
   #   auth:
       # Use environment variables for production
-      # username: ${SHS_SPARK_USERNAME}
-      # password: ${SHS_SPARK_PASSWORD}
+      # username: ${SHS_SERVERS_PRODUCTION_AUTH_USERNAME}
+      # password: ${SHS_SERVERS_PRODUCTION_AUTH_PASSWORD}
       # token: ${SHS_SPARK_TOKEN}
 
   # Staging server example
   # staging:
   #   url: "https://staging-spark-history.company.com:18080"
   #   verify_ssl: true
   #   auth:
-      # username: ${SHS_SPARK_USERNAME}
-      # token: ${SHS_SPARK_TOKEN}
+      # username: ${SHS_SERVERS_STAGING_AUTH_USERNAME}
+      # token: ${SHS_SERVERS_STAGING_AUTH_PASSWORD}
 
   # AWS Glue Spark History Server example
   # glue_ec2:
@@ -43,9 +43,15 @@ mcp:
   debug: true
   address: localhost
 
-# Environment Variables:
-# SHS_SPARK_USERNAME - Default username for authentication
-# SHS_SPARK_PASSWORD - Default password for authentication
-# SHS_SPARK_TOKEN - Default token for authentication
+
+# Available Environment Variables:
 # SHS_MCP_PORT - Port for MCP server (default: 18888)
 # SHS_MCP_DEBUG - Enable debug mode (default: false)
+# SHS_MCP_ADDRESS - Address for MCP server (default: localhost)
+# SHS_MCP_TRANSPORT - MCP transport mode (default: streamable-http)
+# SHS_SERVERS_*_URL - URL for a specific server
+# SHS_SERVERS_*_AUTH_USERNAME - Username for a specific server
+# SHS_SERVERS_*_AUTH_PASSWORD - Password for a specific server
+# SHS_SERVERS_*_AUTH_TOKEN - Token for a specific server
+# SHS_SERVERS_*_VERIFY_SSL - Whether to verify SSL for a specific server (true/false)
+# SHS_SERVERS_*_EMR_CLUSTER_ARN - EMR cluster ARN for a specific server

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "requests~=2.32.4",
     "pydantic~=2.4",
     "boto3~=1.34",
+    "pydantic-settings>=2.9.1",
 ]
 
 [build-system]

src/spark_history_mcp/config/config.py

@@ -2,53 +2,55 @@
 from typing import Dict, List, Literal, Optional
 
 import yaml
-from pydantic import BaseModel, Field
+from pydantic import Field
+from pydantic_settings import (
+    BaseSettings,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+)
 
 
-class AuthConfig(BaseModel):
+class AuthConfig(BaseSettings):
     """Authentication configuration for the Spark server."""
 
-    username: str = Field(None, alias="username")
-    password: str = Field(None, alias="password")
-    token: str = Field(None, alias="token")
+    username: Optional[str] = Field(None)
+    password: Optional[str] = Field(None)
+    token: Optional[str] = Field(None)
 
-    def __init__(self, **data):
-        # Support environment variables for sensitive data
-        if not data.get("username"):
-            data["username"] = os.getenv("SHS_SPARK_USERNAME")
-        if not data.get("password"):
-            data["password"] = os.getenv("SHS_SPARK_PASSWORD")
-        if not data.get("token"):
-            data["token"] = os.getenv("SHS_SPARK_TOKEN")
-        super().__init__(**data)
 
-
-class ServerConfig(BaseModel):
+class ServerConfig(BaseSettings):
     """Server configuration for the Spark server."""
 
     url: Optional[str] = None
-    auth: AuthConfig = Field(None, alias="auth")
-    default: bool = Field(False, alias="default")
-    verify_ssl: bool = Field(True, alias="verify_ssl")
+    auth: AuthConfig = Field(default_factory=AuthConfig, exclude=True)
+    default: bool = False
+    verify_ssl: bool = True
     emr_cluster_arn: Optional[str] = None  # EMR specific field
 
 
-class McpConfig(BaseModel):
+class McpConfig(BaseSettings):
     """Configuration for the MCP server."""
 
     transports: List[Literal["stdio", "sse", "streamable-http"]] = Field(
         default_factory=list
     )
-    address: str = Field(default="localhost")
-    port: str = Field(default="18888")
-    debug: bool = Field(default=False)
+    address: Optional[str] = "localhost"
+    port: Optional[int | str] = "18888"
+    debug: Optional[bool] = False
+    model_config = SettingsConfigDict(extra="ignore")
 
 
-class Config(BaseModel):
+class Config(BaseSettings):
     """Configuration for the Spark client."""
 
     servers: Dict[str, ServerConfig]
     mcp: Optional[McpConfig] = None
+    model_config = SettingsConfigDict(
+        env_prefix="SHS_",
+        env_nested_delimiter="_",
+        env_file=".env",
+        env_file_encoding="utf-8",
+    )
 
     @classmethod
     def from_file(cls, file_path: str) -> "Config":
@@ -60,3 +62,14 @@ def from_file(cls, file_path: str) -> "Config":
             config_data = yaml.safe_load(f)
 
         return cls.model_validate(config_data)
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return env_settings, dotenv_settings, init_settings, file_secret_settings

src/spark_history_mcp/core/app.py

@@ -65,12 +65,10 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
     yield AppContext(clients=clients, default_client=default_client)
 
 
-def run():
-    config = Config.from_file("config.yaml")
-
-    mcp.settings.host = os.getenv("SHS_MCP_ADDRESS", config.mcp.address)
-    mcp.settings.port = int(os.getenv("SHS_MCP_PORT", config.mcp.port))
-    mcp.settings.debug = bool(os.getenv("SHS_MCP_DEBUG", config.mcp.debug))
+def run(config: Config):
+    mcp.settings.host = config.mcp.address
+    mcp.settings.port = int(config.mcp.port)
+    mcp.settings.debug = bool(config.mcp.debug)
     mcp.run(transport=os.getenv("SHS_MCP_TRANSPORT", config.mcp.transports[0]))
 
 

src/spark_history_mcp/core/main.py

@@ -1,8 +1,10 @@
 """Main entry point for Spark History Server MCP."""
 
+import json
 import logging
 import sys
 
+from spark_history_mcp.config.config import Config
 from spark_history_mcp.core import app
 
 # Configure logging
@@ -16,7 +18,11 @@ def main():
     """Main entry point."""
     try:
         logger.info("Starting Spark History Server MCP...")
-        app.run()
+        config = Config.from_file("config.yaml")
+        if config.mcp.debug:
+            logger.setLevel(logging.DEBUG)
+        logger.debug(json.dumps(json.loads(config.model_dump_json()), indent=4))
+        app.run(config)
     except Exception as e:
         logger.error(f"Failed to start MCP server: {e}")
         sys.exit(1)