datadog_lambda 5.91.0__py3-none-any.whl

datadog_lambda/__init__.py

@@ -0,0 +1,23 @@
+from datadog_lambda.cold_start import initialize_cold_start_tracing
+from datadog_lambda.logger import initialize_logging
+import os
+
+
+if os.environ.get("DD_INSTRUMENTATION_TELEMETRY_ENABLED") is None:
+    os.environ["DD_INSTRUMENTATION_TELEMETRY_ENABLED"] = "false"
+
+if os.environ.get("DD_API_SECURITY_ENABLED") is None:
+    os.environ["DD_API_SECURITY_ENABLED"] = "False"
+
+initialize_cold_start_tracing()
+
+# The minor version corresponds to the Lambda layer version.
+# E.g.,, version 0.5.0 gets packaged into layer version 5.
+try:
+    import importlib.metadata as importlib_metadata
+except ModuleNotFoundError:
+    import importlib_metadata
+
+__version__ = importlib_metadata.version(__name__)
+
+initialize_logging(__name__)

datadog_lambda/api.py

@@ -0,0 +1,93 @@
+import os
+import logging
+import base64
+from datadog_lambda.extension import should_use_extension
+
+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 should_use_extension
+        and 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 when the extension is present
+        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/cold_start.py

@@ -0,0 +1,257 @@
+import time
+import os
+from typing import List, Hashable
+import logging
+
+logger = logging.getLogger(__name__)
+
+_cold_start = True
+_proactive_initialization = False
+_lambda_container_initialized = False
+_tracer = None
+
+
+def set_cold_start(init_timestamp_ns):
+    """Set the value of the cold start global
+
+    This should be executed once per Lambda execution before the execution
+    """
+    global _cold_start
+    global _lambda_container_initialized
+    global _proactive_initialization
+    global _tracer
+    if not _lambda_container_initialized:
+        now = time.time_ns()
+        if (now - init_timestamp_ns) // 1_000_000_000 > 10:
+            _cold_start = False
+            _proactive_initialization = True
+        else:
+            _cold_start = not _lambda_container_initialized
+    else:
+        _cold_start = False
+        _proactive_initialization = False
+    _lambda_container_initialized = True
+    from ddtrace import tracer as _tracer
+
+
+def is_cold_start():
+    """Returns the value of the global cold_start"""
+    return _cold_start
+
+
+def is_proactive_init():
+    """Returns the value of the global proactive_initialization"""
+    return _proactive_initialization
+
+
+def is_new_sandbox():
+    return is_cold_start() or is_proactive_init()
+
+
+def get_cold_start_tag():
+    """Returns the cold start tag to be used in metrics"""
+    return "cold_start:{}".format(str(is_cold_start()).lower())
+
+
+def get_proactive_init_tag():
+    """Returns the proactive init tag to be used in metrics"""
+    return "proactive_initialization:{}".format(str(is_proactive_init()).lower())
+
+
+class ImportNode(object):
+    def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None):
+        self.module_name = module_name
+        self.full_file_path = full_file_path
+        self.start_time_ns = start_time_ns
+        self.end_time_ns = end_time_ns
+        self.children = []
+        self.context = None
+        if _lambda_container_initialized:
+            self.context = _tracer.context_provider.active()
+
+
+root_nodes: List[ImportNode] = []
+import_stack: List[ImportNode] = []
+already_wrapped_loaders = set()
+
+
+def reset_node_stacks():
+    root_nodes.clear()
+    import_stack.clear()
+
+
+def push_node(module_name, file_path):
+    node = ImportNode(module_name, file_path, time.time_ns())
+    global import_stack
+    if import_stack:
+        import_stack[-1].children.append(node)
+    import_stack.append(node)
+
+
+def pop_node(module_name):
+    global import_stack
+    if not import_stack:
+        return
+    node = import_stack.pop()
+    if node.module_name != module_name:
+        return
+    end_time_ns = time.time_ns()
+    node.end_time_ns = end_time_ns
+    if not import_stack:  # import_stack empty, a root node has been found
+        global root_nodes
+        root_nodes.append(node)
+
+
+def wrap_exec_module(original_exec_module):
+    def wrapped_method(module):
+        should_pop = False
+        try:
+            spec = module.__spec__
+            push_node(spec.name, spec.origin)
+            should_pop = True
+        except Exception:
+            pass
+        try:
+            return original_exec_module(module)
+        finally:
+            if should_pop:
+                pop_node(spec.name)
+
+    return wrapped_method
+
+
+def wrap_find_spec(original_find_spec):
+    def wrapped_find_spec(*args, **kwargs):
+        spec = original_find_spec(*args, **kwargs)
+        if spec is None:
+            return None
+        loader = getattr(spec, "loader", None)
+        if (
+            loader is not None
+            and isinstance(loader, Hashable)
+            and loader not in already_wrapped_loaders
+        ):
+            if hasattr(loader, "exec_module"):
+                try:
+                    loader.exec_module = wrap_exec_module(loader.exec_module)
+                    already_wrapped_loaders.add(loader)
+                except Exception as e:
+                    logger.debug("Failed to wrap the loader. %s", e)
+        return spec
+
+    return wrapped_find_spec
+
+
+def initialize_cold_start_tracing():
+    if (
+        is_new_sandbox()
+        and os.environ.get("DD_TRACE_ENABLED", "true").lower() == "true"
+        and os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true"
+    ):
+        from sys import meta_path
+
+        for importer in meta_path:
+            try:
+                importer.find_spec = wrap_find_spec(importer.find_spec)
+            except Exception:
+                pass
+
+
+class ColdStartTracer(object):
+    def __init__(
+        self,
+        tracer,
+        function_name,
+        current_span_start_time_ns,
+        trace_ctx,
+        min_duration_ms: int,
+        ignored_libs: List[str] = None,
+    ):
+        if ignored_libs is None:
+            ignored_libs = []
+        self._tracer = tracer
+        self.function_name = function_name
+        self.current_span_start_time_ns = current_span_start_time_ns
+        self.min_duration_ms = min_duration_ms
+        self.trace_ctx = trace_ctx
+        self.ignored_libs = ignored_libs
+        self.need_to_reactivate_context = True
+
+    def trace(self, root_nodes: List[ImportNode] = root_nodes):
+        if not root_nodes:
+            return
+        cold_start_span_start_time_ns = root_nodes[0].start_time_ns
+        cold_start_span_end_time_ns = min(
+            root_nodes[-1].end_time_ns, self.current_span_start_time_ns
+        )
+        cold_start_span = self.create_cold_start_span(cold_start_span_start_time_ns)
+        while root_nodes:
+            root_node = root_nodes.pop()
+            parent = root_node.context or cold_start_span
+            self.trace_tree(root_node, parent)
+        self.finish_span(cold_start_span, cold_start_span_end_time_ns)
+
+    def trace_tree(self, import_node: ImportNode, parent_span):
+        if (
+            import_node.end_time_ns - import_node.start_time_ns
+            < self.min_duration_ms * 1e6
+            or import_node.module_name in self.ignored_libs
+        ):
+            return
+
+        span = self.start_span(
+            "aws.lambda.import", import_node.module_name, import_node.start_time_ns
+        )
+        tags = {
+            "resource_names": import_node.module_name,
+            "resource.name": import_node.module_name,
+            "filename": import_node.full_file_path,
+            "operation_name": self.get_operation_name(import_node.full_file_path),
+        }
+        span.set_tags(tags)
+        if parent_span:
+            span.parent_id = parent_span.span_id
+        for child_node in import_node.children:
+            self.trace_tree(child_node, span)
+        self.finish_span(span, import_node.end_time_ns)
+
+    def create_cold_start_span(self, start_time_ns):
+        span = self.start_span("aws.lambda.load", self.function_name, start_time_ns)
+        tags = {
+            "resource_names": self.function_name,
+            "resource.name": self.function_name,
+            "operation_name": "aws.lambda.load",
+        }
+        span.set_tags(tags)
+        return span
+
+    def start_span(self, span_type, resource, start_time_ns):
+        if self.need_to_reactivate_context:
+            self._tracer.context_provider.activate(
+                self.trace_ctx
+            )  # reactivate required after each finish() call
+            self.need_to_reactivate_context = False
+        span_kwargs = {
+            "service": "aws.lambda",
+            "resource": resource,
+            "span_type": span_type,
+        }
+        span = self._tracer.trace(span_type, **span_kwargs)
+        span.start_ns = start_time_ns
+        return span
+
+    def finish_span(self, span, finish_time_ns):
+        span.finish(finish_time_ns / 1e9)
+        self.need_to_reactivate_context = True
+
+    def get_operation_name(self, filename: str):
+        if filename is None:
+            return "aws.lambda.import_core_module"
+        if not isinstance(filename, str):
+            return "aws.lambda.import"
+        if filename.startswith("/opt/"):
+            return "aws.lambda.import_layer"
+        elif filename.startswith("/var/lang/"):
+            return "aws.lambda.import_runtime"
+        else:
+            return "aws.lambda.import"

