From 3a5b7e954b38ea2ae4879507c3a42e67edae3ba3 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Wed, 21 Aug 2024 15:01:28 +0300
Subject: [PATCH] feat: support ISOLATE_LOG_LABELS (#151)

---
 src/isolate/logger.py | 14 ++++++++++++++
 tests/test_logger.py  | 22 ++++++++++++++++++++++
 2 files changed, 36 insertions(+)
 create mode 100644 tests/test_logger.py

diff --git a/src/isolate/logger.py b/src/isolate/logger.py
index a30d1e0..fdd7a1f 100644
--- a/src/isolate/logger.py
+++ b/src/isolate/logger.py
@@ -1,15 +1,29 @@
 import json
+import os
 
 
 # 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):
         record = {
             "isolate_source": source.name,
             "level": level.name,
             "message": message,
+            **self.log_labels,
         }
         print(json.dumps(record))
 
diff --git a/tests/test_logger.py b/tests/test_logger.py
new file mode 100644
index 0000000..b8d624f
--- /dev/null
+++ b/tests/test_logger.py
@@ -0,0 +1,22 @@
+import json
+
+from isolate.logger import IsolateLogger
+
+
+def test_logger(monkeypatch):
+    labels = {
+        "foo": "$MYENVVAR1",
+        "bar": "$MYENVVAR2",
+        "baz": "baz",
+        "qux": "qux",
+    }
+    monkeypatch.setenv("MYENVVAR1", "myenvvar1")
+    monkeypatch.setenv("MYENVVAR2", "myenvvar2")
+    monkeypatch.setenv("ISOLATE_LOG_LABELS", json.dumps(labels))
+    logger = IsolateLogger()
+    assert logger.log_labels == {
+        "foo": "myenvvar1",
+        "bar": "myenvvar2",
+        "baz": "baz",
+        "qux": "qux",
+    }