fastapi-vue-setup/fastapi-vue/fastapi_vue/logging.py

"""Logging integration: tracerite loading and colored access log formatting.

The access log middleware supplies colored fields (``client``, ``status``,
``method``, ``host``, ``path``, ``extra``, ``timing``) via ``extra=``.  When
colors are disabled the ANSI escape codes are stripped from the assembled
output so the same formatting code path produces plain text.
"""

from __future__ import annotations

import logging
import re
import sys
from contextlib import suppress
from copy import deepcopy
from typing import Literal

import tracerite
from uvicorn.config import Config

from .accesslog import AccessLogMiddleware

ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m")

ACCESS_LOG_FMT = "%(client)s %(status)s %(method)s %(host)s%(path)s %(extra)s%(timing)s"

ACCESS_LOGGER = "fastapi_vue.access"


def strip_ansi(text: str) -> str:
    return ANSI_ESCAPE_RE.sub("", text)


_LEVEL_EMOJI = {
    logging.DEBUG: "🐛",
    logging.INFO: "🔷",
    logging.WARNING: "❗",
    logging.ERROR: "🛑",
    logging.CRITICAL: "🚨",
}


def _level_prefix(record: logging.LogRecord) -> str:
    emoji = _LEVEL_EMOJI.get(record.levelno)
    return f"{emoji} " if emoji else f"{record.levelname}: "


class Formatter(logging.Formatter):
    """Formatter for both access records and ordinary log messages.

    Records with the middleware's access fields (``client`` etc.) are
    formatted from those; anything else gets an emoji level prefix
    (``LEVEL: `` fallback for unknown levels) in place of uvicorn's
    ``levelprefix``.  ANSI codes are stripped when colors are disabled.

    Instantiation always loads tracerite, and with ``access=True`` also
    installs the access-log middleware: ``dictConfig`` builds formatters while
    uvicorn applies ``log_config``, which happens before the app is loaded —
    including in reload/worker subprocesses that re-import the config without
    calling ``fastapi_vue.server.run()`` again.
    """

    def __init__(
        self,
        fmt: str | None = None,
        datefmt: str | None = None,
        style: Literal["%", "{", "$"] = "%",
        use_colors: bool | None = None,
        *,
        access: bool = False,
    ) -> None:
        tracerite.load()
        if access:
            install_access_log()
        if use_colors in (True, False):
            self.use_colors = use_colors
        else:
            self.use_colors = sys.stdout.isatty()
        super().__init__(fmt=fmt, datefmt=datefmt, style=style)

    def formatMessage(self, record: logging.LogRecord) -> str:
        if "client" not in record.__dict__:
            return _level_prefix(record) + record.getMessage()
        formatted = super().formatMessage(record)
        if not self.use_colors:
            formatted = strip_ansi(formatted)
        return formatted


class WebSocketChatterFilter(logging.Filter):
    """Drop stock uvicorn WebSocket handshake/chatter records.

    Stock uvicorn logs WS handshakes (``'%s - "WebSocket %s" ...'``) and the
    websockets library's "connection open/closed" chatter to ``uvicorn.error``,
    ungated by ``access_log``.  Our middleware logs WebSockets itself.
    """

    _PREFIXES = ('%s - "WebSocket ', "connection open", "connection closed", "connection rejected")

    def filter(self, record: logging.LogRecord) -> bool:
        msg = record.msg
        if not isinstance(msg, str):
            return True
        return not msg.startswith(self._PREFIXES)


_installed = False


def install_access_log() -> None:
    """Wrap apps loaded by uvicorn with AccessLogMiddleware, once per process.

    The guard is deliberately module-level: reload/worker subprocesses
    re-import this module, resetting it so the patch is re-applied there.
    """
    global _installed
    if _installed:
        return
    _installed = True

    original_load = Config.load

    def load(self):  # noqa: ANN001, ANN202
        original_load(self)
        if not isinstance(self.loaded_app, AccessLogMiddleware):
            self.loaded_app = AccessLogMiddleware(self.loaded_app)

    Config.load = load  # type: ignore[method-assign]


