Consolidate env reading to single config object.

Author: purple4reina

datadog_lambda/api.py

@@ -1,7 +1,7 @@
 import logging
 import os
 
-from datadog_lambda.fips import fips_mode_enabled
+from datadog_lambda.config import config
 
 logger = logging.getLogger(__name__)
 KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName"
@@ -29,7 +29,6 @@ def decrypt_kms_api_key(kms_client, ciphertext):
     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"
@@ -43,7 +42,7 @@ def decrypt_kms_api_key(kms_client, ciphertext):
         plaintext = kms_client.decrypt(
             CiphertextBlob=decoded_bytes,
             EncryptionContext={
-                KMS_ENCRYPTION_CONTEXT_KEY: function_name,
+                KMS_ENCRYPTION_CONTEXT_KEY: config.function_name,
             },
         )["Plaintext"].decode("utf-8")
 
@@ -66,7 +65,7 @@ def get_api_key() -> str:
     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:
+    if config.fips_mode_enabled:
         logger.debug(
             "FIPS mode is enabled, using FIPS endpoints for secrets management."
         )
@@ -82,7 +81,7 @@ def get_api_key() -> str:
             return ""
         endpoint_url = (
             f"https://secretsmanager-fips.{secrets_region}.amazonaws.com"
-            if fips_mode_enabled
+            if config.fips_mode_enabled
             else None
         )
         secrets_manager_client = _boto3_client(
@@ -95,7 +94,7 @@ def get_api_key() -> str:
         # 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
+            if config.fips_mode_enabled
             else None
         )
         ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint)
@@ -106,7 +105,7 @@ def get_api_key() -> str:
         # 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
+            if config.fips_mode_enabled
             else None
         )
         kms_client = _boto3_client("kms", endpoint_url=fips_endpoint)
@@ -118,7 +117,7 @@ def get_api_key() -> str:
 
 
 def init_api():
-    if not os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
+    if not config.flush_to_log:
         # 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

datadog_lambda/cold_start.py

@@ -1,8 +1,9 @@
 import time
-import os
 from typing import List, Hashable
 import logging
 
+from datadog_lambda.config import config
+
 logger = logging.getLogger(__name__)
 
 _cold_start = True
@@ -86,14 +87,12 @@ def reset_node_stacks():
 
 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()
@@ -102,7 +101,6 @@ def pop_node(module_name):
     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)
 
 
@@ -147,11 +145,7 @@ def wrapped_find_spec(*args, **kwargs):
 
 
 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"
-    ):
+    if is_new_sandbox() and config.trace_enabled and config.cold_start_tracing:
         from sys import meta_path
 
         for importer in meta_path:

datadog_lambda/config.py

@@ -0,0 +1,73 @@
+# 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.
+
+import logging
+import os
+
+
+def _get_env(key, default=None, cast=None):
+    """Get an environment variable with a default value."""
+    val = os.environ.get(key, default)
+    if cast is not None:
+        try:
+            val = cast(val)
+        except ValueError:
+            raise ValueError(f"Invalid value for {key}: {val}")
+            return cast(default)
+    return val
+
+
+def as_bool(val):
+    """Convert a string to a boolean."""
+    if isinstance(val, bool):
+        return val
+    if isinstance(val, str):
+        val = val.lower()
+        if val in ("true", "1", "yes"):
+            return True
+        elif val in ("false", "0", "no"):
+            return False
+    raise ValueError(f"Invalid boolean value: {val}")
+
+
+class config:
+
+    function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME")
+    flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
+    trace_enabled = os.environ.get("DD_TRACE_ENABLED", "true").lower() == "true"
+    cold_start_tracing = (
+        os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true"
+    )
+    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"
+    )
+    log_level = (os.environ.get("DD_LOG_LEVEL") or "INFO").upper()
+    flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
+    enhanced_metrics_enabled = (
+        os.environ.get("DD_ENHANCED_METRICS", "true").lower() == "true"
+    )
+    is_in_tests = os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true"
+    add_span_pointers = os.environ.get(
+        "DD_BOTOCORE_ADD_SPAN_POINTERS", "true"
+    ).lower() in ("true", "1")
+    otel_enabled = os.environ.get("DD_TRACE_OTEL_ENABLED", "false").lower() == "true"
+    is_lambda_context = bool(function_name)
+    telemetry_enabled = (
+        os.environ.get("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "false").lower()
+        == "true"
+    )
+
+
+if config.is_gov_region or config.fips_mode_enabled:
+    logger = logging.getLogger(__name__)
+    logger.debug(
+        "Python Lambda Layer FIPS mode is %s.",
+        "enabled" if config.fips_mode_enabled else "not enabled",
+    )

datadog_lambda/fips.py

@@ -1,5 +1,6 @@
 import logging
-import os
+
+from datadog_lambda.config import config
 
 try:
     _level_mappping = logging.getLevelNamesMapping()
@@ -18,10 +19,9 @@
 
 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)
