datadog_lambda 6.111.0__py3-none-any.whl → 7.112.0__py3-none-any.whl

datadog_lambda/__init__.py

@@ -17,3 +17,10 @@ from datadog_lambda.logger import initialize_logging  # noqa: E402
 
 
 initialize_logging(__name__)
+
+
+from datadog_lambda.patch import patch_all  # noqa: E402
+
+# Patch third-party libraries for tracing, must be done before importing any
+# handler code.
+patch_all()

datadog_lambda/asm.py

@@ -0,0 +1,184 @@
+from copy import deepcopy
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from ddtrace.contrib.internal.trace_utils import _get_request_header_client_ip
+from ddtrace.internal import core
+from ddtrace.trace import Span
+
+from datadog_lambda.trigger import (
+    EventSubtypes,
+    EventTypes,
+    _EventSource,
+    _http_event_types,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _to_single_value_headers(headers: Dict[str, List[str]]) -> Dict[str, str]:
+    """
+    Convert multi-value headers to single-value headers.
+    If a header has multiple values, join them with commas.
+    """
+    single_value_headers = {}
+    for key, values in headers.items():
+        single_value_headers[key] = ", ".join(values)
+    return single_value_headers
+
+
+def _merge_single_and_multi_value_headers(
+    single_value_headers: Dict[str, str],
+    multi_value_headers: Dict[str, List[str]],
+):
+    """
+    Merge single-value headers with multi-value headers.
+    If a header exists in both, we merge them removing duplicates
+    """
+    merged_headers = deepcopy(multi_value_headers)
+    for key, value in single_value_headers.items():
+        if key not in merged_headers:
+            merged_headers[key] = [value]
+        elif value not in merged_headers[key]:
+            merged_headers[key].append(value)
+    return _to_single_value_headers(merged_headers)
+
+
+def asm_set_context(event_source: _EventSource):
+    """Add asm specific items to the ExecutionContext.
+
+    This allows the AppSecSpanProcessor to know information about the event
+    at the moment the span is created and skip it when not relevant.
+    """
+    if event_source.event_type not in _http_event_types:
+        core.set_item("appsec_skip_next_lambda_event", True)
+
+
+def asm_start_request(
+    span: Span,
+    event: Dict[str, Any],
+    event_source: _EventSource,
+    trigger_tags: Dict[str, str],
+):
+    if event_source.event_type not in _http_event_types:
+        return
+
+    request_headers: Dict[str, str] = {}
+    peer_ip: Optional[str] = None
+    request_path_parameters: Optional[Dict[str, Any]] = None
+    route: Optional[str] = None
+
+    if event_source.event_type == EventTypes.ALB:
+        headers = event.get("headers")
+        multi_value_request_headers = event.get("multiValueHeaders")
+        if multi_value_request_headers:
+            request_headers = _to_single_value_headers(multi_value_request_headers)
+        else:
+            request_headers = headers or {}
+
+        raw_uri = event.get("path")
+        parsed_query = event.get("multiValueQueryStringParameters") or event.get(
+            "queryStringParameters"
+        )
+
+    elif event_source.event_type == EventTypes.LAMBDA_FUNCTION_URL:
+        request_headers = event.get("headers", {})
+        peer_ip = event.get("requestContext", {}).get("http", {}).get("sourceIp")
+        raw_uri = event.get("rawPath")
+        parsed_query = event.get("queryStringParameters")
+
+    elif event_source.event_type == EventTypes.API_GATEWAY:
+        request_context = event.get("requestContext", {})
+        request_path_parameters = event.get("pathParameters")
+        route = trigger_tags.get("http.route")
+
+        if event_source.subtype == EventSubtypes.API_GATEWAY:
+            request_headers = event.get("headers", {})
+            peer_ip = request_context.get("identity", {}).get("sourceIp")
+            raw_uri = event.get("path")
+            parsed_query = event.get("multiValueQueryStringParameters")
+
+        elif event_source.subtype == EventSubtypes.HTTP_API:
+            request_headers = event.get("headers", {})
+            peer_ip = request_context.get("http", {}).get("sourceIp")
+            raw_uri = event.get("rawPath")
+            parsed_query = event.get("queryStringParameters")
+
+        elif event_source.subtype == EventSubtypes.WEBSOCKET:
+            request_headers = _to_single_value_headers(
+                event.get("multiValueHeaders", {})
+            )
+            peer_ip = request_context.get("identity", {}).get("sourceIp")
+            raw_uri = event.get("path")
+            parsed_query = event.get("multiValueQueryStringParameters")
+
+        else:
+            return
+
+    else:
+        return
+
+    body = event.get("body")
+    is_base64_encoded = event.get("isBase64Encoded", False)
+
+    request_ip = _get_request_header_client_ip(request_headers, peer_ip, True)
+    if request_ip is not None:
+        span.set_tag_str("http.client_ip", request_ip)
+        span.set_tag_str("network.client.ip", request_ip)
+
+    core.dispatch(
+        # The matching listener is registered in ddtrace.appsec._handlers
+        "aws_lambda.start_request",
+        (
+            span,
+            request_headers,
+            request_ip,
+            body,
+            is_base64_encoded,
+            raw_uri,
+            route,
+            trigger_tags.get("http.method"),
+            parsed_query,
+            request_path_parameters,
+        ),
+    )
+
+
+def asm_start_response(
+    span: Span,
+    status_code: str,
+    event_source: _EventSource,
+    response: Union[Dict[str, Any], str, None],
+):
+    if event_source.event_type not in _http_event_types:
+        return
+
+    if isinstance(response, dict) and (
+        "headers" in response or "multiValueHeaders" in response
+    ):
+        headers = response.get("headers", {})
+        multi_value_request_headers = response.get("multiValueHeaders")
+        if isinstance(multi_value_request_headers, dict) and isinstance(headers, dict):
+            response_headers = _merge_single_and_multi_value_headers(
+                headers, multi_value_request_headers
+            )
+        elif isinstance(headers, dict):
+            response_headers = headers
+        else:
+            response_headers = {
+                "content-type": "application/json",
+            }
+    else:
+        response_headers = {
+            "content-type": "application/json",
+        }
+
+    core.dispatch(
+        # The matching listener is registered in ddtrace.appsec._handlers
+        "aws_lambda.start_response",
+        (
+            span,
+            status_code,
+            response_headers,
+        ),
+    )

datadog_lambda/config.py

@@ -95,6 +95,7 @@ class Config:
     data_streams_enabled = _get_env(
         "DD_DATA_STREAMS_ENABLED", "false", as_bool, depends_on_tracing=True
     )
+    appsec_enabled = _get_env("DD_APPSEC_ENABLED", "false", as_bool)
 
     is_gov_region = _get_env("AWS_REGION", "", lambda x: x.startswith("us-gov-"))
 

datadog_lambda/tracing.py

@@ -67,6 +67,24 @@ HIGHER_64_BITS = "HIGHER_64_BITS"
 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)
