datadog_lambda 6.106.0__tar.gz → 6.108.0__tar.gz

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datadog_lambda
-Version: 6.106.0
+Version: 6.108.0
 Summary: The Datadog AWS Lambda Library
 Home-page: https://github.com/DataDog/datadog-lambda-python
 License: Apache-2.0
@@ -17,9 +17,9 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: dev
-Requires-Dist: boto3 (>=1.34.0,<2.0.0) ; 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)
+Requires-Dist: ddtrace (>=2.20.0,<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.108.0/datadog_lambda/api.py

@@ -0,0 +1,145 @@
+import logging
+import os
+
+from datadog_lambda.fips import fips_mode_enabled
+
+logger = logging.getLogger(__name__)
+KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName"
+api_key = None
+
+
+def decrypt_kms_api_key(kms_client, ciphertext):
+    import base64
+
+    from botocore.exceptions import ClientError
+
+    """
+    Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS.
+    For this to work properly, the Lambda function must have the appropriate IAM permissions.
+
+    Args:
+        kms_client: The KMS client to use for decryption
+        ciphertext (string): The base64-encoded ciphertext to decrypt
+    """
+    decoded_bytes = base64.b64decode(ciphertext)
+
+    """
+    When the API key is encrypted using the AWS console, the function name is added as an
+    encryption context. When the API key is encrypted using the AWS CLI, no encryption context
+    is added. We need to try decrypting the API key both with and without the encryption context.
+    """
+    # Try without encryption context, in case API key was encrypted using the AWS CLI
+    function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME")
+    try:
+        plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)[
+            "Plaintext"
+        ].decode("utf-8")
+    except ClientError:
+        logger.debug(
+            "Failed to decrypt ciphertext without encryption context, \
+            retrying with encryption context"
+        )
+        # Try with encryption context, in case API key was encrypted using the AWS Console
+        plaintext = kms_client.decrypt(
+            CiphertextBlob=decoded_bytes,
+            EncryptionContext={
+                KMS_ENCRYPTION_CONTEXT_KEY: function_name,
+            },
+        )["Plaintext"].decode("utf-8")
+
+    return plaintext
+
+
+def get_api_key() -> str:
+    """
+    Gets the Datadog API key from the environment variables or secrets manager.
+    Extracts the result to a global value to avoid repeated calls to the
+    secrets manager from different products.
+    """
+    global api_key
+    if api_key:
+        return api_key
+
+    DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "")
+    DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "")
+    DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "")
+    DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", ""))
+
+    LAMBDA_REGION = os.environ.get("AWS_REGION", "")
+    if fips_mode_enabled:
+        logger.debug(
+            "FIPS mode is enabled, using FIPS endpoints for secrets management."
+        )
+
+    if DD_API_KEY_SECRET_ARN:
+        # Secrets manager endpoints: https://docs.aws.amazon.com/general/latest/gr/asm.html
+        try:
+            secrets_region = DD_API_KEY_SECRET_ARN.split(":")[3]
+        except Exception:
+            logger.debug(
+                "Invalid secret arn in DD_API_KEY_SECRET_ARN. Unable to get API key."
+            )
+            return ""
+        endpoint_url = (
+            f"https://secretsmanager-fips.{secrets_region}.amazonaws.com"
+            if fips_mode_enabled
+            else None
+        )
+        secrets_manager_client = _boto3_client(
+            "secretsmanager", endpoint_url=endpoint_url, region_name=secrets_region
+        )
+        api_key = secrets_manager_client.get_secret_value(
+            SecretId=DD_API_KEY_SECRET_ARN
+        )["SecretString"]
+    elif DD_API_KEY_SSM_NAME:
+        # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html
+        fips_endpoint = (
+            f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com"
+            if fips_mode_enabled
+            else None
+        )
+        ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint)
+        api_key = ssm_client.get_parameter(
+            Name=DD_API_KEY_SSM_NAME, WithDecryption=True
+        )["Parameter"]["Value"]
+    elif DD_KMS_API_KEY:
+        # KMS endpoints: https://docs.aws.amazon.com/general/latest/gr/kms.html
+        fips_endpoint = (
+            f"https://kms-fips.{LAMBDA_REGION}.amazonaws.com"
+            if fips_mode_enabled
+            else None
+        )
+        kms_client = _boto3_client("kms", endpoint_url=fips_endpoint)
+        api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY)
+    else:
+        api_key = DD_API_KEY
+
+    return api_key
+
+
+def init_api():
+    if not os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
+        # Make sure that this package would always be lazy-loaded/outside from the critical path
+        # since underlying packages are quite heavy to load
+        # and useless with the extension unless sending metrics with timestamps
+        from datadog import api
+
+        if not api._api_key:
+            api._api_key = get_api_key()
+
+        logger.debug("Setting DATADOG_API_KEY of length %d", len(api._api_key))
+
+        # Set DATADOG_HOST, to send data to a non-default Datadog datacenter
+        api._api_host = os.environ.get(
+            "DATADOG_HOST", "https://api." + os.environ.get("DD_SITE", "datadoghq.com")
+        )
+        logger.debug("Setting DATADOG_HOST to %s", api._api_host)
+
+        # Unmute exceptions from datadog api client, so we can catch and handle them
+        api._mute = False
+
+
+def _boto3_client(*args, **kwargs):
+    import botocore.session
+
+    return botocore.session.get_session().create_client(*args, **kwargs)

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/dogstatsd.py

