diff --git a/datadog_lambda/config.py b/datadog_lambda/config.py index ce4924af..b92c37ee 100644 --- a/datadog_lambda/config.py +++ b/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) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index b3f79a96..893df2c4 100644 --- a/datadog_lambda/tracing.py +++ b/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 diff --git a/tests/test_tracing.py b/tests/test_tracing.py index fc18f6e5..f0579da7 100644 --- a/tests/test_tracing.py +++ b/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)