@@ -859,7 +897,7 @@ def create_inferred_span_from_lambda_function_url_event(event, context):
     InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync")
     if span:
         span.set_tags(tags)
-        span.start_ns = int(request_time_epoch) * 1e6
+        span.start_ns = int(request_time_epoch * 1e6)
     return span
 
 

datadog_lambda/version.py

@@ -1 +1 @@
-__version__ = "6.111.0"
+__version__ = "7.112.0"

datadog_lambda/wrapper.py

@@ -9,7 +9,7 @@ import ujson as json
 from importlib import import_module
 from time import time_ns
 
-from datadog_lambda.dsm import set_dsm_context
+from datadog_lambda.asm import asm_set_context, asm_start_response, asm_start_request
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -25,7 +25,6 @@ from datadog_lambda.constants import (
     Headers,
 )
 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.tag_object import tag_object
 from datadog_lambda.tracing import (
@@ -142,8 +141,6 @@ class _LambdaDecorator(object):
             os.environ[DD_REQUESTS_SERVICE_NAME] = os.environ.get(
                 DD_SERVICE, "aws.lambda"
             )
-            # Patch third-party libraries for tracing
-            patch_all()
 
             # Enable LLM Observability
             if config.llmobs_enabled:
@@ -165,10 +162,9 @@ class _LambdaDecorator(object):
             self.response = self.func(event, context, **kwargs)
             return self.response
         except Exception:
-            if not should_use_extension:
-                from datadog_lambda.metric import submit_errors_metric
+            from datadog_lambda.metric import submit_errors_metric
 
-                submit_errors_metric(context)
+            submit_errors_metric(context)
 
             if self.span:
                 self.span.set_traceback()
@@ -240,8 +236,10 @@ class _LambdaDecorator(object):
                     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)
+
                 self.span = create_function_execution_span(
                     context=context,
                     function_name=config.function_name,
@@ -253,6 +251,8 @@ class _LambdaDecorator(object):
                     parent_span=self.inferred_span,
                     span_pointers=calculate_span_pointers(event_source, event),
                 )
