feat: lambda support for DSM

datadog_lambda/tracing.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/wrapper.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)

pyproject.toml

@@ -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 }