fix: use DD_SERVICE for inferred spans when integration service names are removed

datadog_lambda/config.py

@@ -75,6 +75,13 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False):
     add_span_pointers = _get_env("DD_BOTOCORE_ADD_SPAN_POINTERS", "true", as_bool)
     trace_extractor = _get_env("DD_TRACE_EXTRACTOR")
 
+    aws_service_representation_enabled = _get_env(
+        "DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED", "true", as_bool
+    )
+    remove_integration_service_names_enabled = _get_env(
+        "DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED", "false", as_bool
+    )
+
     enhanced_metrics_enabled = _get_env("DD_ENHANCED_METRICS", "true", as_bool)
 
     flush_in_thread = _get_env("DD_FLUSH_IN_THREAD", "false", as_bool)

datadog_lambda/tracing.py

@@ -900,11 +900,14 @@ def determine_service_name(
     if mapped_service:
         return mapped_service
 
+    # When integration service names are removed, inferred (synthetic) spans use
+    # the base service name (DD_SERVICE) instead of the AWS resource/instance
+    # representation.
+    if config.remove_integration_service_names_enabled and config.service:
+        return config.service
+
     # Check if AWS service representation is disabled
-    aws_service_representation = os.environ.get(
-        "DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED", ""
-    ).lower()
-    if aws_service_representation in ("false", "0"):
+    if not config.aws_service_representation_enabled:
         return fallback
 
     # Use extracted_key if it exists and is not empty, otherwise use fallback

tests/test_tracing.py

@@ -1421,6 +1421,13 @@ def test_set_dd_trace_py_root_none_context(self):
 class TestServiceMapping(unittest.TestCase):
     def setUp(self):
         self.service_mapping = {}
+        # These tests exercise the AWS service-representation / service-mapping
+        # resolution, which only applies when DD_SERVICE is not set. Pin
+        # config.service to None so the tests are deterministic regardless of
+        # whether DD_SERVICE is present in the environment (e.g. in CI).
+        service_patcher = patch("datadog_lambda.config.Config.service", None)
+        service_patcher.start()
+        self.addCleanup(service_patcher.stop)
 
     def get_service_mapping(self):
         return global_service_mapping
@@ -1467,6 +1474,7 @@ def test_set_service_mapping(self):
         self.set_service_mapping(new_service_mapping)
         self.assertEqual(self.get_service_mapping(), new_service_mapping)
 
+    @patch("datadog_lambda.config.Config.service", None)
     def test_determine_service_name(self):
         # Prepare the environment
         os.environ["DD_SERVICE_MAPPING"] = "api1:service1,api2:service2"
@@ -1501,44 +1509,111 @@ def test_determine_service_name(self):
             "default",
         )
 
-        # Test with DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED set to false
-        os.environ["DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED"] = "false"
-        self.assertEqual(
-            determine_service_name(
-                self.get_service_mapping(), "api4", "api4", "extracted", "fallback"
-            ),
-            "fallback",
-        )
-
-        # Test with DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED set to 0
-        os.environ["DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED"] = "0"
-        self.assertEqual(
-            determine_service_name(
-                self.get_service_mapping(), "api4", "api4", "extracted", "fallback"
-            ),
-            "fallback",
-        )
+        # Test with DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED disabled
+        with patch(
+            "datadog_lambda.config.Config.aws_service_representation_enabled", False
+        ):
+            self.assertEqual(
+                determine_service_name(
+                    self.get_service_mapping(),
+                    "api4",
+                    "api4",
+                    "extracted",
+                    "fallback",
+                ),
+                "fallback",
+            )
 
-        # Test with DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED not set (default behavior)
-        if "DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED" in os.environ:
-            del os.environ["DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED"]
-        self.assertEqual(
-            determine_service_name(
-                self.get_service_mapping(), "api4", "api4", "extracted", "fallback"
-            ),
-            "extracted",
-        )
+        # Test with DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED enabled (default)
+        with patch(
+            "datadog_lambda.config.Config.aws_service_representation_enabled", True
+        ):
+            self.assertEqual(
+                determine_service_name(
+                    self.get_service_mapping(),
+                    "api4",
+                    "api4",
+                    "extracted",
+                    "fallback",
+                ),
+                "extracted",
+            )
 