+                if config.appsec_enabled:
+                    asm_start_request(self.span, event, event_source, self.trigger_tags)
             else:
                 set_correlation_ids()
             if config.profiling_enabled and is_new_sandbox():
@@ -285,6 +285,15 @@ class _LambdaDecorator(object):
 
                 if status_code:
                     self.span.set_tag("http.status_code", status_code)
+
+                if config.appsec_enabled:
+                    asm_start_response(
+                        self.span,
+                        status_code,
+                        self.event_source,
+                        response=self.response,
+                    )
+
                 self.span.finish()
 
             if self.inferred_span:

{datadog_lambda-6.111.0.dist-info → datadog_lambda-7.112.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datadog_lambda
-Version: 6.111.0
+Version: 7.112.0
 Summary: The Datadog AWS Lambda Library
 Home-page: https://github.com/DataDog/datadog-lambda-python
 License: Apache-2.0
@@ -19,7 +19,7 @@ Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: dev
 Requires-Dist: botocore (>=1.34.0,<2.0.0) ; extra == "dev"
 Requires-Dist: datadog (>=0.51.0,<1.0.0)
-Requires-Dist: ddtrace (>=2.20.0,<4)
+Requires-Dist: ddtrace (>=3.10.2,<4)
 Requires-Dist: flake8 (>=5.0.4,<6.0.0) ; extra == "dev"
 Requires-Dist: pytest (>=8.0.0,<9.0.0) ; extra == "dev"
 Requires-Dist: pytest-benchmark (>=4.0,<5.0) ; extra == "dev"

{datadog_lambda-6.111.0.dist-info → datadog_lambda-7.112.0.dist-info}/RECORD

@@ -1,10 +1,10 @@
-datadog_lambda/__init__.py,sha256=F_KG5XbNq9QEzbioec7DXzyTv3oYJHxeEgXE88dxXRQ,637
+datadog_lambda/__init__.py,sha256=uAqki08Fd26eUXZtoptJX68kTFhK-ZkCyLVCtE0MS4k,802
 datadog_lambda/api.py,sha256=FqQ9vjNSxD5BxjeHDDKbim_NIN08dshW-52wNzqw5yo,5249
+datadog_lambda/asm.py,sha256=gX4ySlmLCGnOWRf0LpMJxDElpQ2su5LnphJJy4Lrc-Y,6115
 datadog_lambda/cold_start.py,sha256=NfE4EtRUHuc2ub26U8sAwMVEDNINbw1OTh2MAklQGyk,8032
-datadog_lambda/config.py,sha256=sWcnhpqJF1pWIjgnxxdPOfJWDQT11ZOsKI8pC6TyrJs,5191
+datadog_lambda/config.py,sha256=QoEZrSL_zSHY5lV9RznAwPONEZe0i3mCWp2jPoO8aEw,5260
 datadog_lambda/constants.py,sha256=0y6O9s_8RLflYR507SDMQjKmYY16tr1yi2KQuuF1GaY,1696
 datadog_lambda/dogstatsd.py,sha256=-2HiU3xvV_beXh8rI9oIXxSUmVcamI0KqAl7gc9LkYA,5216
-datadog_lambda/dsm.py,sha256=TwRGl7bpWMpkv1YJtRsjsCma4QTI3xKrpOvj2GmMO_A,1219
 datadog_lambda/extension.py,sha256=ZU64QpA2K9K9C0jfqusBgpiWQe0QA2dcJCNk7UgjVfw,621
 datadog_lambda/handler.py,sha256=VhdNUb_lpPRJ73h-fdDxxqSm0Ic_lkt4x_9Rt0amVV4,1312
 datadog_lambda/logger.py,sha256=nGxNMouF7wcjmoPsgivzzjNLvSy3WbGtKElxOvITZDg,766
@@ -17,14 +17,14 @@ datadog_lambda/statsd_writer.py,sha256=wUNG8phIreNOGCc5VjYZV-FNogvKfiz9kN33QldqZ
 datadog_lambda/tag_object.py,sha256=NNvMVi07kNYRBMPpB7Lk7_-z5ccmxo2UmTWnvWVPhiM,2089
 datadog_lambda/tags.py,sha256=wy6uH8eAGMn7cfZEdHpL9uEGoM85bVyyXhYwSQtfHHc,2532
 datadog_lambda/thread_stats_writer.py,sha256=VfD6G65zNHDA0nABhGOO0it7yqvn-p6ereQ329DA7r8,2894
-datadog_lambda/tracing.py,sha256=4rKF_m4ou26DYDU2nS7TS6zcjBIdn04nVLx6xSQblRc,54301
+datadog_lambda/tracing.py,sha256=aalT3t_ISPvPZIZzkrSK6OnW1ixY9ijsefjUQQEKQk0,56071
 datadog_lambda/trigger.py,sha256=7EhNoCtJbnuYaqCvv1i2tAqZxfXx7ccsnMkwnWvmd7w,14212
-datadog_lambda/version.py,sha256=FwGzine_pPPhRX5S_HqlSuRRxSU07s3Dg_-ioiPf974,24
-datadog_lambda/wrapper.py,sha256=SCt9yeWnn9N0QLRlkjsjWSh-7mdhyC6Spsdg4Ehzx08,13119
+datadog_lambda/version.py,sha256=HDtiDcfxQBWPKo0xjxP1Hvuj0QipaSpVwhvqFrmieB0,24
+datadog_lambda/wrapper.py,sha256=siHNfNSLb3iIwN0vRnW67CtnOirqogbFG54zjWPVLAU,13377
 datadog_lambda/xray.py,sha256=jvA4Fk76PLMgsjUoUZ7gp2otv53hFt39Nvso1ZNaivg,3749
-datadog_lambda-6.111.0.dist-info/LICENSE,sha256=4yQmjpKp1MKL7DdRDPVHkKYc2W0aezm5SIDske8oAdM,11379
-datadog_lambda-6.111.0.dist-info/LICENSE-3rdparty.csv,sha256=9CDAR1GKawwTbZkqt1RP0uwEcaRM3RhOeTB5tWXr8Ts,1381
-datadog_lambda-6.111.0.dist-info/METADATA,sha256=INJYjzOzOw9SdxlS9l6JCmBAC1kcR_MiGKAh0QRPW6I,7727
-datadog_lambda-6.111.0.dist-info/NOTICE,sha256=Jue-d8mQ1ENIHDZdYc2-X8mVYtScXb8pzF1pTLN-kRc,141
-datadog_lambda-6.111.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-datadog_lambda-6.111.0.dist-info/RECORD,,
+datadog_lambda-7.112.0.dist-info/LICENSE,sha256=4yQmjpKp1MKL7DdRDPVHkKYc2W0aezm5SIDske8oAdM,11379
+datadog_lambda-7.112.0.dist-info/LICENSE-3rdparty.csv,sha256=9CDAR1GKawwTbZkqt1RP0uwEcaRM3RhOeTB5tWXr8Ts,1381
+datadog_lambda-7.112.0.dist-info/METADATA,sha256=YzLOAslZUkBtDAPKfgPyYnpbXKUwd92D1zZbctuFBR4,7727
+datadog_lambda-7.112.0.dist-info/NOTICE,sha256=Jue-d8mQ1ENIHDZdYc2-X8mVYtScXb8pzF1pTLN-kRc,141
+datadog_lambda-7.112.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+datadog_lambda-7.112.0.dist-info/RECORD,,

datadog_lambda/dsm.py

@@ -1,38 +0,0 @@
-from datadog_lambda import logger
-from datadog_lambda.trigger import EventTypes
-
-
-def set_dsm_context(event, event_source):
-
-    if event_source.equals(EventTypes.SQS):
-        _dsm_set_sqs_context(event)
-
-
-def _dsm_set_sqs_context(event):
-    from datadog_lambda.wrapper import format_err_with_traceback
-    from ddtrace.internal.datastreams import data_streams_processor
-    from ddtrace.internal.datastreams.processor import DsmPathwayCodec
-    from ddtrace.internal.datastreams.botocore import (
-        get_datastreams_context,
-        calculate_sqs_payload_size,
-    )
-
-    records = event.get("Records")
-    if records is None:
-        return
-    processor = data_streams_processor()
-
-    for record in records:
-        try:
-            queue_arn = record.get("eventSourceARN", "")
-
-            contextjson = get_datastreams_context(record)
-            payload_size = calculate_sqs_payload_size(record)
-
-            ctx = DsmPathwayCodec.decode(contextjson, processor)
-            ctx.set_checkpoint(
-                ["direction:in", f"topic:{queue_arn}", "type:sqs"],
-                payload_size=payload_size,
-            )
-        except Exception as e:
-            logger.error(format_err_with_traceback(e))