From f8ec3123286e63148123ccba781dcd699f705e1d Mon Sep 17 00:00:00 2001
From: Fijxu <fijxu@nadeko.net>
Date: Thu, 19 Sep 2024 21:24:20 -0300
Subject: [PATCH] Logger: Add color support for different log levels

---
 config/config.example.yml       |  9 +++++++++
 src/invidious.cr                |  5 ++++-
 src/invidious/config.cr         |  2 ++
 src/invidious/helpers/logger.cr | 19 +++++++++++++++++--
 4 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/config/config.example.yml b/config/config.example.yml
index 219aa03f..37b932ea 100644
--- a/config/config.example.yml
+++ b/config/config.example.yml
@@ -222,6 +222,15 @@ https_only: false
 ##
 #log_level: Info
 
+##
+## Enables colors in logs. Useful for debugging purposes
+## This is overridden if "-k" or "--colorize" 
+## are passed on the command line.
+##
+## Accepted values: true, false
+## Default: false
+##
+#colorize_logs: false
 
 # -----------------------------
 #  Features
diff --git a/src/invidious.cr b/src/invidious.cr
index 3804197e..d9a479d1 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -117,6 +117,9 @@ Kemal.config.extra_options do |parser|
   parser.on("-l LEVEL", "--log-level=LEVEL", "Log level, one of #{LogLevel.values} (default: #{CONFIG.log_level})") do |log_level|
     CONFIG.log_level = LogLevel.parse(log_level)
   end
+  parser.on("-k", "--colorize", "Colorize logs") do
+    CONFIG.colorize_logs = true
+  end
   parser.on("-v", "--version", "Print version") do
     puts SOFTWARE.to_pretty_json
     exit
@@ -133,7 +136,7 @@ if CONFIG.output.upcase != "STDOUT"
   FileUtils.mkdir_p(File.dirname(CONFIG.output))
 end
 OUTPUT = CONFIG.output.upcase == "STDOUT" ? STDOUT : File.open(CONFIG.output, mode: "a")
-LOGGER = Invidious::LogHandler.new(OUTPUT, CONFIG.log_level)
+LOGGER = Invidious::LogHandler.new(OUTPUT, CONFIG.log_level, CONFIG.colorize_logs)
 
 # Check table integrity
 Invidious::Database.check_integrity(CONFIG)
diff --git a/src/invidious/config.cr b/src/invidious/config.cr
index c4ddcdb3..d8543d35 100644
--- a/src/invidious/config.cr
+++ b/src/invidious/config.cr
@@ -68,6 +68,8 @@ class Config
   property output : String = "STDOUT"
   # Default log level, valid YAML values are ints and strings, see src/invidious/helpers/logger.cr
   property log_level : LogLevel = LogLevel::Info
+  # Enables colors in logs. Useful for debugging purposes
+  property colorize_logs : Bool = false
   # Database configuration with separate parameters (username, hostname, etc)
   property db : DBConfig? = nil
 
diff --git a/src/invidious/helpers/logger.cr b/src/invidious/helpers/logger.cr
index b443073e..36a3a7f9 100644
--- a/src/invidious/helpers/logger.cr
+++ b/src/invidious/helpers/logger.cr
@@ -1,3 +1,5 @@
+require "colorize"
+
 enum LogLevel
   All   = 0
   Trace = 1
@@ -10,7 +12,7 @@ enum LogLevel
 end
 
 class Invidious::LogHandler < Kemal::BaseLogHandler
-  def initialize(@io : IO = STDOUT, @level = LogLevel::Debug)
+  def initialize(@io : IO = STDOUT, @level = LogLevel::Debug, @color : Bool = true)
   end
 
   def call(context : HTTP::Server::Context)
@@ -39,10 +41,23 @@ class Invidious::LogHandler < Kemal::BaseLogHandler
     @io.flush
   end
 
+  def color(level)
+    case level
+    when LogLevel::Trace then :cyan
+    when LogLevel::Debug then :green
+    when LogLevel::Info  then :white
+    when LogLevel::Warn  then :yellow
+    when LogLevel::Error then :red
+    when LogLevel::Fatal then :magenta
+    else                      :default
+    end
+  end
+
   {% for level in %w(trace debug info warn error fatal) %}
     def {{level.id}}(message : String)
       if LogLevel::{{level.id.capitalize}} >= @level
-        puts("#{Time.utc} [{{level.id}}] #{message}")
+        puts("#{Time.utc} [{{level.id}}] #{message}".colorize(color(LogLevel::{{level.id.capitalize}})).toggle(@color))
+
       end
     end
   {% end %}