datadog_lambda/constants.py

@@ -0,0 +1,53 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2019 Datadog, Inc.
+
+# Datadog trace sampling priority
+
+
+class SamplingPriority(object):
+    USER_REJECT = -1
+    AUTO_REJECT = 0
+    AUTO_KEEP = 1
+    USER_KEEP = 2
+
+
+# Datadog trace headers
+class TraceHeader(object):
+    TRACE_ID = "x-datadog-trace-id"
+    PARENT_ID = "x-datadog-parent-id"
+    SAMPLING_PRIORITY = "x-datadog-sampling-priority"
+
+
+# X-Ray subsegment to save Datadog trace metadata
+class XraySubsegment(object):
+    NAME = "datadog-metadata"
+    TRACE_KEY = "trace"
+    LAMBDA_FUNCTION_TAGS_KEY = "lambda_function_tags"
+    NAMESPACE = "datadog"
+
+
+# TraceContextSource of datadog context. The DD_MERGE_XRAY_TRACES
+# feature uses this to determine when to use X-Ray as the parent
+# trace.
+class TraceContextSource(object):
+    XRAY = "xray"
+    EVENT = "event"
+    DDTRACE = "ddtrace"
+
+
+# X-Ray deamon
+class XrayDaemon(object):
+    XRAY_TRACE_ID_HEADER_NAME = "_X_AMZN_TRACE_ID"
+    XRAY_DAEMON_ADDRESS = "AWS_XRAY_DAEMON_ADDRESS"
+    FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME"
+
+
+class Headers(object):
+    Parent_Span_Finish_Time = "x-datadog-parent-span-finish-time"
+
+    # For one request from the client, the event.requestContext.requestIds in the authorizer lambda
+    # invocation and the main function invocation are IDENTICAL. Therefore we can use it to tell
+    # whether current invocation is the actual original authorizing request or a cached request.
+    Authorizing_Request_Id = "x-datadog-authorizing-requestid"