@@ -1,11 +1,10 @@
+import errno
 import logging
 import os
-import socket
-import errno
 import re
+import socket
 from threading import Lock
 
-
 MIN_SEND_BUFFER_SIZE = 32 * 1024
 log = logging.getLogger("datadog_lambda.dogstatsd")
 
@@ -55,14 +54,21 @@ class DogStatsd(object):
 
         return sock
 
-    def distribution(self, metric, value, tags=None):
+    def distribution(self, metric, value, tags=None, timestamp=None):
         """
-        Send a global distribution value, optionally setting tags.
+        Send a global distribution value, optionally setting tags. The optional
+        timestamp should be an integer representing seconds since the epoch
+        (January 1, 1970, 00:00:00 UTC).
 
         >>> statsd.distribution("uploaded.file.size", 1445)
         >>> statsd.distribution("album.photo.count", 26, tags=["gender:female"])
+        >>> statsd.distribution(
+        >>>     "historic.file.count",
+        >>>     5,
+        >>>     timestamp=int(datetime(2020, 2, 14, 12, 0, 0).timestamp()),
+        >>> )
         """
-        self._report(metric, "d", value, tags)
+        self._report(metric, "d", value, tags, timestamp)
 
     def close_socket(self):
         """
@@ -84,20 +90,21 @@ class DogStatsd(object):
             for tag in tag_list
         ]
 
-    def _serialize_metric(self, metric, metric_type, value, tags):
+    def _serialize_metric(self, metric, metric_type, value, tags, timestamp):
         # Create/format the metric packet
-        return "%s:%s|%s%s" % (
+        return "%s:%s|%s%s%s" % (
             metric,
             value,
             metric_type,
             ("|#" + ",".join(self.normalize_tags(tags))) if tags else "",
+            ("|T" + str(timestamp)) if timestamp is not None else "",
         )
 
-    def _report(self, metric, metric_type, value, tags):
+    def _report(self, metric, metric_type, value, tags, timestamp):
         if value is None:
             return
 
-        payload = self._serialize_metric(metric, metric_type, value, tags)
+        payload = self._serialize_metric(metric, metric_type, value, tags, timestamp)
 
         # Send it
         self._send_to_server(payload)

datadog_lambda-6.108.0/datadog_lambda/fips.py

@@ -0,0 +1,19 @@
+import logging
+import os
+
+is_gov_region = os.environ.get("AWS_REGION", "").startswith("us-gov-")
+
+fips_mode_enabled = (
+    os.environ.get(
+        "DD_LAMBDA_FIPS_MODE",
+        "true" if is_gov_region else "false",
+    ).lower()
+    == "true"
+)
+
+if is_gov_region or fips_mode_enabled:
+    logger = logging.getLogger(__name__)
+    logger.debug(
+        "Python Lambda Layer FIPS mode is %s.",
+        "enabled" if fips_mode_enabled else "not enabled",
+    )

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/handler.py

@@ -3,7 +3,6 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2020 Datadog, Inc.
 
-from __future__ import absolute_import
 from importlib import import_module
 
 import os

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/metric.py

@@ -3,37 +3,66 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+import enum
+import logging
 import os
 import time
-import logging
-import ujson as json
 from datetime import datetime, timedelta
 
+import ujson as json
+
 from datadog_lambda.extension import should_use_extension
-from datadog_lambda.tags import get_enhanced_metrics_tags, dd_lambda_layer_tag
+from datadog_lambda.fips import fips_mode_enabled
+from datadog_lambda.tags import dd_lambda_layer_tag, get_enhanced_metrics_tags
 
 logger = logging.getLogger(__name__)
 
-lambda_stats = None
-extension_thread_stats = None
 
-flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
+class MetricsHandler(enum.Enum):
+    EXTENSION = "extension"
+    FORWARDER = "forwarder"
+    DATADOG_API = "datadog_api"
+    NO_METRICS = "no_metrics"
 
-if should_use_extension:
+
+def _select_metrics_handler():
+    if should_use_extension:
+        return MetricsHandler.EXTENSION
+    if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
+        return MetricsHandler.FORWARDER
+
+    if fips_mode_enabled:
+        logger.debug(
+            "With FIPS mode enabled, the Datadog API metrics handler is unavailable."
+        )
+        return MetricsHandler.NO_METRICS
+
+    return MetricsHandler.DATADOG_API
+
+
+metrics_handler = _select_metrics_handler()
+logger.debug("identified primary metrics handler as %s", metrics_handler)
+
+
+lambda_stats = None
+if metrics_handler == MetricsHandler.EXTENSION:
     from datadog_lambda.statsd_writer import StatsDWriter
 
     lambda_stats = StatsDWriter()
-else:
+
+elif metrics_handler == MetricsHandler.DATADOG_API:
     # Periodical flushing in a background thread is NOT guaranteed to succeed
     # and leads to data loss. When disabled, metrics are only flushed at the
     # end of invocation. To make metrics submitted from a long-running Lambda
     # function available sooner, consider using the Datadog Lambda extension.
-    from datadog_lambda.thread_stats_writer import ThreadStatsWriter
     from datadog_lambda.api import init_api
+    from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 
+    flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
     init_api()
     lambda_stats = ThreadStatsWriter(flush_in_thread)
 
+
 enhanced_metrics_enabled = (
     os.environ.get("DD_ENHANCED_METRICS", "true").lower() == "true"
 )
@@ -44,16 +73,19 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
     Submit a data point to Datadog distribution metrics.
     https://docs.datadoghq.com/graphing/metrics/distributions/
 
-    When DD_FLUSH_TO_LOG is True, write metric to log, and
-    wait for the Datadog Log Forwarder Lambda function to submit
-    the metrics asynchronously.
+    If the Datadog Lambda Extension is present, metrics are submitted to its
+    dogstatsd endpoint.
+
+    When DD_FLUSH_TO_LOG is True or force_async is True, write metric to log,
+    and wait for the Datadog Log Forwarder Lambda function to submit the
+    metrics asynchronously.
 
     Otherwise, the metrics will be submitted to the Datadog API
     periodically and at the end of the function execution in a
     background thread.
 
-    Note that if the extension is present, it will override the DD_FLUSH_TO_LOG value
-    and always use the layer to send metrics to the extension
+    Note that if the extension is present, it will override the DD_FLUSH_TO_LOG
+    value and always use the layer to send metrics to the extension
     """
     if not metric_name or not isinstance(metric_name, str):
         logger.warning(
@@ -71,56 +103,54 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
         )
         return
 
