From ae58a8eb14662988bf311ff4ccc28a92326bec2d Mon Sep 17 00:00:00 2001
From: Qubasa <luis.nixos@gmail.com>
Date: Fri, 22 Sep 2023 18:34:43 +0200
Subject: [PATCH] CLI: Added custom logger

---
 pkgs/clan-cli/clan_cli/custom_logger.py       | 42 +++++++++++++++++++
 pkgs/clan-cli/clan_cli/webui/app.py           |  4 ++
 .../clan_cli/webui/routers/machines.py        |  4 ++
 pkgs/clan-cli/clan_cli/webui/routers/vms.py   |  7 +++-
 pkgs/clan-cli/clan_cli/webui/server.py        |  4 +-
 5 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 pkgs/clan-cli/clan_cli/custom_logger.py

diff --git a/pkgs/clan-cli/clan_cli/custom_logger.py b/pkgs/clan-cli/clan_cli/custom_logger.py
new file mode 100644
index 000000000..242e8028a
--- /dev/null
+++ b/pkgs/clan-cli/clan_cli/custom_logger.py
@@ -0,0 +1,42 @@
+import logging
+import datetime
+
+class CustomFormatter(logging.Formatter):
+
+    grey = "\x1b[38;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    green = "\u001b[32m"
+    blue = "\u001b[34m"
+
+    def format_str(color):
+        reset = "\x1b[0m"
+        return f"%(asctime)s - %(name)s - {color} %(levelname)s {reset} - %(message)s (%(filename)s:%(lineno)d)"
+
+    FORMATS = {
+        logging.DEBUG: format_str(blue),
+        logging.INFO: format_str(green),
+        logging.WARNING: format_str(yellow),
+        logging.ERROR: format_str(red),
+        logging.CRITICAL: format_str(bold_red)
+    }
+
+    def formatTime(self, record,datefmt=None):
+        now = datetime.datetime.now()
+        now = now.strftime("%H:%M:%S")
+        return now
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        formatter.formatTime = self.formatTime
+        return formatter.format(record)
+
+
+def register(level):
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.DEBUG)
+    ch.setFormatter(CustomFormatter())
+    logging.basicConfig(level=level, handlers=[ch])
+
diff --git a/pkgs/clan-cli/clan_cli/webui/app.py b/pkgs/clan-cli/clan_cli/webui/app.py
index 2eb53aade..8b7f645a9 100644
--- a/pkgs/clan-cli/clan_cli/webui/app.py
+++ b/pkgs/clan-cli/clan_cli/webui/app.py
@@ -2,7 +2,9 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.routing import APIRoute
 from fastapi.staticfiles import StaticFiles
+import logging
 
+from .. import custom_logger 
 from .assets import asset_path
 from .routers import flake, health, machines, root, vms
 
@@ -35,4 +37,6 @@ def setup_app() -> FastAPI:
     return app
 
 
+custom_logger.register(logging.getLogger('uvicorn').level)
 app = setup_app()
+
diff --git a/pkgs/clan-cli/clan_cli/webui/routers/machines.py b/pkgs/clan-cli/clan_cli/webui/routers/machines.py
index 61fd5edc7..a2b7778ba 100644
--- a/pkgs/clan-cli/clan_cli/webui/routers/machines.py
+++ b/pkgs/clan-cli/clan_cli/webui/routers/machines.py
@@ -19,6 +19,10 @@ from ..schemas import (
     Status,
 )
 
+# Logging setup
+import logging
+log = logging.getLogger(__name__)
+
 router = APIRouter()
 
 
diff --git a/pkgs/clan-cli/clan_cli/webui/routers/vms.py b/pkgs/clan-cli/clan_cli/webui/routers/vms.py
index 99b124c93..8a55350b8 100644
--- a/pkgs/clan-cli/clan_cli/webui/routers/vms.py
+++ b/pkgs/clan-cli/clan_cli/webui/routers/vms.py
@@ -3,15 +3,16 @@ import json
 import shlex
 from typing import Annotated, AsyncIterator
 
-from fastapi import APIRouter, Body, HTTPException, Request, status
+
+from fastapi import APIRouter, Body, HTTPException, Request, status, logger
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse, StreamingResponse
 
 from ...nix import nix_build, nix_eval
 from ..schemas import VmConfig, VmInspectResponse
 
-router = APIRouter()
 
+router = APIRouter()
 
 class NixBuildException(HTTPException):
     def __init__(self, msg: str, loc: list = ["body", "flake_attr"]):
@@ -120,8 +121,10 @@ command output:
 {stderr}
             """
         )
+import logging
 
 
 @router.post("/api/vms/create")
 async def create_vm(vm: Annotated[VmConfig, Body()]) -> StreamingResponse:
     return StreamingResponse(vm_build(vm))
+
diff --git a/pkgs/clan-cli/clan_cli/webui/server.py b/pkgs/clan-cli/clan_cli/webui/server.py
index 1b7164f30..1c052a191 100644
--- a/pkgs/clan-cli/clan_cli/webui/server.py
+++ b/pkgs/clan-cli/clan_cli/webui/server.py
@@ -7,12 +7,11 @@ import webbrowser
 from contextlib import ExitStack, contextmanager
 from pathlib import Path
 from threading import Thread
-from typing import Iterator
+from typing import (Iterator, Dict, Any)
 
 # XXX: can we dynamically load this using nix develop?
 from uvicorn import run
 
-logger = logging.getLogger(__name__)
 
 
 def defer_open_browser(base_url: str) -> None:
@@ -87,5 +86,6 @@ def start_server(args: argparse.Namespace) -> None:
             port=args.port,
             log_level=args.log_level,
             reload=args.reload,
+            access_log=args.log_level == "debug",
             headers=headers,
         )