Add env.

Author: purple4reina

datadog_lambda/config.py

@@ -40,6 +40,8 @@ def as_bool(val):
 class Config:
 
     service = _get_env("DD_SERVICE")
+    env = _get_env("DD_ENV")
+
     add_span_pointers = _get_env("DD_BOTOCORE_ADD_SPAN_POINTERS", "true", as_bool)
     cold_start_tracing = _get_env("DD_COLD_START_TRACING", "true", as_bool)
     enhanced_metrics_enabled = _get_env("DD_ENHANCED_METRICS", "true", as_bool)

datadog_lambda/wrapper.py

@@ -83,8 +83,6 @@ def get_env_as_int(env_key, default_value: int) -> int:
     os.environ.get(DD_CAPTURE_LAMBDA_PAYLOAD, "false").lower() == "true"
 )
 
-env_env_var = os.environ.get(DD_ENV, None)
-
 init_timestamp_ns = time_ns()
 
 """
@@ -177,7 +175,7 @@ def __init__(self, func):
                     logger.debug(f"Malformatted for env {DD_COLD_START_TRACE_SKIP_LIB}")
             self.response = None
             if config.profiling_enabled:
-                self.prof = profiler.Profiler(env=env_env_var, service=config.service)
+                self.prof = profiler.Profiler(env=config.env, service=config.service)
             if config.trace_extractor:
                 extractor_parts = config.trace_extractor.rsplit(".", 1)
                 if len(extractor_parts) == 2:

tests/test_config.py

@@ -164,6 +164,9 @@ def set_env(key, value):
     ("DD_EXCEPTION_REPLAY_ENABLED", "exception_replay_enabled", "1", True),  # CHANGED
     ("DD_EXCEPTION_REPLAY_ENABLED", "exception_replay_enabled", "0", False),
     ("DD_EXCEPTION_REPLAY_ENABLED", "exception_replay_enabled", "purple", False),
+    ("DD_ENV", "env", None, None),
+    ("DD_ENV", "env", "", ""),
+    ("DD_ENV", "env", "my_env", "my_env"),
 )