-    flush_to_logs = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
     tags = [] if tags is None else list(tags)
     tags.append(dd_lambda_layer_tag)
 
-    if should_use_extension and timestamp is not None:
-        # The extension does not support timestamps for distributions so we create a
-        # a thread stats writer to submit metrics with timestamps to the API
-        timestamp_ceiling = int(
-            (datetime.now() - timedelta(hours=4)).timestamp()
-        )  # 4 hours ago
-        if isinstance(timestamp, datetime):
-            timestamp = int(timestamp.timestamp())
-        if timestamp_ceiling > timestamp:
-            logger.warning(
-                "Timestamp %s is older than 4 hours, not submitting metric %s",
-                timestamp,
-                metric_name,
-            )
-            return
-        global extension_thread_stats
-        if extension_thread_stats is None:
-            from datadog_lambda.thread_stats_writer import ThreadStatsWriter
-            from datadog_lambda.api import init_api
-
-            init_api()
-            extension_thread_stats = ThreadStatsWriter(flush_in_thread)
-
-        extension_thread_stats.distribution(
-            metric_name, value, tags=tags, timestamp=timestamp
-        )
-        return
+    if metrics_handler == MetricsHandler.EXTENSION:
+        if timestamp is not None:
+            if isinstance(timestamp, datetime):
+                timestamp = int(timestamp.timestamp())
+
+            timestamp_floor = int((datetime.now() - timedelta(hours=4)).timestamp())
+            if timestamp < timestamp_floor:
+                logger.warning(
+                    "Timestamp %s is older than 4 hours, not submitting metric %s",
+                    timestamp,
+                    metric_name,
+                )
+                return
 
