feat: lambda support for DSM (#622)

* removed current dsm implementation

* new dsm lambda implementation

* check env variable in proper way

* add tests

* add more detailed comment to function

* fixed lint for tests

* remove not needed tests

* fix

* error handling

* fix

* some fixes

* renamed to extract_context_with_datastreams

* changed to explicit check of dsm context

* move dsm tests to test_tracing.py

* add wanted tests

* caught sns -> sqs bug

* revert back to original tracing.py implementation

* fix lint

* revert spacing stuff to original

* remove unneccessary checks, still set checkpoints even when dsm context fails to be propagated

* remove not needed comment

* fixes

* lambda functions not allowed by lint

* use lambda function, add checks before checkpoint

* remove unneccesary import

* move if statement with least work first

* changed function name to original, arn exception handle w test, return None instead of {}

* some fixes

* remove comments that are not needed

* fix

* fix

* extra check

* remove unneccesary work associated with event_source

* fix lint

* add tests for empty arn logic

* more descriptive name

* fix lint

* fix

* moved arn check to checkpoint, remove comments, add variable dec

* kinesis fix, tests fix

* remove not needed test

* formatting fix

* made tests clearer and more comprehensive

* fix

* made comments more obvious on point of tests

* update toml file

* edit build make more sense

* made sure tests match with table

* fix

* removed not needed comments

* addc comment

* add some more kinesis tests

* add logger check in exception test cases

* fix comment

* fix

* updated tests to use inner checkpoint check, keep build at <4

* fix

* change to 3.10.2

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -67,6 +67,24 @@
 LOWER_64_BITS = "LOWER_64_BITS"
 
 
+def _dsm_set_checkpoint(context_json, event_type, arn):
+    if not config.data_streams_enabled:
+        return
+
+    if not arn:
+        return
+
+    try:
+        from ddtrace.data_streams import set_consume_checkpoint
+
+        carrier_get = lambda k: context_json and context_json.get(k)  # noqa: E731
+        set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False)
+    except Exception as e:
+        logger.debug(
+            f"DSM:Failed to set consume checkpoint for {event_type} {arn}: {e}"
+        )
+
+
 def _convert_xray_trace_id(xray_trace_id):
     """
     Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
@@ -202,7 +220,9 @@ def create_sns_event(message):
     }
 
 
-def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
+def extract_context_from_sqs_or_sns_event_or_context(
+    event, lambda_context, event_source
+):
     """
     Extract Datadog trace context from an SQS event.
 
@@ -214,7 +234,10 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     Lambda Context.
 
     Falls back to lambda context if no trace data is found in the SQS message attributes.
+    Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
+    source_arn = ""
+    event_type = "sqs" if event_source.equals(EventTypes.SQS) else "sns"
 
     # EventBridge => SQS
     try:
@@ -226,6 +249,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
 
     try:
         first_record = event.get("Records")[0]
+        source_arn = first_record.get("eventSourceARN", "")
 
         # logic to deal with SNS => SQS event
         if "body" in first_record:
@@ -241,6 +265,9 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
         msg_attributes = first_record.get("messageAttributes")
         if msg_attributes is None:
             sns_record = first_record.get("Sns") or {}
+            # SNS->SQS event would extract SNS arn without this check
+            if event_source.equals(EventTypes.SNS):
+                source_arn = sns_record.get("TopicArn", "")
             msg_attributes = sns_record.get("MessageAttributes") or {}
         dd_payload = msg_attributes.get("_datadog")
         if dd_payload:
@@ -272,8 +299,9 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                         logger.debug(
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
-
-                return propagator.extract(dd_data)
+                context = propagator.extract(dd_data)
+                _dsm_set_checkpoint(dd_data, event_type, source_arn)
+                return context
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
             # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
@@ -296,9 +324,13 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                             span_id=int(x_ray_context["parent_id"], 16),
                             sampling_priority=float(x_ray_context["sampled"]),
                         )
+        # Still want to set a DSM checkpoint even if DSM context not propagated
+        _dsm_set_checkpoint(None, event_type, source_arn)
         return extract_context_from_lambda_context(lambda_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
+        # Still want to set a DSM checkpoint even if DSM context not propagated
+        _dsm_set_checkpoint(None, event_type, source_arn)
         return extract_context_from_lambda_context(lambda_context)
 
 
@@ -357,9 +389,12 @@ def extract_context_from_eventbridge_event(event, lambda_context):
 def extract_context_from_kinesis_event(event, lambda_context):
     """
     Extract datadog trace context from a Kinesis Stream's base64 encoded data string
+    Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
+    source_arn = ""
     try:
         record = get_first_record(event)
+        source_arn = record.get("eventSourceARN", "")
         kinesis = record.get("kinesis")
         if not kinesis:
             return extract_context_from_lambda_context(lambda_context)
@@ -373,10 +408,13 @@ def extract_context_from_kinesis_event(event, lambda_context):
             data_obj = json.loads(data_str)
             dd_ctx = data_obj.get("_datadog")
             if dd_ctx:
-                return propagator.extract(dd_ctx)
+                context = propagator.extract(dd_ctx)
+                _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn)
+                return context
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
-
+    # Still want to set a DSM checkpoint even if DSM context not propagated
+    _dsm_set_checkpoint(None, "kinesis", source_arn)
     return extract_context_from_lambda_context(lambda_context)
 
 
@@ -594,7 +632,7 @@ def extract_dd_trace_context(
         )
     elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS):
         context = extract_context_from_sqs_or_sns_event_or_context(
-            event, lambda_context
+            event, lambda_context, event_source
         )
     elif event_source.equals(EventTypes.EVENTBRIDGE):
         context = extract_context_from_eventbridge_event(event, lambda_context)

datadog_lambda/tracing.py

@@ -10,7 +10,6 @@
 from time import time_ns
 
 from datadog_lambda.asm import asm_set_context, asm_start_response, asm_start_request
-from datadog_lambda.dsm import set_dsm_context
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -237,8 +236,6 @@ def _before(self, event, context):
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, config.decode_authorizer_context
                     )
-                if config.data_streams_enabled:
-                    set_dsm_context(event, event_source)
 
                 if config.appsec_enabled:
                     asm_set_context(event_source)

datadog_lambda/wrapper.py

@@ -28,7 +28,7 @@ classifiers = [
 python = ">=3.8.0,<4"
 datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
-ddtrace = ">=2.20.0,<4"
+ddtrace = ">=3.10.2,<4"
 ujson = ">=5.9.0"
 botocore = { version = "^1.34.0", optional = true }
 requests = { version ="^2.22.0", optional = true }