Refactor correlation causation ids

chocs_middleware/trace/logger.py

@@ -3,6 +3,7 @@
 import traceback
 from dataclasses import is_dataclass, asdict
 from datetime import date, datetime, time
+from decimal import Decimal
 from inspect import istraceback
 from typing import Dict, Optional, IO, Union, Any, List
 
@@ -67,20 +68,20 @@ def default(self, data: Any) -> Any:
 
 
 class JsonFormatter(logging.Formatter):
-    def __init__(self, json_encoder: json.JSONEncoder = JsonEncoder(), message_format: str = "[{level}] {timestamp} {msg}"):
+    def __init__(
+        self, json_encoder: json.JSONEncoder = JsonEncoder(), message_format: str = "[{level}] {timestamp} {msg}"
+    ):
         self.json_encoder = json_encoder
         self.message_format = message_format
         super(JsonFormatter, self).__init__()
 
     @staticmethod
     def get_message(record: logging.LogRecord) -> Any:
-        if isinstance(record.msg, str):
+        if isinstance(record.msg, (str, float, int, bool, Decimal, type(None))):
             return record.msg
 
         if record.levelname != "DEBUG":
-            msg = f"Dumping objects is prohibited at `{record.levelname}` log level."
-            record.levelname = "ERROR"
-            return msg
+            return repr(record.msg)
 
         return record.msg
 

chocs_middleware/trace/middleware.py

@@ -86,13 +86,13 @@ def detect_sentry(self):
     def handle(self, request: HttpRequest, next: MiddlewareHandler) -> HttpResponse:
 
         # Retrieve values from headers
-        request_id = str(request.headers.get("x-request-id", self.generate_id()))
+        request_id = str(self.generate_id())
         correlation_id = str(request.headers.get("x-correlation-id", request_id))
         causation_id = str(request.headers.get("x-causation-id", request_id))
 
         # Update headers when needed
-        if "x-request-id" not in request.headers:
-            request.headers["x-request-id"] = request_id
+        request.headers["x-request-id"] = request_id
+
         if "x-correlation-id" not in request.headers:
             request.headers["x-correlation-id"] = correlation_id
         if "x-causation-id" not in request.headers:
@@ -102,18 +102,21 @@ def handle(self, request: HttpRequest, next: MiddlewareHandler) -> HttpResponse:
             "method": str(request.method),
             "path": str(request.path),
             "route": str(request.route.route),
-            "id": request_id,
-            "correlation_id": correlation_id,
-            "causation_id": causation_id,
         }
 
         Logger.set_tag("request", request_tag)
+        Logger.set_tag("x-request-id", request_id)
+        Logger.set_tag("x-correlation-id", correlation_id)
+        Logger.set_tag("x-causation-id", causation_id)
 
         # Populate sentry tags
         if self._use_sentry:
             from sentry_sdk import set_tag
 
             set_tag("request", request_tag)
+            set_tag("x-request-id", request_id)
+            set_tag("x-correlation-id", correlation_id)
+            set_tag("x-causation-id", causation_id)
 
         # Automatically add extra headers to requests library
         if self._use_http and self._http_strategy == HttpStrategy.REQUESTS:
@@ -134,7 +137,6 @@ def wrapped_request(method, url, **kwargs):
                 if "headers" not in kwargs:
                     kwargs["headers"] = {}
 
-                kwargs["headers"]["x-request-id"] = self.generate_id()
                 kwargs["headers"]["x-causation-id"] = request_id
                 kwargs["headers"]["x-correlation-id"] = correlation_id
 
@@ -155,7 +157,6 @@ def urllib_wrapped_request(self_, method, url, fields=None, headers=None, **urlo
                 if not headers:
                     headers = {}
 
-                headers["x-request-id"] = self.generate_id()
                 headers["x-causation-id"] = request_id
                 headers["x-correlation-id"] = correlation_id
 
@@ -167,7 +168,5 @@ def urllib_wrapped_request(self_, method, url, fields=None, headers=None, **urlo
 
         response = next(request)
         response.headers.set("x-request-id", request_id)
-        response.headers.set("x-causation-id", causation_id)
-        response.headers.set("x-correlation-id", correlation_id)
 
         return response

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "chocs_middleware.trace"
-version = "0.3.1"
+version = "0.4.0"
 description = "Http tracing middleware for chocs library."
 authors = ["Dawid Kraczkowski <dawid.kraczkowski@gmail.com>"]
 license = "MIT"

tests/trace/test_logger.py

@@ -133,7 +133,7 @@ def test_can_log_a_dict() -> None:
     assert log["log_message"] == {"test": "ok"}
 
 
-def test_fail_log_a_dict_in_non_debug_level() -> None:
+def test_log_a_dict_in_non_debug_level_as_string() -> None:
     # given
     logger_stream = StringIO()
     logger = Logger.get("test_cannot_log_a_dict", log_stream=logger_stream)
@@ -147,8 +147,8 @@ def test_fail_log_a_dict_in_non_debug_level() -> None:
     json_payload = record[record.find('\t'):]
     log = json.loads(json_payload)
 
-    assert log["level"] == "ERROR"
-    assert log["log_message"] != {"test": "ok"}
+    assert log["level"] == "INFO"
+    assert log["log_message"] == str({"test": "ok"})
 
 
 def test_can_retrieve_same_logger_multiple_times() -> None:

tests/trace/test_middleware.py

@@ -19,7 +19,6 @@ def requests_support_mock(url, request):
 
         assert "x-correlation-id" in request.headers
         assert "x-causation-id" in request.headers
-        assert "x-request-id" in request.headers
 
         return "ok"
 
@@ -43,7 +42,6 @@ def urllib_support_mock(*args, **kwargs):
     headers = kwargs["headers"]
     assert "x-correlation-id" in headers
     assert "x-causation-id" in headers
-    assert "x-request-id" in headers
 
     return urllib3.HTTPResponse("ok")
 
@@ -80,10 +78,8 @@ def requests_support_mock(url, request):
 
         assert "x-correlation-id" in request.headers
         assert "x-causation-id" in request.headers
-        assert "x-request-id" in request.headers
         assert request.headers.get("x-correlation-id")[0:13] == "service-name-"
         assert request.headers.get("x-causation-id")[0:13] == "service-name-"
-        assert request.headers.get("x-request-id")[0:13] == "service-name-"
 
         return "ok"