-    if should_use_extension:
         logger.debug(
             "Sending metric %s value %s to Datadog via extension", metric_name, value
         )
         lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
+
+    elif force_async or (metrics_handler == MetricsHandler.FORWARDER):
+        write_metric_point_to_stdout(metric_name, value, timestamp=timestamp, tags=tags)
+
+    elif metrics_handler == MetricsHandler.DATADOG_API:
+        lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
+
+    elif metrics_handler == MetricsHandler.NO_METRICS:
+        logger.debug(
+            "Metric %s cannot be submitted because the metrics handler is disabled",
+            metric_name,
+        ),
+
     else:
-        if flush_to_logs or force_async:
-            write_metric_point_to_stdout(
-                metric_name, value, timestamp=timestamp, tags=tags
-            )
-        else:
-            lambda_stats.distribution(
-                metric_name, value, tags=tags, timestamp=timestamp
-            )
+        # This should be qutie impossible, but let's at least log a message if
+        # it somehow happens.
+        logger.debug(
+            "Metric %s cannot be submitted because the metrics handler is not configured: %s",
+            metric_name,
+            metrics_handler,
+        )
 
 
-def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]):
+def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=None):
     """Writes the specified metric point to standard output"""
+    tags = tags or []
+
     logger.debug(
         "Sending metric %s value %s to Datadog via log forwarder", metric_name, value
     )
@@ -138,19 +168,8 @@ def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]):
 
 
 def flush_stats(lambda_context=None):
-    lambda_stats.flush()
-
-    if extension_thread_stats is not None:
-        tags = None
-        if lambda_context is not None:
-            tags = get_enhanced_metrics_tags(lambda_context)
-            split_arn = lambda_context.invoked_function_arn.split(":")
-            if len(split_arn) > 7:
-                # Get rid of the alias
-                split_arn.pop()
-            arn = ":".join(split_arn)
-            tags.append("function_arn:" + arn)
-        extension_thread_stats.flush(tags)
+    if lambda_stats is not None:
+        lambda_stats.flush()
 
 
 def submit_enhanced_metric(metric_name, lambda_context):
@@ -188,3 +207,17 @@ def submit_errors_metric(lambda_context):
         lambda_context (object): Lambda context dict passed to the function by AWS
     """
     submit_enhanced_metric("errors", lambda_context)
+
+
+def submit_dynamodb_stream_type_metric(event):
+    stream_view_type = (
+        event.get("Records", [{}])[0].get("dynamodb", {}).get("StreamViewType")
+    )
+    if stream_view_type:
+        lambda_metric(
+            "datadog.serverless.dynamodb.stream.type",
+            1,
+            timestamp=None,
+            tags=[f"streamtype:{stream_view_type}"],
+            force_async=True,
+        )

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/span_pointers.py

@@ -6,6 +6,8 @@ from typing import Optional
 
 from ddtrace._trace._span_pointer import _SpanPointerDirection
 from ddtrace._trace._span_pointer import _SpanPointerDescription
+
+from datadog_lambda.metric import submit_dynamodb_stream_type_metric
 from datadog_lambda.trigger import EventTypes
 
 
@@ -28,6 +30,8 @@ def calculate_span_pointers(
                 return _calculate_s3_span_pointers_for_event(event)
 
             elif event_source.equals(EventTypes.DYNAMODB):
+                # Temporary metric. TODO eventually remove(@nhulston)
+                submit_dynamodb_stream_type_metric(event)
                 return _calculate_dynamodb_span_pointers_for_event(event)
 
     except Exception as e:

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/stats_writer.py

@@ -1,5 +1,5 @@
 class StatsWriter:
-    def distribution(self, metric_name, value, tags=[], timestamp=None):
+    def distribution(self, metric_name, value, tags=None, timestamp=None):
         raise NotImplementedError()
 
     def flush(self):

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/statsd_writer.py

@@ -1,5 +1,5 @@
-from datadog_lambda.stats_writer import StatsWriter
 from datadog_lambda.dogstatsd import statsd
+from datadog_lambda.stats_writer import StatsWriter
 
 
 class StatsDWriter(StatsWriter):
@@ -7,8 +7,8 @@ class StatsDWriter(StatsWriter):
     Writes distribution metrics using StatsD protocol
     """
 