+    level = _level_mappping.get(config.log_level)
     if level is None:
         logger.setLevel(logging.INFO)
-        logger.warning("Invalid log level: %s Defaulting to INFO", str_level)
+        logger.warning("Invalid log level: %s Defaulting to INFO", config.log_level)
     else:
         logger.setLevel(level)

datadog_lambda/logger.py

@@ -5,14 +5,13 @@
 
 import enum
 import logging
-import os
 import time
 from datetime import datetime, timedelta
 
 import ujson as json
 
+from datadog_lambda.config import config
 from datadog_lambda.extension import should_use_extension
-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__)
@@ -28,10 +27,10 @@ class MetricsHandler(enum.Enum):
 def _select_metrics_handler():
     if should_use_extension:
         return MetricsHandler.EXTENSION
-    if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
+    if config.flush_to_log:
         return MetricsHandler.FORWARDER
 
-    if fips_mode_enabled:
+    if config.fips_mode_enabled:
         logger.debug(
             "With FIPS mode enabled, the Datadog API metrics handler is unavailable."
         )
@@ -58,14 +57,8 @@ def _select_metrics_handler():
     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"
-)
+    lambda_stats = ThreadStatsWriter(config.flush_in_thread)
 
 
 def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=False):
@@ -191,7 +184,7 @@ def submit_enhanced_metric(metric_name, lambda_context):
         metric_name (str): metric name w/o enhanced prefix i.e. "invocations" or "errors"
         lambda_context (object): Lambda context dict passed to the function by AWS
     """
-    if not enhanced_metrics_enabled:
+    if not config.enhanced_metrics_enabled:
         logger.debug(
             "Not submitting enhanced metric %s because enhanced metrics are disabled",
             metric_name,

datadog_lambda/metric.py

@@ -3,7 +3,6 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
-import os
 import sys
 import logging
 import zlib
@@ -13,6 +12,7 @@
 from wrapt.importer import when_imported
 from ddtrace import patch_all as patch_all_dd
 
+from datadog_lambda import config
 from datadog_lambda.tracing import (
     get_dd_trace_context,
     dd_tracing_enabled,
@@ -44,8 +44,7 @@ def _patch_for_integration_tests():
     Patch `requests` to log the outgoing requests for integration tests.
     """
     global _integration_tests_patched
-    is_in_tests = os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true"
-    if not _integration_tests_patched and is_in_tests:
+    if not _integration_tests_patched and config.is_in_tests:
         wrap("requests", "Session.send", _log_request)
         _integration_tests_patched = True
 

datadog_lambda/patch.py

@@ -1,28 +1,23 @@
 from itertools import chain
 import logging
-import os
 from typing import List
 from typing import Optional
 
 from ddtrace._trace._span_pointer import _SpanPointerDirection
 from ddtrace._trace._span_pointer import _SpanPointerDescription
 
+from datadog_lambda import config
 from datadog_lambda.metric import submit_dynamodb_stream_type_metric
 from datadog_lambda.trigger import EventTypes
 
 
 logger = logging.getLogger(__name__)
 
 
-dd_botocore_add_span_pointers = os.environ.get(
-    "DD_BOTOCORE_ADD_SPAN_POINTERS", "true"
-).lower() in ("true", "1")
-
-
 def calculate_span_pointers(
     event_source,
     event,
-    botocore_add_span_pointers=dd_botocore_add_span_pointers,
+    botocore_add_span_pointers=config.add_span_pointers,
 ) -> List[_SpanPointerDescription]:
     try:
         if botocore_add_span_pointers: