Add make_inferred_span.

Author: purple4reina

datadog_lambda/config.py

@@ -9,23 +9,26 @@
 logger = logging.getLogger(__name__)
 
 
-def _get_env(key, default=None, cast=None):
+def _get_env(key, default=None, cast=None, depends_on_tracing=False):
     @property
     def _getter(self):
         if not hasattr(self, prop_key):
-            val = os.environ.get(key, default)
-            if cast is not None:
-                try:
-                    val = cast(val)
-                except (ValueError, TypeError):
-                    logger.warning(
-                        "Failed to cast environment variable '%s' with value '%s' to type %s. Using default value '%s'.",
-                        key,
-                        val,
-                        cast.__name__,
-                        default,
-                    )
-                    val = default
+            if depends_on_tracing and not config.trace_enabled:
+                val = False
+            else:
+                val = os.environ.get(key, default)
+                if cast is not None:
+                    try:
+                        val = cast(val)
+                    except (ValueError, TypeError):
+                        logger.warning(
+                            "Failed to cast environment variable '%s' with value '%s' to type %s. Using default value '%s'.",
+                            key,
+                            val,
+                            cast.__name__,
+                            default,
+                        )
+                        val = default
             setattr(self, prop_key, val)
         return getattr(self, prop_key)
 
@@ -75,6 +78,9 @@ class Config:
     llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool)
     exception_replay_enabled = _get_env("DD_EXCEPTION_REPLAY_ENABLED", "false", as_bool)
 
+    make_inferred_span = _get_env("DD_TRACE_MANAGED_SERVICES", "true", as_bool,
+                                   depends_on_tracing=True)
+
     local_test = _get_env("DD_LOCAL_TEST", "false", as_bool)
 
     @property

datadog_lambda/wrapper.py

@@ -126,9 +126,6 @@ def __init__(self, func):
             depends_on_dd_tracing_enabled = (
                 lambda original_boolean: config.trace_enabled and original_boolean
             )
-            self.make_inferred_span = depends_on_dd_tracing_enabled(
-                os.environ.get(DD_TRACE_MANAGED_SERVICES, "true").lower() == "true"
-            )
             self.encode_authorizer_context = depends_on_dd_tracing_enabled(
                 os.environ.get(DD_ENCODE_AUTHORIZER_CONTEXT, "true").lower() == "true"
             )
@@ -252,7 +249,7 @@ def _before(self, event, context):
 
             if config.trace_enabled:
                 set_dd_trace_py_root(trace_context_source, config.merge_xray_traces)
-                if self.make_inferred_span:
+                if config.make_inferred_span:
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, self.decode_authorizer_context
                     )

tests/test_config.py

@@ -109,6 +109,21 @@ def test_config_from_environ(env_key, conf_key, env_val, conf_val, setenv):
     assert getattr(config, conf_key) == conf_val
 
 
+_test_config_from_environ_depends_on_tracing = (
+    *_test_as_bool("DD_TRACE_MANAGED_SERVICES", "make_inferred_span", default=True),
+)
+
+@pytest.mark.parametrize('env_key,conf_key,env_val,conf_val', _test_config_from_environ_depends_on_tracing)
+@pytest.mark.parametrize('trace_enabled', [True, False])
+def test_config_from_environ_depends_on_tracing(env_key, conf_key, env_val, conf_val, setenv, trace_enabled):
+    setenv(env_key, env_val)
+    setenv("DD_TRACE_ENABLED", "true" if trace_enabled else "false")
+    if trace_enabled:
+        assert getattr(config, conf_key) is conf_val
+    else:
+        assert getattr(config, conf_key) is False
+
+
 _test_fips_mode_from_environ = (
     (None, None, False),
     (None, "", False),

tests/test_wrapper.py

@@ -506,6 +506,7 @@ def lambda_handler(event, context):
         self.assertEqual(os.environ.get("DD_REQUESTS_SERVICE_NAME"), "myAwesomeService")
         del os.environ["DD_SERVICE"]
 
+    @patch("datadog_lambda.config.Config.make_inferred_span", False)
     def test_encode_authorizer_span(self):
         @wrapper.datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -532,7 +533,6 @@ def lambda_handler(event, context):
         trace_ctx.sampling_priority = 1
         test_span.finish()
         lambda_handler.inferred_span = test_span
-        lambda_handler.make_inferred_span = False
         result = lambda_handler(lambda_event, lambda_context)
         raw_inject_data = result["context"]["_datadog"]
         self.assertIsInstance(raw_inject_data, str)
@@ -562,11 +562,9 @@ def return_type_test(event, context):
 class TestLambdaDecoratorSettings(unittest.TestCase):
     @patch("datadog_lambda.config.Config.trace_enabled", False)
     def test_some_envs_should_depend_on_dd_tracing_enabled(self):
-        os.environ[wrapper.DD_TRACE_MANAGED_SERVICES] = "true"
         os.environ[wrapper.DD_ENCODE_AUTHORIZER_CONTEXT] = "true"
         os.environ[wrapper.DD_DECODE_AUTHORIZER_CONTEXT] = "true"
         decorator = wrapper._LambdaDecorator(func=None)
-        self.assertFalse(decorator.make_inferred_span)
         self.assertFalse(decorator.encode_authorizer_context)
         self.assertFalse(decorator.decode_authorizer_context)