fix: ensures xray_trace_id is refreshed

Author: heitorlessa

aws_lambda_powertools/logging/formatter.py

@@ -30,28 +30,28 @@ 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", "xray_trace_id", "timestamp"])
-        )
+        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.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
-
+    def _build_root_keys(self, **kwargs):
         return {
             "level": "%(levelname)s",
             "location": "%(funcName)s:%(lineno)d",
-            "xray_trace_id": trace_id,
             "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)
 
@@ -90,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}
 

docs/content/core/logger.mdx

@@ -305,7 +305,7 @@ This can be fixed by either ensuring both has the `service` value as `payment`,
 
 You might want to continue to use the same date formatting style, or override `location` to display the `package.function_name:line_number` as you previously had.
 
-Logger allows you to either change the format or suppress the following keys altogether at the initialization: `location`, `timestamp`, `level`, `xray_trace_id`, and `datefmt`
+Logger allows you to either change the format or suppress the following keys altogether at the initialization: `location`, `timestamp`, `level`, and `datefmt`
 
 ```python
 from aws_lambda_powertools import Logger
@@ -317,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`, `xray_trace_id`, 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

@@ -201,3 +201,29 @@ def test_log_dict_xray_is_not_present_when_tracing_is_disabled(stdout, monkeypat
 
     # 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")