From f9da090705da09e1e3bdcf6fa8784c5cf51fc3fd Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Sat, 23 Dec 2023 21:19:40 +0100
Subject: [PATCH] Enable logging to Grafana Loki

---
 akibabot/logging.py | 51 +++++++++++++++++++++++++++++++++++++++++++++
 akibabot/main.py    | 25 +++++++++++-----------
 akibabot/xmpp.py    |  7 ++++++-
 pyproject.toml      |  6 ++++++
 4 files changed, 76 insertions(+), 13 deletions(-)
 create mode 100644 akibabot/logging.py

diff --git a/akibabot/logging.py b/akibabot/logging.py
new file mode 100644
index 0000000..3280b93
--- /dev/null
+++ b/akibabot/logging.py
@@ -0,0 +1,51 @@
+import sys
+import logging
+import logging.handlers
+from multiprocessing import Queue
+
+
+def init_loki_logging(config: dict[str, str]):
+    """Configures the logging to Loki."""
+    import logging_loki
+
+    queue = Queue(-1)
+    handler = logging.handlers.QueueHandler(queue)
+    handler_loki = logging_loki.LokiHandler(
+        url=config["loki"]["url"],
+        version="1",
+        tags={
+            "application": "akibabot",
+        },
+    )
+    formatter = logging.Formatter(
+        "name=%(name)s level=%(levelname)s time=%(asctime)s %(message)s",
+        datefmt="%Y-%m-%d-%H:%M:%S",
+    )
+    logging.handlers.QueueListener(queue, handler_loki)
+    logging.getLogger().addHandler(handler_loki)
+
+
+def init_logging(config: dict[str, str]):
+    """Initialises the logging"""
+    # Check if we should use Loki and configure it, if enabled.
+    log_config = config.get("logging", {})
+    loki = config.get("loki", {})
+    if loki.get("enable", False):
+        init_loki_logging(config)
+
+    # Configure logging to the console, if enabled.
+    logger = logging.getLogger()
+    if log_config.get("console", True):
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        )
+        console_handler = logging.StreamHandler(stream=sys.stdout)
+        console_handler.setFormatter(formatter)
+        logger.addHandler(console_handler)
+
+    # Set the log level
+    level = {
+        "info": logging.INFO,
+        "debug": logging.DEBUG,
+    }.get(log_config.get("level", "info"))
+    logger.setLevel(level)
diff --git a/akibabot/main.py b/akibabot/main.py
index b107605..8d70eda 100644
--- a/akibabot/main.py
+++ b/akibabot/main.py
@@ -10,6 +10,7 @@ import toml
 import akibapass_downloader.login
 import akibapass_downloader.episode
 
+from akibabot.logging import init_logging
 from akibabot.xmpp import send_notification
 
 
@@ -69,23 +70,21 @@ def get_episode_numbers(
     ]
 
 
-def setup_logging():
-    """Sets up the logging."""
-    logging.basicConfig(level=logging.INFO)
-
-
 def main():
-    # Setup logging
-    setup_logging()
-
     # Load config
     config_file = sys.argv[1]
     config = toml.load(config_file)
+
+    # Load secrets
+    log = logging.getLogger("akibabot")
     if (secrets_file := config["general"].get("secrets_file", None)) is not None:
-        logging.info("Processing secrets config file...")
+        log.info("Processing secrets config file...")
         secrets = toml.load(secrets_file)
         config.update(secrets)
 
+    # Setup logging
+    init_logging(config)
+
     # Sanity checks
     xmpp_from_jid = config["xmpp"]["jid"]
     xmpp_to_jid = config["xmpp"]["to"]
@@ -95,7 +94,7 @@ def main():
     storage_path = Path(config["general"]["path"])
     shows: list[Show] = []
     for show in config["shows"]:
-        logging.info("Processing %s", show["name"])
+        log.info("Processing %s", show["name"])
         max_episodes = show["max_episodes"]
         local_episodes = get_episode_numbers(
             show["name"],
@@ -117,9 +116,11 @@ def main():
 
     # Exit early when there is no show left to handle.
     if not shows:
-        print("Nothing to do.")
+        log.info("Nothing to do.")
         sys.exit(0)
 
+    return
+
     # Load cookies, if available
     cookies = {}
     cookies_path = Path(config["general"]["cookie_path"])
@@ -140,7 +141,7 @@ def main():
 
     # Iterate over the shows to process and download the new episodes
     for show in shows:
-        logging.debug("Processing %s", show.name)
+        logging.info("Processing %s", show.name)
         episodes_path = build_episodes_path(
             show.name,
             show.year,
diff --git a/akibabot/xmpp.py b/akibabot/xmpp.py
index f553b9d..91ca4e0 100644
--- a/akibabot/xmpp.py
+++ b/akibabot/xmpp.py
@@ -1,3 +1,5 @@
+import logging
+
 import aioxmpp
 
 
@@ -7,6 +9,7 @@ async def send_notification(
     """Sends a notification of a new show download to @to using the account @jid:@password.
     @show_name is the name of show. @episode is the number of the episode that was just downloaded.
     """
+    logging.getLogger("akibabot.xmpp").debug("Sending notification via XMPP")
     client = aioxmpp.PresenceManagedClient(
         aioxmpp.JID.fromstr(jid),
         aioxmpp.make_security_layer(password),
@@ -17,5 +20,7 @@ async def send_notification(
             to=aioxmpp.JID.fromstr(to),
             type_=aioxmpp.MessageType.CHAT,
         )
-        msg.body[None] = f"Episode {episode} of {show_name} is now downloaded and available!"
+        msg.body[
+            None
+        ] = f"Episode {episode} of {show_name} is now downloaded and available!"
         await stream.send(msg)
diff --git a/pyproject.toml b/pyproject.toml
index 028ce58..91fa68b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,12 @@ dev = [
     "pylint"
 ]
 
+# For logging to Loki
+loki = [
+    # Loki logging
+    "python-logging-loki"
+]
+
 [project.scripts]
 akibabot = "akibabot.main:main"