fix(logger): strip xray_trace_id when explicitly disabled (aws-powertools#2852)

heitorlessa · web-flow · commit 7dbdd81ad572 · 2023-07-26T15:15:18.000+02:00
diff --git a/aws_lambda_powertools/logging/formatter.py b/aws_lambda_powertools/logging/formatter.py
@@ -236,8 +236,12 @@ def _build_default_keys():
             "timestamp": "%(asctime)s",
         }
 
-    @staticmethod
-    def _get_latest_trace_id():
+    def _get_latest_trace_id(self):
+        xray_trace_id_key = self.log_format.get("xray_trace_id", "")
+        if xray_trace_id_key is None:
+            # key is explicitly disabled; ignore it. e.g., Logger(xray_trace_id=None)
+            return None
+
         xray_trace_id = os.getenv(constants.XRAY_TRACE_ID_ENV)
         return xray_trace_id.split(";")[0].replace("Root=", "") if xray_trace_id else None
 
diff --git a/tests/functional/test_logger_powertools_formatter.py b/tests/functional/test_logger_powertools_formatter.py
@@ -251,6 +251,27 @@ def test_log_dict_xray_is_updated_when_tracing_id_changes(stdout, monkeypatch, s
     monkeypatch.delenv(name="_X_AMZN_TRACE_ID")
 
 
+def test_log_dict_xray_is_not_present_when_explicitly_disabled(
+    stdout: io.StringIO,
+    monkeypatch: pytest.MonkeyPatch,
+    service_name: str,
+):
+    # GIVEN a logger is initialized within a Lambda function with X-Ray enabled
+    # and X-Ray Trace ID key is explicitly disabled
+    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(service=service_name, stream=stdout, xray_trace_id=None)
+
+    # 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_custom_std_log_attribute(stdout, service_name):
     # GIVEN a logger where we have a standard log attr process
     # https://docs.python.org/3/library/logging.html#logrecord-attributes
