Source code for ccat_data_transfer.logging_utils

from contextvars import ContextVar
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, Any
from sqlalchemy.orm import Session
from ccat_ops_db import models
import logging
import os
import sys
from .exceptions import DataTransferError
from .config.config import ccat_data_transfer_settings


# Per-task correlation context (#117). Set when a Celery task starts (see
# CCATEnhancedSQLAlchemyTask.__call__) and merged into every structured log line
# emitted while the task runs, so logs are correlatable to their operation with
# no call-site changes. The key is (operation_type, operation_id) — the pair the
# whole system already correlates on — plus the Celery task_id.
_log_correlation: ContextVar[dict] = ContextVar("log_correlation", default={})




def set_log_correlation(**fields: Any):
    """Set the correlation fields auto-injected into structured log lines.

    Returns the ContextVar token so the caller can restore the previous context
    with :func:`reset_log_correlation` once the task body returns.
    """
    return _log_correlation.set(dict(fields))





def reset_log_correlation(token) -> None:
    """Restore the correlation context to the state before ``set_log_correlation``."""
    _log_correlation.reset(token)





def get_log_correlation() -> dict:
    """Return the currently-active correlation fields (empty outside a task)."""
    return _log_correlation.get()





class BBCPLogHandler:


    def __init__(self, base_log_path: str = ccat_data_transfer_settings.BBCP_LOG_PATH):
        self.base_log_path = Path(base_log_path)
        self.base_log_path.mkdir(parents=True, exist_ok=True)
        os.chmod(self.base_log_path, 0o750)




    def get_log_path(
        self, transfer_id: int, timestamp: Optional[datetime] = None
    ) -> Path:
        """Generate structured log path for BBCP transfer."""
        if timestamp is None:
            timestamp = datetime.now()

        date_path = self.base_log_path / timestamp.strftime("%Y/%m/%d")
        date_path.mkdir(parents=True, exist_ok=True)
        os.chmod(date_path, 0o750)

        return (
            date_path
            / f"bbcp_transfer_{transfer_id}_{timestamp.strftime('%Y%m%d_%H%M%S')}.log"
        )




    def store_bbcp_output(
        self,
        session: Session,
        data_transfer: models.DataTransfer,
        stdout: bytes,
        stderr: bytes,
        success: bool,
        timestamp: Optional[datetime] = None,
    ) -> models.DataTransferLog:
        """Store BBCP output and create minimal log entry."""
        if timestamp is None:
            timestamp = datetime.now()

        # Get path for log file
        log_path = self.get_log_path(data_transfer.id, timestamp)

        # Store combined output with restricted permissions
        fd = os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o640)
        with os.fdopen(fd, "w") as f:
            f.write("=== STDOUT ===\n")
            f.write(stdout)
            f.write("\n=== STDERR ===\n")
            f.write(stderr)

        # Create log entry first
        log_entry = models.DataTransferLog(
            data_transfer_id=data_transfer.id,
            data_transfer=data_transfer,
            timestamp=timestamp,
            log_path=str(log_path),
            status="success" if success else "failed",
        )

        session.add(log_entry)
        return log_entry




def _serialize_value(value: Any) -> Any:
    """Helper function to serialize values for logging"""
    if hasattr(value, "__dict__"):
        return str(value)
    return value


def _single_line(value) -> str:
    """Collapse all whitespace (incl. newlines) to single spaces for a one-line logfmt value."""
    return " ".join(str(value).split())




class LogfmtFormatter(logging.Formatter):
    """Formats log records as a single logfmt line.

    Produces: ts=<utc-iso8601> level=<lowercase> logger=<short-module> <record.getMessage()>

    Tracebacks are appended on following lines (no leading ts= prefix) so that
    downstream multiline collectors that key on '^ts=' fold them into the event.
    """



    def format(self, record: logging.LogRecord) -> str:
        # Guard: already formatted by this formatter (e.g. double-emit protection)
        msg = record.getMessage()
        if msg.startswith("ts="):
            return msg

        ts = datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat()
        level = record.levelname.lower()
        # Strip the package prefix so the logger field stays short and readable.
        logger_name = record.name.removeprefix("ccat_data_transfer.")

        line = f"ts={ts} level={level} logger={logger_name} {msg}"

        if record.exc_info:
            line += "\n" + self.formatException(record.exc_info)
        if record.stack_info:
            line += "\n" + self.formatStack(record.stack_info)

        return line






def configure_logging(level: Optional[Any] = None) -> None:
    """Install exactly one LogfmtFormatter handler on the package base logger.

    Idempotent: if a LogfmtFormatter handler is already present on the base
    logger, only the level is updated.  This function is safe to call multiple
    times (issue #124 will call it each poll cycle to change level at runtime).

    Args:
        level: logging level name ("INFO", "DEBUG", ...) or int constant.
               Defaults to settings.LOG_LEVEL ("INFO" unless overridden).
    """
    base = logging.getLogger("ccat_data_transfer")

    # Resolve level
    if level is None:
        raw = getattr(ccat_data_transfer_settings, "LOG_LEVEL", "INFO")
        resolved_level = logging.getLevelName(str(raw).upper())
    elif isinstance(level, str):
        resolved_level = logging.getLevelName(level.upper())
    else:
        resolved_level = level

    # getLevelName returns "Level <name>" string for unknown names (Python 3.4+).
    # Fall back to INFO so a misconfigured LOG_LEVEL never silences the logger.
    if not isinstance(resolved_level, int):
        resolved_level = logging.INFO

    base.setLevel(resolved_level)

    # Idempotent handler install: only add when no LogfmtFormatter handler exists
    has_logfmt = any(
        isinstance(h.formatter, LogfmtFormatter) for h in base.handlers
    )
    if not has_logfmt:
        # Clear any stale handlers (e.g. from old get_structured_logger calls)
        base.handlers = []
        handler = logging.StreamHandler()
        handler.setFormatter(LogfmtFormatter())
        base.addHandler(handler)

    # Prevent lines from also appearing on the root logger (third-party loggers
    # like sqlalchemy/celery are NOT affected -- they have their own hierarchy).
    base.propagate = False






