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)


class Formatter(logging.Formatter):
    """Formatter for the combined HTTP/WebSocket access log.

    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:
        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, and a
    filter dropping stock uvicorn's WebSocket chatter from ``uvicorn.error``.
    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")

    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