nix-builder-autoscaler/agent/nix_builder_autoscaler/logging.py

"""Structured JSON logging setup."""

from __future__ import annotations

import json
import logging
import sys
from datetime import UTC, datetime
from typing import Any


class JSONFormatter(logging.Formatter):
    """Format log records as single-line JSON."""

    EXTRA_FIELDS = (
        "slot_id",
        "reservation_id",
        "instance_id",
        "request_id",
        "error",
        "category",
        "count",
        "ids",
        "idle_seconds",
    )

    def format(self, record: logging.LogRecord) -> str:
        """Format a log record as JSON."""
        entry: dict[str, Any] = {
            "ts": datetime.now(UTC).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }
        for field in self.EXTRA_FIELDS:
            val = getattr(record, field, None)
            if val is not None:
                entry[field] = val
        if record.exc_info and record.exc_info[1] is not None:
            entry["exception"] = self.formatException(record.exc_info)
        return json.dumps(entry, default=str)


def setup_logging(level: str = "INFO") -> None:
    """Configure the root logger with JSON output to stderr.

    Args:
        level: Log level name (DEBUG, INFO, WARNING, ERROR).
    """
    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(JSONFormatter())

    root = logging.getLogger()
    root.handlers.clear()
    root.addHandler(handler)
    root.setLevel(getattr(logging, level.upper(), logging.INFO))