From 0ebc16a78a424da86072f7444a6f65536388a3b4 Mon Sep 17 00:00:00 2001
From: simon-schoonjans <simon.schoonjans@waylay.io>
Date: Fri, 8 Sep 2023 09:03:14 +0200
Subject: [PATCH] feat: add support for props_to_labels

---
 README.md                |  2 ++
 logging_loki/emitter.py  | 18 ++++++++++++++----
 logging_loki/handlers.py |  5 +++--
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index c74f63b..42718fd 100644
--- a/README.md
+++ b/README.md
@@ -29,6 +29,7 @@ handler = logging_loki.LokiHandler(
     headers={"X-Scope-OrgID": "example-id"},
     auth=("username", "password"),
     version="1",
+    props_to_labels: Optional[list[str]] = ["foo"]
 )
 
 logger = logging.getLogger("my-logger")
@@ -62,6 +63,7 @@ handler_loki = logging_loki.LokiHandler(
     headers={"X-Scope-OrgID": "example-id"},
     auth=("username", "password"),
     version="1",
+    props_to_labels: Optional[list[str]] = ["foo"]
 )
 logging.handlers.QueueListener(queue, handler_loki)
 
diff --git a/logging_loki/emitter.py b/logging_loki/emitter.py
index 978559d..3fb0aff 100644
--- a/logging_loki/emitter.py
+++ b/logging_loki/emitter.py
@@ -32,7 +32,14 @@ class LokiEmitter(abc.ABC):
     label_replace_with = const.label_replace_with
     session_class = requests.Session
 
-    def __init__(self, url: str, tags: Optional[dict] = None, headers: Optional[dict] = None, auth: BasicAuth = None, as_json: bool = False):
+    def __init__(self, 
+        url: str, 
+        tags: Optional[dict] = None, 
+        headers: Optional[dict] = None, 
+        auth: BasicAuth = None, 
+        as_json: bool = False,
+        props_to_labels: Optional[list[str]] = None
+    ):
         """
         Create new Loki emitter.
 
@@ -52,6 +59,8 @@ def __init__(self, url: str, tags: Optional[dict] = None, headers: Optional[dict
         self.auth = auth
         #: Optional bool, send record as json?
         self.as_json = as_json
+        #: Optional list, send record as json?
+        self.props_to_labels = props_to_labels or []
 
         self._session: Optional[requests.Session] = None
         self._lock = threading.Lock()
@@ -106,10 +115,11 @@ def build_tags(self, record: logging.LogRecord) -> Dict[str, Any]:
         tags[self.level_tag] = record.levelname.lower()
         tags[self.logger_tag] = record.name
 
-        extra_tags = getattr(record, "tags", {})
-        if not isinstance(extra_tags, dict):
-            return tags
+        extra_tags = {k: getattr(record, k) for k in self.props_to_labels if getattr(record, k)}
+        if isinstance(passed_tags := getattr(record, "tags", {}), dict):
+            extra_tags = extra_tags | passed_tags
 
+        
         for tag_name, tag_value in extra_tags.items():
             cleared_name = self.format_label(tag_name)
             if cleared_name:
diff --git a/logging_loki/handlers.py b/logging_loki/handlers.py
index 4651f3c..af363c8 100644
--- a/logging_loki/handlers.py
+++ b/logging_loki/handlers.py
@@ -45,7 +45,8 @@ def __init__(
         headers: Optional[dict] = None,
         auth: Optional[emitter.BasicAuth] = None,
         version: Optional[str] = None,
-        as_json: Optional[bool] = False
+        as_json: Optional[bool] = False,
+        props_to_labels: Optional[list[str]] = None
     ):
         """
         Create new Loki logging handler.
@@ -71,7 +72,7 @@ def __init__(
         version = version or const.emitter_ver
         if version not in self.emitters:
             raise ValueError("Unknown emitter version: {0}".format(version))
-        self.emitter = self.emitters[version](url, tags, headers, auth, as_json)
+        self.emitter = self.emitters[version](url, tags, headers, auth, as_json, props_to_labels)
 
     def handleError(self, record):  # noqa: N802
         """Close emitter and let default handler take actions on error."""