-    def distribution(self, metric_name, value, tags=[], timestamp=None):
-        statsd.distribution(metric_name, value, tags=tags)
+    def distribution(self, metric_name, value, tags=None, timestamp=None):
+        statsd.distribution(metric_name, value, tags=tags, timestamp=timestamp)
 
     def flush(self):
         pass

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/thread_stats_writer.py

@@ -3,6 +3,7 @@ import logging
 # Make sure that this package would always be lazy-loaded/outside from the critical path
 # since underlying packages are quite heavy to load and useless when the extension is present
 from datadog.threadstats import ThreadStats
+
 from datadog_lambda.stats_writer import StatsWriter
 
 logger = logging.getLogger(__name__)
@@ -17,7 +18,7 @@ class ThreadStatsWriter(StatsWriter):
         self.thread_stats = ThreadStats(compress_payload=True)
         self.thread_stats.start(flush_in_thread=flush_in_thread)
 
-    def distribution(self, metric_name, value, tags=[], timestamp=None):
+    def distribution(self, metric_name, value, tags=None, timestamp=None):
         self.thread_stats.distribution(
             metric_name, value, tags=tags, timestamp=timestamp
         )

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/tracing.py

@@ -2,10 +2,8 @@
 # under the Apache License Version 2.0.
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
-import hashlib
 import logging
 import os
-import base64
 import traceback
 import ujson as json
 from datetime import datetime, timezone
@@ -39,6 +37,7 @@ from datadog_lambda.trigger import (
     _EventSource,
     parse_event_source,
     get_first_record,
+    is_step_function_event,
     EventTypes,
     EventSubtypes,
 )
@@ -258,6 +257,8 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
             dd_json_data = None
             dd_json_data_type = dd_payload.get("Type") or dd_payload.get("dataType")
             if dd_json_data_type == "Binary":
+                import base64
+
                 dd_json_data = dd_payload.get("binaryValue") or dd_payload.get("Value")
                 if dd_json_data:
                     dd_json_data = base64.b64decode(dd_json_data)
@@ -271,6 +272,15 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
 
             if dd_json_data:
                 dd_data = json.loads(dd_json_data)
+
+                if is_step_function_event(dd_data):
+                    try:
+                        return extract_context_from_step_functions(dd_data, None)
+                    except Exception:
+                        logger.debug(
+                            "Failed to extract Step Functions context from SQS/SNS event."
+                        )
+
                 return propagator.extract(dd_data)
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
@@ -313,6 +323,15 @@ def _extract_context_from_eventbridge_sqs_event(event):
     body = json.loads(body_str)
     detail = body.get("detail")
     dd_context = detail.get("_datadog")
+
+    if is_step_function_event(dd_context):
+        try:
+            return extract_context_from_step_functions(dd_context, None)
+        except Exception:
+            logger.debug(
+                "Failed to extract Step Functions context from EventBridge to SQS event."
+            )
+
     return propagator.extract(dd_context)
 
 
@@ -320,12 +339,23 @@ def extract_context_from_eventbridge_event(event, lambda_context):
     """
     Extract datadog trace context from an EventBridge message's Details.
     This is only possible if Details is a JSON string.
+
+    If we find a Step Function context, try to extract the trace context from
+    that header.
     """
     try:
         detail = event.get("detail")
         dd_context = detail.get("_datadog")
         if not dd_context:
             return extract_context_from_lambda_context(lambda_context)
