baoshi-measure-cargo/request_logging.py

import json
import logging
import os
import time
from datetime import datetime
from logging.handlers import RotatingFileHandler

from fastapi import FastAPI, Request
from starlette.concurrency import iterate_in_threadpool


def setup_request_logging(app: FastAPI, max_response_len: int, max_bytes: int, backup_count: int) -> None:
    logger = logging.getLogger("cargo_height.request")
    _setup_request_logger(logger, max_bytes=max_bytes, backup_count=backup_count)

    @app.middleware("http")
    async def request_log_middleware(request: Request, call_next):
        request_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
        start = time.perf_counter()
        try:
            response = await call_next(request)
        except Exception:
            elapsed_ms = (time.perf_counter() - start) * 1000
            logger.exception(
                "request_time=%s method=%s path=%s status=%s duration_ms=%.2f response=%s",
                request_time,
                request.method,
                request.url.path,
                500,
                elapsed_ms,
                "internal_error",
            )
            raise

        body = b""
        async for chunk in response.body_iterator:
            body += chunk
        response.body_iterator = iterate_in_threadpool(iter([body]))

        response_text = _parse_response_text(body, response.headers.get("content-type", ""))
        if len(response_text) > max_response_len:
            response_text = response_text[:max_response_len] + "...(truncated)"

        elapsed_ms = (time.perf_counter() - start) * 1000
        logger.info(
            "request_time=%s method=%s path=%s status=%s duration_ms=%.2f response=%s",
            request_time,
            request.method,
            request.url.path,
            response.status_code,
            elapsed_ms,
            response_text,
        )
        return response


def _setup_request_logger(logger: logging.Logger, max_bytes: int, backup_count: int) -> None:
    log_dir = os.path.join(os.getcwd(), "Log")
    os.makedirs(log_dir, exist_ok=True)
    log_file = os.path.join(log_dir, "request.log")

    logger.setLevel(logging.INFO)
    logger.propagate = False
    if logger.handlers:
        return

    formatter = logging.Formatter(
        "%(asctime)s [%(levelname)s] %(name)s - %(message)s",
        "%Y-%m-%d %H:%M:%S",
    )

    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=max_bytes,
        backupCount=backup_count,
        encoding="utf-8",
    )
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)

    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.INFO)
    stream_handler.setFormatter(formatter)

    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)


def _parse_response_text(body: bytes, content_type: str) -> str:
    if not body:
        return ""
    if "application/json" in content_type:
        try:
            return json.dumps(json.loads(body), ensure_ascii=False)
        except Exception:
            return body.decode("utf-8", errors="replace")
    return body.decode("utf-8", errors="replace")