feat(tracing): Add span links from SNS messages (#2363)

* feat(tracing): Add span links from SNS messages

* fix: Update SNS capture_serverless test

---------

Co-authored-by: Riccardo Magliocchetti <riccardo.magliocchetti@gmail.com>

Author: jackylamhk

elasticapm/contrib/serverless/aws.py

@@ -235,11 +235,20 @@ def __enter__(self):
                 transaction_name = "RECEIVE {}".format(record["eventSourceARN"].split(":")[5])
 
         if "Records" in self.event:
+            # SQS
             links = [
                 TraceParent.from_string(record["messageAttributes"]["traceparent"]["stringValue"])
                 for record in self.event["Records"][:1000]
                 if "messageAttributes" in record and "traceparent" in record["messageAttributes"]
             ]
+            # SNS
+            links += [
+                TraceParent.from_string(record["Sns"]["MessageAttributes"]["traceparent"]["Value"])
+                for record in self.event["Records"][:1000]
+                if "Sns" in record
+                and "MessageAttributes" in record["Sns"]
+                and "traceparent" in record["Sns"]["MessageAttributes"]
+            ]
         else:
             links = []
 

tests/contrib/serverless/aws_sns_test_data.json

@@ -27,6 +27,10 @@
                     "City": {
                         "Type": "String",
                         "Value": "Any City"
+                    },
+                    "traceparent": {
+                        "Type": "String",
+                        "Value": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00"
                     }
                 }
             }

tests/contrib/serverless/aws_tests.py

@@ -344,6 +344,8 @@ def test_func(event, context):
     assert transaction["span_count"]["started"] == 1
     assert transaction["context"]["message"]["headers"]["Population"] == "1250800"
     assert transaction["context"]["message"]["headers"]["City"] == "Any City"
+    assert len(transaction["links"]) == 1
+    assert transaction["links"][0] == {"trace_id": "0af7651916cd43dd8448eb211c80319c", "span_id": "b7ad6b7169203331"}
 
 
 def test_capture_serverless_sqs(event_sqs, context, elasticapm_client):