Merge pull request #150 from heitorlessa/feat/logger-add-xray-trace-id

feat: add xray_trace_id key when tracing is active #137

Author: heitorlessa

aws_lambda_powertools/logging/formatter.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 
 
 class JsonFormatter(logging.Formatter):
@@ -30,13 +31,26 @@ def __init__(self, **kwargs):
         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.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))
+
+    def _build_root_keys(self, **kwargs):
+        return {
+            "level": "%(levelname)s",
+            "location": "%(funcName)s:%(lineno)d",
+            "timestamp": "%(asctime)s",
+            **kwargs,
+        }
+
+    @staticmethod
+    def _get_latest_trace_id():
+        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 trace_id
 
     def update_formatter(self, **kwargs):
         self.format_dict.update(kwargs)
@@ -76,6 +90,9 @@ def format(self, record):  # noqa: A003
         if record.exc_text:
             log_dict["exception"] = record.exc_text
 
+        # fetch latest X-Ray Trace ID, if any
+        log_dict.update({"xray_trace_id": self._get_latest_trace_id()})
+
         # Filter out top level key with values that are None
         log_dict = {k: v for k, v in log_dict.items() if v is not None}
 

aws_lambda_powertools/logging/logger.py

@@ -235,14 +235,14 @@ def handler(event, context):
 
         @functools.wraps(lambda_handler)
         def decorate(event, context):
+            lambda_context = build_lambda_context_model(context)
+            cold_start = _is_cold_start()
+            self.structure_logs(append=True, cold_start=cold_start, **lambda_context.__dict__)
+
             if log_event:
                 logger.debug("Event received")
                 self.info(event)
 
-            lambda_context = build_lambda_context_model(context)
-            cold_start = _is_cold_start()
-
-            self.structure_logs(append=True, cold_start=cold_start, **lambda_context.__dict__)
             return lambda_handler(event, context)
 
         return decorate

docs/content/core/logger.mdx

@@ -47,7 +47,7 @@ Logger(service="payment", level="INFO")
 
 ## Standard structured keys
 
-Your Logger will always include the following keys to your structured logging:
+Your Logger will include the following keys to your structured logging, by default:
 
 Key | Type | Example | Description
 ------------------------------------------------- | ------------------------------------------------- | --------------------------------------------------------------------------------- | -------------------------------------------------
@@ -57,6 +57,7 @@ Key | Type | Example | Description
 **service** | str | "payment" | Service name defined. "service_undefined" will be used if unknown
 **sampling_rate** | int |  0.1 | Debug logging sampling rate in percentage e.g. 10% in this case
 **message** | any |  "Collecting payment" | Log statement value. Unserializable JSON values will be casted to string
+**xray_trace_id** | str | "1-5759e988-bd862e3fe1be46a994272793" | X-Ray Trace ID when Lambda function has enabled Tracing
 
 ## Capturing Lambda context info
 
@@ -316,7 +317,7 @@ logger = Logger(stream=stdout, location="[%(funcName)s] %(module)s", datefmt="fa
 logger = Logger(stream=stdout, location=None) # highlight-line
 ```
 
-Alternatively, you can also change the order of the following log record keys via the `log_record_order` parameter: `level`, location`, message`, and timestamp`
+Alternatively, you can also change the order of the following log record keys via the `log_record_order` parameter: `level`, `location`, `message`, and `timestamp`
 
 ```python
 from aws_lambda_powertools import Logger

tests/functional/test_aws_lambda_logging.py

@@ -170,3 +170,60 @@ 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
+
+
+def test_log_dict_xray_is_updated_when_tracing_id_changes(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")
+
+    # and Trace ID changes to mimick a new invocation
+    trace_id_2 = "1-5759e988-bd862e3fe1be46a949393982437"
+    trace_header_2 = f"Root={trace_id_2};Parent=53995c3f42cd8ad8;Sampled=1"
+    monkeypatch.setenv(name="_X_AMZN_TRACE_ID", value=trace_header_2)
+
+    logger.info("foo bar")
+
+    log_dict, log_dict_2 = [json.loads(line.strip()) for line in stdout.getvalue().split("\n") if line]
+
+    # THEN `xray_trace_id`` key should be different in both invocations
+    assert log_dict["xray_trace_id"] == trace_id
+    assert log_dict_2["xray_trace_id"] == trace_id_2
+
+    monkeypatch.delenv(name="_X_AMZN_TRACE_ID")