From debe1b9e80595ba0c369a2ea9450bf3116b81289 Mon Sep 17 00:00:00 2001
From: Matteo Ferrando <matteo.ferrando2@gmail.com>
Date: Sun, 15 Sep 2024 12:25:07 -0500
Subject: [PATCH] refactor: make IsolateLogger more explicit

---
 src/isolate/logger.py        | 41 ++++++++++++++++++++++++------------
 src/isolate/server/server.py |  5 +++--
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/isolate/logger.py b/src/isolate/logger.py
index fdd7a1f..473fe07 100644
--- a/src/isolate/logger.py
+++ b/src/isolate/logger.py
@@ -1,24 +1,17 @@
 import json
 import os
 
+from isolate.logs import LogLevel, LogSource
+
 
 # NOTE: we probably should've created a proper `logging.getLogger` here,
 # but it handling `source` would be not trivial, so we are better off
 # just keeping it simple for now.
 class IsolateLogger:
-    def __init__(self):
-        self.log_labels = {}
-        raw = os.getenv("ISOLATE_LOG_LABELS")
-        if raw:
-            labels = json.loads(raw)
-            for key, value in labels.items():
-                if value.startswith("$"):
-                    expanded = os.getenv(value[1:])
-                else:
-                    expanded = value
-                self.log_labels[key] = expanded
-
-    def log(self, level, message, source):
+    def __init__(self, log_labels: dict[str, str]):
+        self.log_labels = log_labels
+
+    def log(self, level: LogLevel, message: str, source: LogSource):
         record = {
             "isolate_source": source.name,
             "level": level.name,
@@ -27,5 +20,25 @@ def log(self, level, message, source):
         }
         print(json.dumps(record))
 
+    @classmethod
+    def with_env_expanded(cls, labels: dict[str, str]):
+        for key, value in labels.items():
+            if value.startswith("$"):
+                expanded = os.getenv(value[1:])
+            else:
+                expanded = value
+            if expanded is not None:
+                labels[key] = expanded
+
+        return cls(labels)
+
+
+_labels = {}
+try:
+    raw = os.getenv("ISOLATE_LOG_LABELS")
+    if raw:
+        _labels: dict[str, str] = json.loads(raw)
+except BaseException:
+    print("Failed to parse ISOLATE_LOG_LABELS")
 
-logger = IsolateLogger()
+ENV_LOGGER = IsolateLogger.with_env_expanded(labels=_labels)
diff --git a/src/isolate/server/server.py b/src/isolate/server/server.py
index 67fb2a4..7bdb441 100644
--- a/src/isolate/server/server.py
+++ b/src/isolate/server/server.py
@@ -29,7 +29,7 @@
 from isolate.backends.virtualenv import VirtualPythonEnvironment
 from isolate.connections.grpc import AgentError, LocalPythonGRPC
 from isolate.connections.grpc.configuration import get_default_options
-from isolate.logger import logger
+from isolate.logger import ENV_LOGGER, IsolateLogger
 from isolate.logs import Log, LogLevel, LogSource
 from isolate.server import definitions, health
 from isolate.server.health_server import HealthServicer
@@ -453,9 +453,10 @@ def _proxy_to_queue(
 @dataclass
 class LogHandler:
     messages: Queue
+    logger: IsolateLogger = ENV_LOGGER
 
     def handle(self, log: Log) -> None:
-        logger.log(log.level, log.message, source=log.source)
+        self.logger.log(log.level, log.message, source=log.source)
         self._add_log_to_queue(log)
 
     def _add_log_to_queue(self, log: Log) -> None: