add span links from SQS messages in AWS lambda transactions (#1662)

* add span links from SQS messages in AWS lambda transactions

closes #1473

* update changelog

Co-authored-by: Colton Myers <[email protected]>

Author: beniwohli

CHANGELOG.asciidoc

@@ -39,6 +39,7 @@ endif::[]
  * Add support for instrumenting the Elasticsearch 8 Python client {pull}1642[#1642]
  * Add `*principal*` to default `sanitize_field_names` configuration {pull}1664[#1664]
  * Add docs and better support for custom metrics, including in AWS Lambda {pull}1643[#1643]
+ * Add support for capturing span links from AWS SQS in AWS Lambda {pull}1662[#1662]
 
 [float]
 ===== Bug fixes

elasticapm/base.py

@@ -302,6 +302,7 @@ def begin_transaction(
         :param trace_parent: an optional TraceParent object for distributed tracing
         :param start: override the start timestamp, mostly useful for testing
         :param auto_activate: whether to set this transaction in execution_context
+        :param links: a sequence of traceparent objects to causally link this transaction with
         :return: the started transaction object
         """
         if self.config.is_recording:

elasticapm/contrib/serverless/aws.py

@@ -184,7 +184,16 @@ def __enter__(self):
                 transaction_type = "messaging"
                 transaction_name = "RECEIVE {}".format(record["eventSourceARN"].split(":")[5])
 
-        self.transaction = self.client.begin_transaction(transaction_type, trace_parent=trace_parent)
+        if "Records" in self.event:
+            links = [
+                TraceParent.from_string(record["messageAttributes"]["traceparent"]["stringValue"])
+                for record in self.event["Records"][:1000]
+                if "messageAttributes" in record and "traceparent" in record["messageAttributes"]
+            ]
+        else:
+            links = []
+
+        self.transaction = self.client.begin_transaction(transaction_type, trace_parent=trace_parent, links=links)
         elasticapm.set_transaction_name(transaction_name, override=False)
         if self.source in SERVERLESS_HTTP_REQUEST:
             elasticapm.set_context(

tests/contrib/serverless/aws_sqs_test_data.json

@@ -28,6 +28,12 @@
                     "stringListValues": [],
                     "binaryListValues": [],
                     "dataType": "String"
+                },
+                "traceparent": {
+                    "stringValue": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00",
+                    "stringListValues": [],
+                    "binaryListValues": [],
+                    "dataType": "String"
                 }
             },
             "md5OfBody": "5eb63bbbe01eeed093cb22bb8f5acdc3",
@@ -36,4 +42,4 @@
             "awsRegion": "us-east-1"
         }
     ]
-}
+}

tests/contrib/serverless/aws_tests.py

@@ -324,6 +324,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_s3_batch(event_s3_batch, context, elasticapm_client):