149 lines
4.3 KiB
Python
149 lines
4.3 KiB
Python
"""
|
|
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) |