kami_spider_monorepo/observability/logging.py

"""
Structured logging with trace context propagation.
Integrates with OpenTelemetry for log correlation.
"""

import logging
import json
from typing import Any, MutableMapping
from datetime import datetime
from contextvars import ContextVar
from core.config import settings


# Context variable for trace ID
trace_id_var: ContextVar[str] = ContextVar("trace_id", default="")


class TraceContextFormatter(logging.Formatter):
    """
    Custom formatter that adds trace context to log records.
    
    Formats logs as JSON in production, human-readable in development.
    """
    
    def __init__(self, use_json: bool = True):
        """
        Initialize formatter.
        
        Args:
            use_json: Whether to format as JSON (True) or human-readable (False)
        """
        super().__init__()
        self.use_json = use_json
    
    def format(self, record: logging.LogRecord) -> str:
        """
        Format log record with trace context.
        
        Args:
            record: Log record
        
        Returns:
            str: Formatted log message
        """
        # Get trace ID from context
        trace_id = trace_id_var.get()
        
        if self.use_json:
            # JSON format for production
            log_data = {
                "timestamp": datetime.utcnow().isoformat() + "Z",
                "level": record.levelname,
                "logger": record.name,
                "message": record.getMessage(),
                "trace_id": trace_id,
                "module": record.module,
                "function": record.funcName,
                "line": record.lineno,
            }
            
            # Add exception info if present
            if record.exc_info:
                log_data["exception"] = self.formatException(record.exc_info)
            
            # Add extra fields
            if hasattr(record, "extra"):
                log_data["extra"] = getattr(record, "extra")
            
            return json.dumps(log_data)
        else:
            # Human-readable format for development
            timestamp = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
            trace_part = f" [trace_id={trace_id}]" if trace_id else ""
            message = f"{timestamp} - {record.levelname:8s} - {record.name:30s}{trace_part} - {record.getMessage()}"
            
            if record.exc_info:
                message += "\n" + self.formatException(record.exc_info)
            
            return message


def setup_logging() -> None:
    """
    Configure application logging.
    
    This should be called at application startup.
    """
    # Determine if we should use JSON format
    use_json = settings.is_production
    
    # Create formatter
    formatter = TraceContextFormatter(use_json=use_json)
    
    # Configure root logger
    root_logger = logging.getLogger()
    root_logger.setLevel(settings.log_level)
    
    # Remove existing handlers
    root_logger.handlers.clear()
    
    # Create console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(settings.log_level)
    console_handler.setFormatter(formatter)
    
    # Add handler to root logger
    root_logger.addHandler(console_handler)
    
    # Set levels for third-party loggers
    logging.getLogger("uvicorn").setLevel(logging.INFO)
    logging.getLogger("uvicorn.access").setLevel(logging.INFO)
    logging.getLogger("uvicorn.error").setLevel(logging.INFO)
    logging.getLogger("fastapi").setLevel(logging.INFO)
    
    # Reduce noise from OpenTelemetry
    logging.getLogger("opentelemetry").setLevel(logging.WARNING)


def set_trace_id(trace_id: str) -> None:
    """
    Set trace ID in context for current request.
    
    Args:
        trace_id: Trace ID to set
    """
    trace_id_var.set(trace_id)


def get_trace_id() -> str:
    """
    Get trace ID from context.
    
    Returns:
        str: Current trace ID or empty string
    """
    return trace_id_var.get()


def get_logger(name: str) -> logging.Logger:
    """
    Get logger instance.
    
    Args:
        name: Logger name (usually __name__)
    
    Returns:
        logging.Logger: Logger instance
    """
    return logging.getLogger(name)


class LoggerAdapter(logging.LoggerAdapter):
    """
    Logger adapter that automatically includes trace context.
    """
    
    def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> tuple[str, MutableMapping[str, Any]]:
        """
        Process log message to add trace context.
        
        Args:
            msg: Log message
            kwargs: Keyword arguments
        
        Returns:
            tuple: Processed message and kwargs
        """
        trace_id = trace_id_var.get()
        if trace_id:
            extra = kwargs.get("extra", {})
            extra["trace_id"] = trace_id
            kwargs["extra"] = extra
        return msg, kwargs


def get_logger_with_trace(name: str) -> LoggerAdapter:
    """
    Get logger adapter with automatic trace context.
    
    Args:
        name: Logger name (usually __name__)
    
    Returns:
        LoggerAdapter: Logger adapter instance
    """
    logger = logging.getLogger(name)
    return LoggerAdapter(logger, {})