+
+        try:
+            return extract_context_from_step_functions(dd_context, None)
+        except Exception:
+            logger.debug(
+                "Failed to extract Step Functions context from EventBridge event."
+            )
+
         return propagator.extract(dd_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
@@ -343,6 +373,8 @@ def extract_context_from_kinesis_event(event, lambda_context):
             return extract_context_from_lambda_context(lambda_context)
         data = kinesis.get("data")
         if data:
+            import base64
+
             b64_bytes = data.encode("ascii")
             str_bytes = base64.b64decode(b64_bytes)
             data_str = str_bytes.decode("ascii")
@@ -357,6 +389,8 @@ def extract_context_from_kinesis_event(event, lambda_context):
 
 
 def _deterministic_sha256_hash(s: str, part: str) -> int:
+    import hashlib
+
     sha256_hash = hashlib.sha256(s.encode()).hexdigest()
     # First two chars is '0b'. zfill to ensure 256 bits, but we only care about the first 128 bits
     binary_hash = bin(int(sha256_hash, 16))[2:].zfill(256)
@@ -424,7 +458,7 @@ def _generate_sfn_trace_id(execution_id: str, part: str):
 def extract_context_from_step_functions(event, lambda_context):
     """
     Only extract datadog trace context when Step Functions Context Object is injected
-    into lambda's event dict.
+    into lambda's event dict. Unwrap "Payload" if it exists to handle Legacy Lambda cases.
 
     If '_datadog' header is present, we have two cases:
       1. Root is a Lambda and we use its traceID
@@ -435,25 +469,25 @@ def extract_context_from_step_functions(event, lambda_context):
     object.
     """
     try:
+        event = event.get("Payload", event)
+        event = event.get("_datadog", event)
+
         meta = {}
-        dd_data = event.get("_datadog")
 
-        if dd_data and dd_data.get("serverless-version") == "v1":
-            if "x-datadog-trace-id" in dd_data:  # lambda root
-                trace_id = int(dd_data.get("x-datadog-trace-id"))
-                high_64_bit_trace_id = _parse_high_64_bits(
-                    dd_data.get("x-datadog-tags")
-                )
+        if event.get("serverless-version") == "v1":
+            if "x-datadog-trace-id" in event:  # lambda root
+                trace_id = int(event.get("x-datadog-trace-id"))
+                high_64_bit_trace_id = _parse_high_64_bits(event.get("x-datadog-tags"))
                 if high_64_bit_trace_id:
                     meta["_dd.p.tid"] = high_64_bit_trace_id
             else:  # sfn root
-                root_execution_id = dd_data.get("RootExecutionId")
+                root_execution_id = event.get("RootExecutionId")
                 trace_id = _generate_sfn_trace_id(root_execution_id, LOWER_64_BITS)
                 meta["_dd.p.tid"] = _generate_sfn_trace_id(
                     root_execution_id, HIGHER_64_BITS
                 )
 
-            parent_id = _generate_sfn_parent_id(dd_data)
+            parent_id = _generate_sfn_parent_id(event)
         else:
             execution_id = event.get("Execution").get("Id")
             trace_id = _generate_sfn_trace_id(execution_id, LOWER_64_BITS)
@@ -472,20 +506,6 @@ def extract_context_from_step_functions(event, lambda_context):
         return extract_context_from_lambda_context(lambda_context)
 
 
-def is_legacy_lambda_step_function(event):
-    """
-    Check if the event is a step function that called a legacy lambda
-    """
-    if not isinstance(event, dict) or "Payload" not in event:
-        return False
-
-    event = event.get("Payload")
-    return isinstance(event, dict) and (
-        "_datadog" in event
-        or ("Execution" in event and "StateMachine" in event and "State" in event)
-    )
-
-
 def extract_context_custom_extractor(extractor, event, lambda_context):
     """
     Extract Datadog trace context using a custom trace extractor function
@@ -535,6 +555,8 @@ def get_injected_authorizer_data(event, is_http_api) -> dict:
         if not dd_data_raw:
             return None
 
+        import base64
+
         injected_data = json.loads(base64.b64decode(dd_data_raw))
 
         # Lambda authorizer's results can be cached. But the payload will still have the injected
@@ -1309,8 +1331,18 @@ def create_inferred_span_from_eventbridge_event(event, context):
         synchronicity="async",
         tag_source="self",
     )
-    dt_format = "%Y-%m-%dT%H:%M:%SZ"
+
     timestamp = event.get("time")
+    dt_format = "%Y-%m-%dT%H:%M:%SZ"
+
+    # Use more granular timestamp from upstream Step Function if possible
+    try:
+        if is_step_function_event(event.get("detail")):
+            timestamp = event["detail"]["_datadog"]["State"]["EnteredTime"]
+            dt_format = "%Y-%m-%dT%H:%M:%S.%fZ"
+    except (TypeError, KeyError, AttributeError):
+        logger.debug("Error parsing timestamp from Step Functions event")
+
     dt = datetime.strptime(timestamp, dt_format)
 
     tracer.set_tags(_dd_origin)
@@ -1320,6 +1352,11 @@ def create_inferred_span_from_eventbridge_event(event, context):
     if span:
         span.set_tags(tags)
     span.start = dt.replace(tzinfo=timezone.utc).timestamp()
+
+    # Since inferred span will later parent Lambda, preserve Lambda's current parent
+    if dd_trace_context.span_id:
+        span.parent_id = dd_trace_context.span_id
+
     return span
 
 

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/trigger.py

@@ -3,7 +3,6 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
-import base64
 import gzip
 import ujson as json
 from io import BytesIO, BufferedReader
@@ -146,9 +145,7 @@ def parse_event_source(event: dict) -> _EventSource:
     if event.get("source") == "aws.events" or has_event_categories:
         event_source = _EventSource(EventTypes.CLOUDWATCH_EVENTS)
 
-    if (
-        "_datadog" in event and event.get("_datadog").get("serverless-version") == "v1"
-    ) or ("Execution" in event and "StateMachine" in event and "State" in event):
+    if is_step_function_event(event):
         event_source = _EventSource(EventTypes.STEPFUNCTIONS)
 
     event_record = get_first_record(event)
@@ -244,6 +241,8 @@ def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> s
 
     # e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz
     if source.event_type == EventTypes.CLOUDWATCH_LOGS:
+        import base64
+
         with gzip.GzipFile(
             fileobj=BytesIO(base64.b64decode(event.get("awslogs", {}).get("data")))
         ) as decompress_stream:
@@ -369,3 +368,29 @@ def extract_http_status_code_tag(trigger_tags, response):
         status_code = response.status_code
 
     return str(status_code)
+
+
+def is_step_function_event(event):
+    """
+    Check if the event is a step function that invoked the current lambda.
+
+    The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a
+    "_datadog" for JSONata style context propagation.
+
+    The actual event must contain "Execution", "StateMachine", and "State" fields.
+    """
+    event = event.get("Payload", event)
+
+    # JSONPath style
+    if "Execution" in event and "StateMachine" in event and "State" in event:
+        return True
+
+    # JSONata style
+    dd_context = event.get("_datadog")
+    return (
+        dd_context
+        and "Execution" in dd_context
+        and "StateMachine" in dd_context
+        and "State" in dd_context
+        and "serverless-version" in dd_context
+    )

datadog_lambda-6.108.0/datadog_lambda/version.py

@@ -0,0 +1 @@
+__version__ = "6.108.0"

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/datadog_lambda/wrapper.py

@@ -2,7 +2,6 @@
 # under the Apache License Version 2.0.
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
-import base64
 import os
 import logging
 import traceback
@@ -23,11 +22,6 @@ from datadog_lambda.constants import (
     XraySubsegment,
     Headers,
 )
-from datadog_lambda.metric import (
-    flush_stats,
-    submit_invocations_metric,
-    submit_errors_metric,
-)
 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
@@ -45,7 +39,6 @@ from datadog_lambda.tracing import (
     is_authorizer_response,
     tracer,
     propagator,
-    is_legacy_lambda_step_function,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -242,7 +235,11 @@ class _LambdaDecorator(object):
             self.response = self.func(event, context, **kwargs)
             return self.response
         except Exception:
-            submit_errors_metric(context)
+            if not should_use_extension:
+                from datadog_lambda.metric import submit_errors_metric
+
+                submit_errors_metric(context)
+
             if self.span:
                 self.span.set_traceback()
             raise
@@ -268,6 +265,9 @@ class _LambdaDecorator(object):
         injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns
         if request_id is not None:
             injected_headers[Headers.Authorizing_Request_Id] = request_id
+
+        import base64
+
         datadog_data = base64.b64encode(
             json.dumps(injected_headers, escape_forward_slashes=False).encode()
         ).decode()
@@ -278,9 +278,12 @@ class _LambdaDecorator(object):
         try:
             self.response = None
             set_cold_start(init_timestamp_ns)
-            submit_invocations_metric(context)
-            if is_legacy_lambda_step_function(event):
-                event = event["Payload"]
+
+            if not should_use_extension:
+                from datadog_lambda.metric import submit_invocations_metric
+
+                submit_invocations_metric(context)
+
             self.trigger_tags = extract_trigger_tags(event, context)
             # Extract Datadog trace context and source from incoming requests
             dd_context, trace_context_source, event_source = extract_dd_trace_context(
@@ -379,6 +382,8 @@ class _LambdaDecorator(object):
                     logger.debug("Failed to create cold start spans. %s", e)
 
             if not self.flush_to_log or should_use_extension:
+                from datadog_lambda.metric import flush_stats
+
                 flush_stats(context)
             if should_use_extension and self.local_testing_mode:
                 # when testing locally, the extension does not know when an

{datadog_lambda-6.106.0 → datadog_lambda-6.108.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "datadog_lambda"
-version = "6.106.0"
+version = "6.108.0"
 description = "The Datadog AWS Lambda Library"
 authors = ["Datadog, Inc. <dev@datadoghq.com>"]
 license = "Apache-2.0"
@@ -28,9 +28,9 @@ classifiers = [
 python = ">=3.8.0,<4"
 datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
-ddtrace = ">=2.20.0"
+ddtrace = ">=2.20.0,<4"
 ujson = ">=5.9.0"
-boto3 = { version = "^1.34.0", optional = true }
+botocore = { version = "^1.34.0", optional = true }
 requests = { version ="^2.22.0", optional = true }
 pytest = { version= "^8.0.0", optional = true }
 pytest-benchmark = { version = "^4.0", optional = true }
@@ -38,7 +38,7 @@ flake8 = { version = "^5.0.4", optional = true }
 
 [tool.poetry.extras]
 dev = [
-    "boto3",
+    "botocore",
     "flake8",
     "pytest",
     "pytest-benchmark",

datadog_lambda-6.106.0/datadog_lambda/api.py

@@ -1,89 +0,0 @@
-import os
-import logging
-import base64
-
-logger = logging.getLogger(__name__)
-KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName"
-
-
-def decrypt_kms_api_key(kms_client, ciphertext):
-    from botocore.exceptions import ClientError
-
-    """
-    Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS.
-    For this to work properly, the Lambda function must have the appropriate IAM permissions.
-
-    Args:
-        kms_client: The KMS client to use for decryption
-        ciphertext (string): The base64-encoded ciphertext to decrypt
-    """
-    decoded_bytes = base64.b64decode(ciphertext)
-
-    """
-    When the API key is encrypted using the AWS console, the function name is added as an
-    encryption context. When the API key is encrypted using the AWS CLI, no encryption context
-    is added. We need to try decrypting the API key both with and without the encryption context.
-    """
-    # Try without encryption context, in case API key was encrypted using the AWS CLI
-    function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME")
-    try:
-        plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)[
-            "Plaintext"
-        ].decode("utf-8")
-    except ClientError:
-        logger.debug(
-            "Failed to decrypt ciphertext without encryption context, \
-            retrying with encryption context"
-        )
-        # Try with encryption context, in case API key was encrypted using the AWS Console
-        plaintext = kms_client.decrypt(
-            CiphertextBlob=decoded_bytes,
-            EncryptionContext={
-                KMS_ENCRYPTION_CONTEXT_KEY: function_name,
-            },
-        )["Plaintext"].decode("utf-8")
-
-    return plaintext
-
-
-def init_api():
-    if not os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
-        # Make sure that this package would always be lazy-loaded/outside from the critical path
-        # since underlying packages are quite heavy to load
-        # and useless with the extension unless sending metrics with timestamps
-        from datadog import api
-
-        if not api._api_key:
-            import boto3
-
-            DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "")
-            DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "")
-            DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "")
-            DD_API_KEY = os.environ.get(
-                "DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")
-            )
-
-            if DD_API_KEY_SECRET_ARN:
-                api._api_key = boto3.client("secretsmanager").get_secret_value(
-                    SecretId=DD_API_KEY_SECRET_ARN
-                )["SecretString"]
-            elif DD_API_KEY_SSM_NAME:
-                api._api_key = boto3.client("ssm").get_parameter(
-                    Name=DD_API_KEY_SSM_NAME, WithDecryption=True
-                )["Parameter"]["Value"]
-            elif DD_KMS_API_KEY:
-                kms_client = boto3.client("kms")
-                api._api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY)
-            else:
-                api._api_key = DD_API_KEY
-
-        logger.debug("Setting DATADOG_API_KEY of length %d", len(api._api_key))
-
-        # Set DATADOG_HOST, to send data to a non-default Datadog datacenter
-        api._api_host = os.environ.get(
-            "DATADOG_HOST", "https://api." + os.environ.get("DD_SITE", "datadoghq.com")
-        )
-        logger.debug("Setting DATADOG_HOST to %s", api._api_host)
-
-        # Unmute exceptions from datadog api client, so we can catch and handle them
-        api._mute = False

datadog_lambda-6.106.0/datadog_lambda/version.py

@@ -1 +0,0 @@
-__version__ = "6.106.0"