SQS: add automatic propagation of AWSTraceHeader attribute when sending message (#12810)

bentsku · web-flow · commit 05a922a4ffd2 · 2025-06-27T22:53:43.000+02:00
diff --git a/localstack-core/localstack/services/sqs/provider.py b/localstack-core/localstack/services/sqs/provider.py
@@ -1679,6 +1679,13 @@ def _create_message_attributes(
             MessageSystemAttributeName.SenderId: context.account_id,  # not the account ID in AWS
             MessageSystemAttributeName.SentTimestamp: str(now(millis=True)),
         }
+        # we are not using the `context.trace_context` here as it is automatically populated
+        # AWS only adds the `AWSTraceHeader` attribute if the header is explicitly present
+        # TODO: check maybe with X-Ray Active mode?
+        if "X-Amzn-Trace-Id" in context.request.headers:
+            result[MessageSystemAttributeName.AWSTraceHeader] = str(
+                context.request.headers["X-Amzn-Trace-Id"]
+            )
 
         if message_system_attributes is not None:
             for attr in message_system_attributes:
diff --git a/tests/aws/services/sqs/test_sqs.py b/tests/aws/services/sqs/test_sqs.py
@@ -3530,6 +3530,45 @@ def test_system_attributes_have_no_effect_on_attr_md5(self, sqs_create_queue, aw
             == "5ae4d5d7636402d80f4eb6d213245a88"
         )
 
+    @markers.aws.validated
+    def test_aws_trace_header_propagation(self, sqs_create_queue, aws_sqs_client, snapshot):
+        def add_xray_header(request, **_kwargs):
+            request.headers["X-Amzn-Trace-Id"] = (
+                "Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1"
+            )
+
+        try:
+            aws_sqs_client.meta.events.register("before-send.sqs.SendMessage", add_xray_header)
+
+            queue_url = sqs_create_queue()
+
+            msg_attrs_provider = {
+                "timestamp": {"StringValue": "1493147359900", "DataType": "Number"}
+            }
+
+            aws_sqs_client.send_message(
+                QueueUrl=queue_url, MessageBody="test", MessageAttributes=msg_attrs_provider
+            )
+
+            response = aws_sqs_client.receive_message(
+                QueueUrl=queue_url,
+                AttributeNames=["SentTimestamp", "AWSTraceHeader"],
+                MaxNumberOfMessages=1,
+                MessageAttributeNames=["All"],
+                VisibilityTimeout=2,
+                WaitTimeSeconds=2,
+            )
+
+            assert len(response["Messages"]) == 1
+            message = response["Messages"][0]
+            snapshot.match("xray-msg", message)
+            assert (
+                message["Attributes"]["AWSTraceHeader"]
+                == "Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1"
+            )
+        finally:
+            aws_sqs_client.meta.events.unregister("before-send.sqs.SendMessage", add_xray_header)
+
     @markers.aws.validated
     def test_inflight_message_requeue(self, sqs_create_queue, aws_client):
         visibility_timeout = 3
diff --git a/tests/aws/services/sqs/test_sqs.snapshot.json b/tests/aws/services/sqs/test_sqs.snapshot.json
@@ -3974,5 +3974,43 @@
         }
       }
     }
+  },
+  "tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_aws_trace_header_propagation[sqs]": {
+    "recorded-date": "27-06-2025, 10:55:55",
+    "recorded-content": {
+      "xray-msg": {
+        "Attributes": {
+          "AWSTraceHeader": "Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1",
+          "SentTimestamp": "timestamp"
+        },
+        "Body": "test",
+        "MD5OfBody": "098f6bcd4621d373cade4e832627b4f6",
+        "MD5OfMessageAttributes": "235c5c510d26fb653d073faed50ae77c",
+        "MessageAttributes": {
+          "timestamp": "timestamp"
+        },
+        "MessageId": "<uuid:1>",
+        "ReceiptHandle": "<receipt-handle:1>"
+      }
+    }
+  },
+  "tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_aws_trace_header_propagation[sqs_query]": {
+    "recorded-date": "27-06-2025, 10:55:56",
+    "recorded-content": {
+      "xray-msg": {
+        "Attributes": {
+          "AWSTraceHeader": "Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1",
+          "SentTimestamp": "timestamp"
+        },
+        "Body": "test",
+        "MD5OfBody": "098f6bcd4621d373cade4e832627b4f6",
+        "MD5OfMessageAttributes": "235c5c510d26fb653d073faed50ae77c",
+        "MessageAttributes": {
+          "timestamp": "timestamp"
+        },
+        "MessageId": "<uuid:1>",
+        "ReceiptHandle": "<receipt-handle:1>"
+      }
+    }
   }
 }
diff --git a/tests/aws/services/sqs/test_sqs.validation.json b/tests/aws/services/sqs/test_sqs.validation.json
@@ -1,4 +1,22 @@
 {
+  "tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_aws_trace_header_propagation[sqs]": {
+    "last_validated_date": "2025-06-27T10:55:55+00:00",
+    "durations_in_seconds": {
+      "setup": 0.47,
+      "call": 0.71,
+      "teardown": 0.15,
+      "total": 1.33
+    }
+  },
+  "tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_aws_trace_header_propagation[sqs_query]": {
+    "last_validated_date": "2025-06-27T10:55:56+00:00",
+    "durations_in_seconds": {
+      "setup": 0.0,
+      "call": 0.68,
+      "teardown": 0.15,
+      "total": 0.83
+    }
+  },
   "tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_change_message_visibility_after_visibility_timeout_expiration[sqs]": {
     "last_validated_date": "2024-04-30T13:33:26+00:00"
   },