datadog_lambda/dogstatsd.py

@@ -0,0 +1,143 @@
+import logging
+import os
+import socket
+import errno
+import re
+from threading import Lock
+
+
+MIN_SEND_BUFFER_SIZE = 32 * 1024
+log = logging.getLogger("datadog_lambda.dogstatsd")
+
+
+class DogStatsd(object):
+    def __init__(self):
+        self._socket_lock = Lock()
+        self.socket_path = None
+        self.host = "localhost"
+        self.port = 8125
+        self.socket = None
+        self.encoding = "utf-8"
+
+    def get_socket(self, telemetry=False):
+        """
+        Return a connected socket.
+
+        Note: connect the socket before assigning it to the class instance to
+        avoid bad thread race conditions.
+        """
+        with self._socket_lock:
+            self.socket = self._get_udp_socket(
+                self.host,
+                self.port,
+            )
+            return self.socket
+
+    @classmethod
+    def _ensure_min_send_buffer_size(cls, sock, min_size=MIN_SEND_BUFFER_SIZE):
+        # Increase the receiving buffer size where needed (e.g. MacOS has 4k RX
+        # buffers which is half of the max packet size that the client will send.
+        if os.name == "posix":
+            try:
+                recv_buff_size = sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
+                if recv_buff_size <= min_size:
+                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, min_size)
+                    log.debug("Socket send buffer increased to %dkb", min_size / 1024)
+            finally:
+                pass
+
+    @classmethod
+    def _get_udp_socket(cls, host, port):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sock.setblocking(0)
+        cls._ensure_min_send_buffer_size(sock)
+        sock.connect((host, port))
+
+        return sock
+
+    def distribution(self, metric, value, tags=None):
+        """
+        Send a global distribution value, optionally setting tags.
+
+        >>> statsd.distribution("uploaded.file.size", 1445)
+        >>> statsd.distribution("album.photo.count", 26, tags=["gender:female"])
+        """
+        self._report(metric, "d", value, tags)
+
+    def close_socket(self):
+        """
+        Closes connected socket if connected.
+        """
+        with self._socket_lock:
+            if self.socket:
+                try:
+                    self.socket.close()
+                except OSError as e:
+                    log.error("Unexpected error: %s", str(e))
+                self.socket = None
+
+    def normalize_tags(self, tag_list):
+        TAG_INVALID_CHARS_RE = re.compile(r"[^\w\d_\-:/\.]", re.UNICODE)
+        TAG_INVALID_CHARS_SUBS = "_"
+        return [
+            re.sub(TAG_INVALID_CHARS_RE, TAG_INVALID_CHARS_SUBS, tag)
+            for tag in tag_list
+        ]
+
+    def _serialize_metric(self, metric, metric_type, value, tags):
+        # Create/format the metric packet
+        return "%s:%s|%s%s" % (
+            metric,
+            value,
+            metric_type,
+            ("|#" + ",".join(self.normalize_tags(tags))) if tags else "",
+        )
+
+    def _report(self, metric, metric_type, value, tags):
+        if value is None:
+            return
+
+        payload = self._serialize_metric(metric, metric_type, value, tags)
+
+        # Send it
+        self._send_to_server(payload)
+
+    def _send_to_server(self, packet):
+        try:
+            mysocket = self.socket or self.get_socket()
+            mysocket.send(packet.encode(self.encoding))
+            return True
+        except socket.timeout:
+            # dogstatsd is overflowing, drop the packets (mimicks the UDP behaviour)
+            pass
+        except (socket.herror, socket.gaierror) as socket_err:
+            log.warning(
+                "Error submitting packet: %s, dropping the packet and closing the socket",
+                socket_err,
+            )
+            self.close_socket()
+        except socket.error as socket_err:
+            if socket_err.errno == errno.EAGAIN:
+                log.debug(
+                    "Socket send would block: %s, dropping the packet", socket_err
+                )
+            elif socket_err.errno == errno.ENOBUFS:
+                log.debug("Socket buffer full: %s, dropping the packet", socket_err)
+            elif socket_err.errno == errno.EMSGSIZE:
+                log.debug(
+                    "Packet size too big (size: %d): %s, dropping the packet",
+                    len(packet.encode(self.encoding)),
+                    socket_err,
+                )
+            else:
+                log.warning(
+                    "Error submitting packet: %s, dropping the packet and closing the socket",
+                    socket_err,
+                )
+                self.close_socket()
+        except Exception as e:
+            log.error("Unexpected error: %s", str(e))
+        return False
+
+
+statsd = DogStatsd()

