feat: add xray_trace_id key when tracing is active #137

Author: heitorlessa

Diff for: aws_lambda_powertools/logging/formatter.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 
 
 class JsonFormatter(logging.Formatter):
@@ -29,14 +30,27 @@ def __init__(self, **kwargs):
         # Set the default unserializable function, by default values will be cast as str.
         self.default_json_formatter = kwargs.pop("json_default", str)
         # Set the insertion order for the log messages
-        self.format_dict = dict.fromkeys(kwargs.pop("log_record_order", ["level", "location", "message", "timestamp"]))
+        self.format_dict = dict.fromkeys(
+            kwargs.pop("log_record_order", ["level", "location", "message", "xray_trace_id", "timestamp"])
+        )
+        self.reserved_keys = ["timestamp", "level", "location"]
         # Set the date format used by `asctime`
         super(JsonFormatter, self).__init__(datefmt=kwargs.pop("datefmt", None))
 
-        self.reserved_keys = ["timestamp", "level", "location"]
-        self.format_dict.update(
-            {"level": "%(levelname)s", "location": "%(funcName)s:%(lineno)d", "timestamp": "%(asctime)s", **kwargs}
-        )
+        self.format_dict.update(self._build_root_keys(**kwargs))
+
+    @staticmethod
+    def _build_root_keys(**kwargs):
+        xray_trace_id = os.getenv("_X_AMZN_TRACE_ID")
+        trace_id = xray_trace_id.split(";")[0].replace("Root=", "") if xray_trace_id else None
+
+        return {
+            "level": "%(levelname)s",
+            "location": "%(funcName)s:%(lineno)d",
+            "xray_trace_id": trace_id,
+            "timestamp": "%(asctime)s",
+            **kwargs,
+        }
 
     def update_formatter(self, **kwargs):
         self.format_dict.update(kwargs)

Diff for: tests/functional/test_aws_lambda_logging.py

@@ -170,3 +170,34 @@ def test_log_dict_key_strip_nones(stdout):
 
     # THEN the keys should only include `level`, `message`, `service`, `sampling_rate`
     assert sorted(log_dict.keys()) == ["level", "message", "sampling_rate", "service"]
+
+
+def test_log_dict_xray_is_present_when_tracing_is_enabled(stdout, monkeypatch):
+    # GIVEN a logger is initialized within a Lambda function with X-Ray enabled
+    trace_id = "1-5759e988-bd862e3fe1be46a994272793"
+    trace_header = f"Root={trace_id};Parent=53995c3f42cd8ad8;Sampled=1"
+    monkeypatch.setenv(name="_X_AMZN_TRACE_ID", value=trace_header)
+    logger = Logger(stream=stdout)
+
+    # WHEN logging a message
+    logger.info("foo")
+
+    log_dict: dict = json.loads(stdout.getvalue())
+
+    # THEN `xray_trace_id`` key should be present
+    assert log_dict["xray_trace_id"] == trace_id
+
+    monkeypatch.delenv(name="_X_AMZN_TRACE_ID")
+
+
+def test_log_dict_xray_is_not_present_when_tracing_is_disabled(stdout, monkeypatch):
+    # GIVEN a logger is initialized within a Lambda function with X-Ray disabled (default)
+    logger = Logger(stream=stdout)
+
+    # WHEN logging a message
+    logger.info("foo")
+
+    log_dict: dict = json.loads(stdout.getvalue())
+
+    # THEN `xray_trace_id`` key should not be present
+    assert "xray_trace_id" not in log_dict