From ef77a757c64c601783da0038d66b19722492e654 Mon Sep 17 00:00:00 2001
From: Nikita Melkozerov <nikita@some.engineering>
Date: Fri, 3 Nov 2023 11:34:04 +0000
Subject: [PATCH] Support providing a logging context

---
 fixcloudutils/logging/json_logger.py | 6 +++++-
 tests/logging_test.py                | 8 ++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/fixcloudutils/logging/json_logger.py b/fixcloudutils/logging/json_logger.py
index 5634373..8868545 100644
--- a/fixcloudutils/logging/json_logger.py
+++ b/fixcloudutils/logging/json_logger.py
@@ -21,7 +21,7 @@
 #  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 import json
 from logging import Formatter, LogRecord
-from typing import Mapping, Optional, Dict
+from typing import Mapping, Optional, Dict, Callable
 
 
 class JsonFormatter(Formatter):
@@ -34,11 +34,13 @@ def __init__(
         fmt_dict: Mapping[str, str],
         time_format: str = "%Y-%m-%dT%H:%M:%S",
         static_values: Optional[Dict[str, str]] = None,
+        get_logging_context: Callable[[], Dict[str, str]] = lambda: {},
     ) -> None:
         super().__init__()
         self.fmt_dict = fmt_dict
         self.time_format = time_format
         self.static_values = static_values or {}
+        self.get_context_values = get_logging_context
         self.__uses_time = "asctime" in self.fmt_dict.values()
 
     def usesTime(self) -> bool:  # noqa: N802
@@ -55,6 +57,8 @@ def prop(name: str) -> str:
 
         message_dict = {fmt_key: prop(fmt_val) for fmt_key, fmt_val in self.fmt_dict.items()}
         message_dict.update(self.static_values)
+        if context := self.get_context_values():
+            message_dict.update(context)
         if record.exc_info:
             if not record.exc_text:
                 record.exc_text = self.formatException(record.exc_info)
diff --git a/tests/logging_test.py b/tests/logging_test.py
index d92dfda..8fadfb6 100644
--- a/tests/logging_test.py
+++ b/tests/logging_test.py
@@ -22,14 +22,18 @@
 import logging
 
 import prometheus_client
+from typing import Dict
 
 from fixcloudutils.logging import JsonFormatter, PrometheusLoggingCounter
 
 
 def test_json_logging() -> None:
-    format = JsonFormatter({"level": "levelname", "message": "message"})
+    def logging_context() -> Dict[str, str]:
+        return {"foo": "bar"}
+
+    format = JsonFormatter({"level": "levelname", "message": "message"}, get_logging_context=logging_context)
     record = logging.getLogger().makeRecord("test", logging.INFO, "test", 1, "test message", (), None)
-    assert format.format(record) == '{"level": "INFO", "message": "test message"}'
+    assert format.format(record) == '{"level": "INFO", "message": "test message", "foo": "bar"}'
 
 
 def test_prometheus_counter() -> None: