COBY : specs + task 1

COBY/utils/logging.py

@@ -0,0 +1,149 @@
+"""
+Logging utilities for the COBY system.
+"""
+
+import logging
+import logging.handlers
+import sys
+import uuid
+from pathlib import Path
+from typing import Optional
+from contextvars import ContextVar
+
+# Context variable for correlation ID
+correlation_id: ContextVar[Optional[str]] = ContextVar('correlation_id', default=None)
+
+
+class CorrelationFilter(logging.Filter):
+    """Add correlation ID to log records"""
+    
+    def filter(self, record):
+        record.correlation_id = correlation_id.get() or 'N/A'
+        return True
+
+
+class COBYFormatter(logging.Formatter):
+    """Custom formatter with correlation ID support"""
+    
+    def __init__(self, include_correlation_id: bool = True):
+        self.include_correlation_id = include_correlation_id
+        
+        if include_correlation_id:
+            fmt = '%(asctime)s - %(name)s - %(levelname)s - [%(correlation_id)s] - %(message)s'
+        else:
+            fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+            
+        super().__init__(fmt, datefmt='%Y-%m-%d %H:%M:%S')
+
+
+def setup_logging(
+    level: str = 'INFO',
+    log_file: Optional[str] = None,
+    max_file_size: int = 100,  # MB
+    backup_count: int = 5,
+    enable_correlation_id: bool = True,
+    console_output: bool = True
+) -> None:
+    """
+    Set up logging configuration for the COBY system.
+    
+    Args:
+        level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        log_file: Path to log file (None = no file logging)
+        max_file_size: Maximum log file size in MB
+        backup_count: Number of backup files to keep
+        enable_correlation_id: Whether to include correlation IDs in logs
+        console_output: Whether to output logs to console
+    """
+    # Convert string level to logging constant
+    numeric_level = getattr(logging, level.upper(), logging.INFO)
+    
+    # Create root logger
+    root_logger = logging.getLogger()
+    root_logger.setLevel(numeric_level)
+    
+    # Clear existing handlers
+    root_logger.handlers.clear()
+    
+    # Create formatter
+    formatter = COBYFormatter(include_correlation_id=enable_correlation_id)
+    
+    # Add correlation filter if enabled
+    correlation_filter = CorrelationFilter() if enable_correlation_id else None
+    
+    # Console handler
+    if console_output:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(numeric_level)
+        console_handler.setFormatter(formatter)
+        if correlation_filter:
+            console_handler.addFilter(correlation_filter)
+        root_logger.addHandler(console_handler)
+    
+    # File handler
+    if log_file:
+        # Create log directory if it doesn't exist
+        log_path = Path(log_file)
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Rotating file handler
+        file_handler = logging.handlers.RotatingFileHandler(
+            log_file,
+            maxBytes=max_file_size * 1024 * 1024,  # Convert MB to bytes
+            backupCount=backup_count
+        )
+        file_handler.setLevel(numeric_level)
+        file_handler.setFormatter(formatter)
+        if correlation_filter:
+            file_handler.addFilter(correlation_filter)
+        root_logger.addHandler(file_handler)
+    
+    # Set specific logger levels
+    logging.getLogger('websockets').setLevel(logging.WARNING)
+    logging.getLogger('urllib3').setLevel(logging.WARNING)
+    logging.getLogger('requests').setLevel(logging.WARNING)
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Get a logger instance with the specified name.
+    
+    Args:
+        name: Logger name (typically __name__)
+        
+    Returns:
+        logging.Logger: Logger instance
+    """
+    return logging.getLogger(name)
+
+
+def set_correlation_id(corr_id: Optional[str] = None) -> str:
+    """
+    Set correlation ID for current context.
+    
+    Args:
+        corr_id: Correlation ID (generates UUID if None)
+        
+    Returns:
+        str: The correlation ID that was set
+    """
+    if corr_id is None:
+        corr_id = str(uuid.uuid4())[:8]  # Short UUID
+    
+    correlation_id.set(corr_id)
+    return corr_id
+
+
+def get_correlation_id() -> Optional[str]:
+    """
+    Get current correlation ID.
+    
+    Returns:
+        str: Current correlation ID or None
+    """
+    return correlation_id.get()
+
+
+def clear_correlation_id() -> None:
+    """Clear correlation ID from current context."""
+    correlation_id.set(None)