Add service.

Author: purple4reina

datadog_lambda/config.py

@@ -27,6 +27,7 @@ def as_bool(val):
 
 class Config:
 
+    service = _get_env("DD_SERVICE")
     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

@@ -158,7 +158,6 @@ def __init__(self, func):
         """Executes when the wrapped function gets wrapped"""
         try:
             self.func = func
-            self.service = os.environ.get(DD_SERVICE, None)
             self.extractor_env = os.environ.get(DD_TRACE_EXTRACTOR, None)
             self.trace_extractor = None
             self.span = None
@@ -200,7 +199,7 @@ def __init__(self, func):
                     logger.debug(f"Malformatted for env {DD_COLD_START_TRACE_SKIP_LIB}")
             self.response = None
             if profiling_env_var:
-                self.prof = profiler.Profiler(env=env_env_var, service=self.service)
+                self.prof = profiler.Profiler(env=env_env_var, service=config.service)
             if self.extractor_env:
                 extractor_parts = self.extractor_env.rsplit(".", 1)
                 if len(extractor_parts) == 2:
@@ -367,8 +366,8 @@ def _after(self, event, context):
                 if status_code:
                     self.inferred_span.set_tag("http.status_code", status_code)
 
-                if self.service:
-                    self.inferred_span.set_tag("peer.service", self.service)
+                if config.service:
+                    self.inferred_span.set_tag("peer.service", config.service)
 
                 if InferredSpanInfo.is_async(self.inferred_span) and self.span:
                     self.inferred_span.finish(finish_time=self.span.start)

tests/test_config.py

@@ -15,6 +15,9 @@ def set_env(key, value):
 
 
 _test_config_from_environ = (
+    ("DD_SERVICE", "service", None, None),
+    ("DD_SERVICE", "service", "", ""),
+    ("DD_SERVICE", "service", "my_service", "my_service"),
     ("AWS_LAMBDA_FUNCTION_NAME", "function_name", None, "function"),
     ("AWS_LAMBDA_FUNCTION_NAME", "function_name", "", ""),
     ("AWS_LAMBDA_FUNCTION_NAME", "function_name", "my_function", "my_function"),