aws_python_custom_logger.py

class LogLevel(Enum):
    CRITICAL = CRITICAL
    DEBUG = DEBUG
    ERROR = ERROR
    FATAL = FATAL
    INFO = INFO
    NOTSET = NOTSET
    WARNING = WARNING


class FormatterJSON(logging.Formatter):
    def format(self, record: LogRecord) -> str:
        record.message = record.getMessage()
        if self.usesTime():
            record.asctime = self.formatTime(record, self.datefmt)
        log_dict = {
            "message": record.message,
            "extra": record.__dict__.get("extra", {}),
            "module": record.module,
            "levelname": record.levelname,
            "time": "%(asctime)s.%(msecs)dZ" % dict(asctime=record.asctime, msecs=record.msecs),
            "aws_request_id": getattr(record, "aws_request_id", "00000000-0000-0000-0000-000000000000"),
        }
        result: str = json.dumps(log_dict, default=lambda value: str(value)) 
        return result + "\n" # the newline char I've found is necesssary in container image Lambas, otherwise multiple logs are grouped together in one CloudWatch row


class CustomLogger(logging.Logger):
    def makeRecord(
        self,
        name: str,
        level: int,
        fn: str,
        lno: int,
        msg: str,
        args: Any,
        exc_info: Union[Tuple[Type[BaseException], BaseException, Optional[Any]], Tuple[None, None, None], None],
        func: Optional[str] = None,
        extra: Optional[Mapping[str, Any]] = None,
        sinfo: Optional[str] = None,
    ) -> LogRecord:
        rv = logging.LogRecord(name, level, fn, lno, msg, args, exc_info, func, sinfo)
        if extra is not None:
            rv.__dict__["extra"] = extra
        return rv


def use_logging(level: LogLevel = LogLevel.INFO) -> logging.Logger:
    custom_logger = logging.getLogger()
    custom_logger.__class__ = CustomLogger
    custom_logger.setLevel(level.value)
    formatter = FormatterJSON("[%(levelname)s]\t%(asctime)s.%(msecs)dZ\t%(levelno)s\t%(x)s\n", "%Y-%m-%dT%H:%M:%S")
    for handler in custom_logger.handlers:
        handler.setFormatter(formatter)
    return custom_logger