def patch_log_config(log_config, *, access_log: bool = True):  # noqa: ANN001, ANN201
    """Patch a uvicorn log_config dict for our logging, best-effort.

    Users presumably base their config on uvicorn's default dict, but any
    shape is tolerated: pieces that do not fit the config's structure are
    silently skipped.  Non-dict configs (e.g. an ini file path) pass through
    untouched.

    Always adds an unreferenced NullHandler whose Formatter instantiation
    loads tracerite in every process uvicorn applies the config in, a
    filter dropping stock uvicorn's WebSocket chatter from ``uvicorn.error``,
    an emoji-level-prefix Formatter in place of uvicorn's stock ``default``
    formatter (a user-supplied one wins), a root logger entry so
    ``logging.info()`` et al. print through the default handler, and a
    no-prefix ``kanta`` logger entry (likewise).  With ``access_log``,
    additionally rewires the ``access`` formatter to our Formatter and
    attaches its handler to our ``fastapi_vue.access`` logger.  We must not
    attach handlers to ``uvicorn.access``: uvicorn gates its own
    protocol-level access logging on ``uvicorn.access.hasHandlers()``.
    """
    if not isinstance(log_config, dict):
        return log_config
    config = deepcopy(log_config)

    with suppress(Exception):
        config["formatters"]["fastapi_vue"] = {"()": "fastapi_vue.logging.Formatter"}
        config["handlers"]["fastapi_vue"] = {
            "class": "logging.NullHandler",
            "formatter": "fastapi_vue",
        }

    with suppress(Exception):
        filters = config.setdefault("filters", {})
        filters["ws_chatter"] = {"()": "fastapi_vue.logging.WebSocketChatterFilter"}
        handler_filters = config["handlers"]["default"].setdefault("filters", [])
        if "ws_chatter" not in handler_filters:
            handler_filters.append("ws_chatter")

    # Emoji level prefixes for ordinary logs, replacing uvicorn's stock
    # default formatter; a user-supplied default formatter is left alone.
    with suppress(Exception):
        default = config["formatters"]["default"]
        if default.get("()") in (None, "uvicorn.logging.DefaultFormatter"):
            config["formatters"]["default"] = {
                "()": "fastapi_vue.logging.Formatter",
                "fmt": "%(message)s",
                "use_colors": None,
            }

    # uvicorn's default config leaves the root logger handlerless, eating
    # logging.info() et al.; route root through uvicorn's default handler.
    with suppress(Exception):
        root = config.setdefault("root", {})
        root.setdefault("level", "INFO")
        root_handlers = root.setdefault("handlers", [])
        if "default" not in root_handlers:
            root_handlers.append("default")

    # kanta-style output (diffs, colored headers) prints without prefixes,
    # like our access log.  A user-supplied "kanta" logger entry wins.
    with suppress(Exception):
        config["formatters"].setdefault("plain", {"fmt": "%(message)s"})
        config["handlers"].setdefault(
            "plain",
            {
                "class": "logging.StreamHandler",
                "formatter": "plain",
                "stream": "ext://sys.stderr",
            },
        )
        config.setdefault("loggers", {}).setdefault(
            "kanta",
            {"handlers": ["plain"], "level": "INFO", "propagate": False},
        )

    if access_log:
        with suppress(Exception):
            config["formatters"]["access"] = {
                "()": "fastapi_vue.logging.Formatter",
                "fmt": ACCESS_LOG_FMT,
                "use_colors": None,
                "access": True,
            }
        with suppress(Exception):
            if "access" in config["handlers"]:
                config.setdefault("loggers", {})[ACCESS_LOGGER] = {
                    "handlers": ["access"],
                    "level": "INFO",
                    "propagate": False,
                }

    return config