[SVLS-5265] S3 Event Handler Span Pointers

Author: apiarian-datadog

datadog_lambda/span_pointers.py

@@ -0,0 +1,74 @@
+from itertools import chain
+from ddtrace._trace.utils_botocore.span_pointers import _aws_s3_object_span_pointer_description
+from ddtrace._trace._span_pointer import _SpanPointerDirection
+from ddtrace._trace._span_pointer import _SpanPointerDescription
+from datadog_lambda.trigger import EventTypes
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_span_pointers(
+    event_source,
+    event,
+) -> list[_SpanPointerDescription]:
+    if event_source.equals(EventTypes.S3):
+        return _calculate_s3_span_pointers_for_event(event)
+
+    return []
+
+
+def _calculate_s3_span_pointers_for_event(event) -> list[_SpanPointerDescription]:
+    # Example event:
+    # https://docs.aws.amazon.com/lambda/latest/dg/with-s3.html
+
+    return list(
+        chain.from_iterable(
+            _calculate_s3_span_pointers_for_event_record(record)
+            for record in event.get("Records", [])
+        )
+    )
+
+
+def _calculate_s3_span_pointers_for_event_record(record) -> list[_SpanPointerDescription]:
+    # Event types:
+    # https://docs.aws.amazon.com/AmazonS3/latest/userguide/notification-how-to-event-types-and-destinations.html
+
+    if record.get("eventName").startswith("ObjectCreated:"):
+        s3_information = record.get("s3", None)
+        if s3_information is not None:
+            return _calculate_s3_span_pointers_for_object_created_s3_information(s3_information)
+
+    return []
+
+
+def _calculate_s3_span_pointers_for_object_created_s3_information(s3_information) -> list[_SpanPointerDescription]:
+    try:
+        bucket = s3_information["bucket"]["name"]
+        key = s3_information["object"]["key"]
+        etag = s3_information["object"]["eTag"]
+
+    except KeyError as e:
+        logger.warning(
+            "missing s3 information required to make a span pointer: %s",
+            str(e),
+        )
+        return []
+
+    try:
+        return [
+            _aws_s3_object_span_pointer_description(
+                pointer_direction=_SpanPointerDirection.UPSTREAM,
+                bucket=bucket,
+                key=key,
+                etag=etag,
+            )
+        ]
+
+    except Exception as e:
+        logger.warning(
+            "failed to generate S3 span pointer: %s",
+            str(e),
+        )
+        return []

datadog_lambda/tracing.py

@@ -1259,6 +1259,7 @@ def create_function_execution_span(
     merge_xray_traces,
     trigger_tags,
     parent_span=None,
+    span_pointers=None,
 ):
     tags = None
     if context:
@@ -1296,6 +1297,14 @@ def create_function_execution_span(
         span.set_tags(tags)
     if parent_span:
         span.parent_id = parent_span.span_id
+    if span_pointers:
+        for span_pointer_description in span_pointers:
+            span._add_span_pointer(
+                pointer_kind=span_pointer_description.pointer_kind,
+                pointer_direction=span_pointer_description.pointer_direction,
+                pointer_hash=span_pointer_description.pointer_hash,
+                extra_attributes=span_pointer_description.extra_attributes,
+            )
     return span
 
 

datadog_lambda/wrapper.py

@@ -30,6 +30,7 @@
 )
 from datadog_lambda.module_name import modify_module_name
 from datadog_lambda.patch import patch_all
+from datadog_lambda.span_pointers import calculate_span_pointers
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
     create_dd_dummy_metadata_subsegment,
@@ -315,6 +316,7 @@ def _before(self, event, context):
                     self.merge_xray_traces,
                     self.trigger_tags,
                     parent_span=self.inferred_span,
+                    span_pointers=calculate_span_pointers(event_source, event),
                 )
             else:
                 set_correlation_ids()