def apply_runtime_log_level(op_config, session) -> None:
    """Read LOG_LEVEL from operational config and apply it. Never raises."""
    try:
        configure_logging(op_config.get(session, "LOG_LEVEL"))
    except Exception:
        pass  # log-level adjustment must never break the caller




class StructuredLogger:
    """Wrapper for structured logging with consistent logfmt formatting.

    Each call to debug/info/warning/error/exception formats the message body as
    ``msg="<text>" key=value ...`` and delegates to the underlying stdlib logger,
    which carries the body through to LogfmtFormatter where the ts/level/logger
    prefix fields are prepended once.
    """



    def __init__(self, logger: logging.Logger):
        self.logger = logger




    def setLevel(self, level: int) -> None:
        """Set the logging level of the logger."""
        self.logger.setLevel(level)




    def getEffectiveLevel(self) -> int:
        """Get the effective logging level of the logger."""
        return self.logger.getEffectiveLevel()


    def _format_message(self, message: str, **kwargs) -> str:
        """Format the message body as logfmt key=value pairs.

        Produces: msg="<message>" key=value ...

        Values are quoted when the string representation contains a space, '=',
        '"', or is empty.  The ts/level/logger prefix fields are added by
        LogfmtFormatter, NOT here -- one occurrence per line guaranteed.
        """
        serialized_kwargs = {k: _serialize_value(v) for k, v in kwargs.items()}

        # Auto-inject the running task's correlation fields (#117) so logs emitted
        # inside a task carry operation_type/operation_id/task_id with no call-site
        # change. Explicit kwargs win over correlation on key collision.
        correlation = {
            k: v for k, v in get_log_correlation().items() if k not in serialized_kwargs
        }
        fields = {**correlation, **serialized_kwargs}

        # msg is always quoted (messages commonly contain spaces)
        formatted_pairs = [f'msg="{message}"']

        for key, value in fields.items():
            str_value = str(value) if not isinstance(value, str) else value
            needs_quotes = isinstance(value, str) and (
                " " in str_value
                or "=" in str_value
                or '"' in str_value
                or str_value == ""
            )
            if needs_quotes:
                formatted_pairs.append(f'{key}="{value}"')
            else:
                formatted_pairs.append(f"{key}={value}")

        return " ".join(formatted_pairs)



    def debug(self, message: str, **kwargs) -> None:
        """Log a debug message."""
        self.logger.debug(self._format_message(message=message, **kwargs))




    def info(self, message: str, **kwargs) -> None:
        """Log an info message."""
        self.logger.info(self._format_message(message=message, **kwargs))




    def warning(self, message: str, **kwargs) -> None:
        """Log a warning message."""
        self.logger.warning(self._format_message(message=message, **kwargs))




    def error(self, message: str, error: Optional[Exception] = None, **kwargs) -> None:
        """Log an error message.

        If error is a BaseException, summary fields (error_type, error) go on the
        main logfmt line and the traceback is carried via exc_info so LogfmtFormatter
        appends it on following lines — keeping the summary single-line.
        If error is a plain string (legacy callers), it is embedded directly as the
        error= field with no traceback.
        """
        fields = dict(kwargs)
        exc_info = None

        if error is not None:
            if isinstance(error, BaseException):
                fields["error_type"] = error.__class__.__name__
                fields["error"] = _single_line(str(error))
                if isinstance(error, DataTransferError):
                    fields["transfer_id"] = error.transfer_id
                # Carry the real traceback via exc_info so LogfmtFormatter emits it
                # on lines after the summary — not embedded inside a quoted value.
                exc_info = error
            else:
                # Legacy: caller passed error=str(e) as a plain string
                fields["error"] = _single_line(error)

        self.logger.error(
            self._format_message(message=message, **fields), exc_info=exc_info
        )




    def critical(self, message: str, **kwargs) -> None:
        """Log a critical message."""
        self.logger.critical(self._format_message(message=message, **kwargs))




    def exception(self, message: str, **kwargs) -> None:
        """Log an exception with traceback at ERROR level.
        Must be called from an except block — logger.exception sets exc_info=True
        itself, so LogfmtFormatter appends the traceback on following lines.
        """
        fields = dict(kwargs)
        exc = sys.exc_info()[1]
        if exc is not None:
            fields["error_type"] = exc.__class__.__name__
            fields["error"] = _single_line(str(exc))
            if isinstance(exc, DataTransferError):
                fields["transfer_id"] = exc.transfer_id
        # logger.exception sets exc_info=True automatically; do NOT pass traceback=
        # as a kwarg — that would embed multiline content in the logfmt summary line.
        self.logger.exception(self._format_message(message=message, **fields))






def get_structured_logger(name: str) -> StructuredLogger:
    """Return a StructuredLogger for the given module name.

    Ensures the package base logger ('ccat_data_transfer') has a LogfmtFormatter
    handler installed via configure_logging() so output always reaches stderr.
    Module-level loggers propagate to the base logger -- no per-logger handler needed.
    """
    # One-time (idempotent) setup of the base logger handler + level
    configure_logging()
    return StructuredLogger(logging.getLogger(name))





def setup_celery_logging():
    """Configure Celery logging to use the same format as our application logging"""
    from celery.signals import setup_logging

    @setup_logging.connect
    def config_loggers(*args, **kwargs):
        # Prevent Celery from creating its own logger
        return True