src/common/logger.py

import hashlib
import logging
import os
import re
import uuid
from contextvars import ContextVar
from datetime import datetime

import colorlog
import pytz
from dotenv import load_dotenv

load_dotenv()

_request_id = ContextVar("request_id", default=None)

VIETNAM_TZ = pytz.timezone("Asia/Ho_Chi_Minh")


def get_colored_uuid(uuid_str: str) -> str:
    colors = [
        "\033[97;41m",
        "\033[97;42m",
        "\033[97;43m",
        "\033[97;44m",
        "\033[97;45m",
        "\033[97;46m",
        "\033[30;102m",
        "\033[30;103m",
        "\033[30;104m",
        "\033[30;105m",
        "\033[30;106m",
        "\033[97;100m",
        "\033[97;101m",
        "\033[97;102m",
        "\033[97;104m",
        "\033[97;105m",
    ]
    hash_value = int(hashlib.md5(uuid_str.encode()).hexdigest()[:8], 16)
    return f"{colors[hash_value % len(colors)]}{uuid_str}\033[0m"


class ColoredFormatter(colorlog.ColoredFormatter):
    def formatTime(self, record, datefmt=None):
        dt = datetime.fromtimestamp(record.created, tz=VIETNAM_TZ)
        return dt.strftime(datefmt) if datefmt else dt.isoformat()

    def format(self, record):
        rid = getattr(record, "request_id", None)
        if rid:
            record.request_id = get_colored_uuid(rid)
        return super().format(record)


class FileFormatter(logging.Formatter):
    def formatTime(self, record, datefmt=None):
        dt = datetime.fromtimestamp(record.created, tz=VIETNAM_TZ)
        return dt.strftime(datefmt) if datefmt else dt.isoformat()

    def format(self, record):
        rid = getattr(record, "request_id", None)
        if rid:
            ansi = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]")
            record.request_id = ansi.sub("", str(rid))
        return super().format(record)


class UUIDFilter(logging.Filter):
    def filter(self, record):
        rid = _request_id.get()
        if rid is None:
            rid = str(uuid.uuid4())[:]
            _request_id.set(rid)
        record.request_id = rid
        return True


def setup_logger(name="APP", log_file="logs/log.log") -> logging.Logger:
    logger = colorlog.getLogger(name)
    if logger.handlers:
        return logger

    logger.setLevel(logging.INFO)
    logger.propagate = False

    os.makedirs(os.path.dirname(log_file), exist_ok=True)

    console_handler = colorlog.StreamHandler()
    console_handler.setFormatter(
        ColoredFormatter(
            "%(asctime)s-%(log_color)s%(levelname)s%(reset)s-%(request_id)s - %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
            log_colors={
                "DEBUG": "cyan",
                "INFO": "green",
                "WARNING": "yellow",
                "ERROR": "red,bold",
                "CRITICAL": "white,bg_red,bold",
            },
        )
    )

    file_handler = logging.FileHandler(log_file, encoding="utf-8")
    file_handler.setFormatter(
        FileFormatter("%(asctime)s-%(levelname)s-%(request_id)s - %(message)s")
    )

    logger.addFilter(UUIDFilter())
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)
    return logger


logger = setup_logger()
request_id_context = _request_id