datadog_lambda/extension.py

@@ -0,0 +1,29 @@
+import logging
+from os import path
+
+AGENT_URL = "http://127.0.0.1:8124"
+FLUSH_PATH = "/lambda/flush"
+EXTENSION_PATH = "/opt/extensions/datadog-agent"
+
+logger = logging.getLogger(__name__)
+
+
+def is_extension_present():
+    if path.exists(EXTENSION_PATH):
+        return True
+    return False
+
+
+def flush_extension():
+    try:
+        import urllib.request
+
+        req = urllib.request.Request(AGENT_URL + FLUSH_PATH, "".encode("ascii"))
+        urllib.request.urlopen(req)
+    except Exception as e:
+        logger.debug("Failed to flush extension, returned with error %s", e)
+        return False
+    return True
+
+
+should_use_extension = is_extension_present()

datadog_lambda/handler.py

@@ -0,0 +1,31 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# 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
+from datadog_lambda.wrapper import datadog_lambda_wrapper
+from datadog_lambda.module_name import modify_module_name
+
+
+class HandlerError(Exception):
+    pass
+
+
+path = os.environ.get("DD_LAMBDA_HANDLER", None)
+if path is None:
+    raise HandlerError(
+        "DD_LAMBDA_HANDLER is not defined. Can't use prebuilt datadog handler"
+    )
+parts = path.rsplit(".", 1)
+if len(parts) != 2:
+    raise HandlerError("Value %s for DD_LAMBDA_HANDLER has invalid format." % path)
+
+
+(mod_name, handler_name) = parts
+modified_mod_name = modify_module_name(mod_name)
+handler_module = import_module(modified_mod_name)
+handler = datadog_lambda_wrapper(getattr(handler_module, handler_name))

datadog_lambda/logger.py

@@ -0,0 +1,27 @@
+import logging
+import os
+
+try:
+    _level_mappping = logging.getLevelNamesMapping()
+except AttributeError:
+    # python 3.8
+    _level_mappping = {name: num for num, name in logging._levelToName.items()}
+# https://docs.datadoghq.com/agent/troubleshooting/debug_mode/?tab=agentv6v7#agent-log-level
+_level_mappping.update(
+    {
+        "TRACE": 5,
+        "WARN": logging.WARNING,
+        "OFF": 100,
+    }
+)
+
+
+def initialize_logging(name):
+    logger = logging.getLogger(name)
+    str_level = (os.environ.get("DD_LOG_LEVEL") or "INFO").upper()
+    level = _level_mappping.get(str_level)
+    if level is None:
+        logger.setLevel(logging.INFO)
+        logger.warning("Invalid log level: %s Defaulting to INFO", str_level)
+    else:
+        logger.setLevel(level)