-        # Test with empty extracted key
-        self.assertEqual(
-            determine_service_name(
-                self.get_service_mapping(), "api4", "api4", "  ", "fallback"
-            ),
-            "fallback",
-        )
+            # Test with empty extracted key
+            self.assertEqual(
+                determine_service_name(
+                    self.get_service_mapping(), "api4", "api4", "  ", "fallback"
+                ),
+                "fallback",
+            )
 
         del os.environ["DD_SERVICE_MAPPING"]
 
+    def test_determine_service_name_with_remove_integration_flag(self):
+        # By default (flag off), inferred spans use the AWS resource name even
+        # when DD_SERVICE is set.
+        with patch("datadog_lambda.config.Config.service", "my-service"):
+            self.assertEqual(
+                determine_service_name(
+                    {}, "queue-name", "lambda_sqs", "queue-name", "sqs"
+                ),
+                "queue-name",
+            )
+
+            # With DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED, inferred
+            # spans use DD_SERVICE instead of the AWS resource name.
+            with patch(
+                "datadog_lambda.config.Config."
+                "remove_integration_service_names_enabled",
+                True,
+            ):
+                self.assertEqual(
+                    determine_service_name(
+                        {}, "queue-name", "lambda_sqs", "queue-name", "sqs"
+                    ),
+                    "my-service",
+                )
+
+                # An explicit service mapping still wins over DD_SERVICE.
+                self.assertEqual(
+                    determine_service_name(
+                        {"lambda_sqs": "mapped-service"},
+                        "queue-name",
+                        "lambda_sqs",
+                        "queue-name",
+                        "sqs",
+                    ),
+                    "mapped-service",
+                )
+
+                # DD_SERVICE is used even when AWS service representation is
+                # disabled.
+                with patch(
+                    "datadog_lambda.config.Config."
+                    "aws_service_representation_enabled",
+                    False,
+                ):
+                    self.assertEqual(
+                        determine_service_name(
+                            {}, "queue-name", "lambda_sqs", "queue-name", "sqs"
+                        ),
+                        "my-service",
+                    )
+
+        # When DD_SERVICE is not set, the flag has no effect (resource name).
+        with patch("datadog_lambda.config.Config.service", None):
+            with patch(
+                "datadog_lambda.config.Config."
+                "remove_integration_service_names_enabled",
+                True,
+            ):
+                self.assertEqual(
+                    determine_service_name(
+                        {}, "queue-name", "lambda_sqs", "queue-name", "sqs"
+                    ),
+                    "queue-name",
+                )
+
     def test_remaps_all_inferred_span_service_names_from_api_gateway_event(self):
         new_service_mapping = {"lambda_api_gateway": "new-name"}
         self.set_service_mapping(new_service_mapping)
@@ -1727,6 +1802,27 @@ def test_remaps_specific_inferred_span_service_names_from_sqs_event(self):
         self.assertEqual(span2.get_tag("operation_name"), "aws.sqs")
         self.assertEqual(span2.service, "different-sqs-url")
 
+    def test_create_inferred_span_uses_dd_service_with_remove_integration_flag(
+        self,
+    ):
+        # With DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED and DD_SERVICE
+        # set, inferred spans use DD_SERVICE instead of the AWS resource name.
+        event_sample_source = "sqs-string-msg-attribute"
+        test_file = event_samples + event_sample_source + ".json"
+        with open(test_file, "r") as event:
+            original_event = json.load(event)
+
+        ctx = get_mock_context()
+        ctx.aws_request_id = "123"
+
+        with patch("datadog_lambda.config.Config.service", "my-dd-service"), patch(
+            "datadog_lambda.config.Config." "remove_integration_service_names_enabled",
+            True,
+        ):
+            span = create_inferred_span(original_event, ctx)
+        self.assertEqual(span.get_tag("operation_name"), "aws.sqs")
+        self.assertEqual(span.service, "my-dd-service")
+
     def test_remaps_all_inferred_span_service_names_from_sns_event(self):
         self.set_service_mapping({"lambda_sns": "new-name"})
         event_sample_source = "sns-string-msg-attribute"
@@ -2555,6 +2651,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
 
 @pytest.mark.parametrize("source,expect", _test_create_inferred_span)
 @patch("ddtrace.trace.Span.finish", autospec=True)
+@patch("datadog_lambda.config.Config.service", None)
 def test_create_inferred_span(mock_span_finish, source, expect):
     with open(f"{event_samples}{source}.json") as f:
         event = json.load(f)