From 8b810c4000092ddd259d15ed04c55cea12bc0676 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Sun, 10 Apr 2022 13:14:52 +0200
Subject: [PATCH] Unify logging
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: András Schmelczer <andras@schmelczer.dev>
---
 good_ai/src/good_ai/good_ai/deploy/serve.py   | 37 ++++++++++++++++++-
 .../src/good_ai/utilities/logger/__init__.py  |  1 +
 .../utilities/logger/custom_formatter.py      | 16 ++++----
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/good_ai/src/good_ai/good_ai/deploy/serve.py b/good_ai/src/good_ai/good_ai/deploy/serve.py
index c23f7e4..8b8d14b 100644
--- a/good_ai/src/good_ai/good_ai/deploy/serve.py
+++ b/good_ai/src/good_ai/good_ai/deploy/serve.py
@@ -1,7 +1,7 @@
 from typing import Any, Callable
 
 import uvicorn
-from fastapi import status
+from fastapi import FastAPI, status
 from fastapi.middleware.wsgi import WSGIMiddleware
 from fastapi.responses import RedirectResponse
 from typing_extensions import Never
@@ -18,6 +18,7 @@ def serve(
     function: Callable[..., Any],
     disable_docs: bool = False,
     disable_metrics: bool = False,
+    configure: Callable[[FastAPI], None]=lambda _:None
 ) -> Never:
     app = create_fastapi_app(function.__name__, disable_docs=disable_docs)
 
@@ -36,4 +37,36 @@ def serve(
             output = t.log_output(result)
         return output
 
-    uvicorn.run(app, host="0.0.0.0", port=5050)
+    configure(app)
+
+    uvicorn.run(app, host="0.0.0.0", port=5050, log_config={
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "()": "good_ai.logger.CustomFormatter",
+                "fmt": "%(asctime)s | %(levelname)8s | %(message)s",
+            },
+            "access": {
+                "()": "good_ai.logger.CustomFormatter",
+                "fmt": '%(asctime)s | %(levelname)8s | %(message)s',  # noqa: E501
+            },
+        },
+        "handlers": {
+            "default": {
+                "formatter": "default",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stderr",
+            },
+            "access": {
+                "formatter": "access",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",
+            },
+        },
+        "loggers": {
+            "uvicorn": {"handlers": ["default"], "level": "INFO"},
+            "uvicorn.error": {"level": "INFO"},
+            "uvicorn.access": {"handlers": ["access"], "level": "INFO", "propagate": False},
+        },
+    })
diff --git a/good_ai/src/good_ai/utilities/logger/__init__.py b/good_ai/src/good_ai/utilities/logger/__init__.py
index 9eb451a..5b95aad 100644
--- a/good_ai/src/good_ai/utilities/logger/__init__.py
+++ b/good_ai/src/good_ai/utilities/logger/__init__.py
@@ -1 +1,2 @@
 from .create_logger import create_logger
+from .custom_formatter import CustomFormatter
diff --git a/good_ai/src/good_ai/utilities/logger/custom_formatter.py b/good_ai/src/good_ai/utilities/logger/custom_formatter.py
index 0760ce3..81cb958 100644
--- a/good_ai/src/good_ai/utilities/logger/custom_formatter.py
+++ b/good_ai/src/good_ai/utilities/logger/custom_formatter.py
@@ -6,16 +6,16 @@ from .colors import BLUE, BOLD_RED, GREY, RED, RESET, YELLOW
 class CustomFormatter(logging.Formatter):
     def __init__(self, fmt: str):
         super().__init__()
-        self.fmt = fmt
-        self.FORMATS = {
-            logging.DEBUG: GREY + self.fmt + RESET,
-            logging.INFO: BLUE + self.fmt + RESET,
-            logging.WARNING: YELLOW + self.fmt + RESET,
-            logging.ERROR: RED + self.fmt + RESET,
-            logging.CRITICAL: BOLD_RED + self.fmt + RESET,
+        self._fmt = fmt
+        self._color_mapping = {
+            logging.DEBUG: GREY + self._fmt + RESET,
+            logging.INFO: BLUE + self._fmt + RESET,
+            logging.WARNING: YELLOW + self._fmt + RESET,
+            logging.ERROR: RED + self._fmt + RESET,
+            logging.CRITICAL: BOLD_RED + self._fmt + RESET,
         }
 
     def format(self, record: logging.LogRecord) -> str:
-        log_fmt = self.FORMATS.get(record.levelno)
+        log_fmt = self._color_mapping.get(record.levelno)
         formatter = logging.Formatter(log_fmt)
         return formatter.format(record)