diff --git a/datadog_lambda/span_pointers.py b/datadog_lambda/span_pointers.py index d5c8bde7..86c9e60b 100644 --- a/datadog_lambda/span_pointers.py +++ b/datadog_lambda/span_pointers.py @@ -1,5 +1,6 @@ from itertools import chain import logging +import os from typing import List from ddtrace._trace._span_pointer import _SpanPointerDirection @@ -10,16 +11,23 @@ logger = logging.getLogger(__name__) +dd_botocore_add_span_pointers = os.environ.get( + "DD_BOTOCORE_ADD_SPAN_POINTERS", "true" +).lower() in ("true", "1") + + def calculate_span_pointers( event_source, event, + botocore_add_span_pointers=dd_botocore_add_span_pointers, ) -> List[_SpanPointerDescription]: try: - if event_source.equals(EventTypes.S3): - return _calculate_s3_span_pointers_for_event(event) + if botocore_add_span_pointers: + if event_source.equals(EventTypes.S3): + return _calculate_s3_span_pointers_for_event(event) - elif event_source.equals(EventTypes.DYNAMODB): - return _calculate_dynamodb_span_pointers_for_event(event) + elif event_source.equals(EventTypes.DYNAMODB): + return _calculate_dynamodb_span_pointers_for_event(event) except Exception as e: logger.warning( diff --git a/tests/test_span_pointers.py b/tests/test_span_pointers.py index 8073a704..f4cd1b26 100644 --- a/tests/test_span_pointers.py +++ b/tests/test_span_pointers.py @@ -14,6 +14,7 @@ class SpanPointersCase(NamedTuple): name: str event_source: _EventSource event: dict + dd_botocore_span_pointers: bool span_pointers: List[_SpanPointerDescription] @pytest.mark.parametrize( @@ -23,12 +24,14 @@ class SpanPointersCase(NamedTuple): name="some unsupported event", event_source=_EventSource(EventTypes.UNKNOWN), event={}, + dd_botocore_span_pointers=True, span_pointers=[], ), SpanPointersCase( name="empty s3 event", event_source=_EventSource(EventTypes.S3), event={}, + dd_botocore_span_pointers=True, span_pointers=[], ), SpanPointersCase( @@ -50,6 +53,7 @@ class SpanPointersCase(NamedTuple): }, ], }, + dd_botocore_span_pointers=True, span_pointers=[ _SpanPointerDescription( pointer_kind="aws.s3.object", @@ -59,6 +63,28 @@ class SpanPointersCase(NamedTuple): ), ], ), + SpanPointersCase( + name="sensible s3 event with dd_botocore_span_pointers disabled", + event_source=_EventSource(EventTypes.S3), + event={ + "Records": [ + { + "eventName": "ObjectCreated:Put", + "s3": { + "bucket": { + "name": "mybucket", + }, + "object": { + "key": "mykey", + "eTag": "123abc", + }, + }, + }, + ], + }, + dd_botocore_span_pointers=False, + span_pointers=[], + ), SpanPointersCase( name="malformed s3 event", event_source=_EventSource(EventTypes.S3), @@ -78,12 +104,14 @@ class SpanPointersCase(NamedTuple): }, ], }, + dd_botocore_span_pointers=True, span_pointers=[], ), SpanPointersCase( name="empty dynamodb event", event_source=_EventSource(EventTypes.DYNAMODB), event={}, + dd_botocore_span_pointers=True, span_pointers=[], ), SpanPointersCase( @@ -109,6 +137,7 @@ class SpanPointersCase(NamedTuple): }, ], }, + dd_botocore_span_pointers=True, span_pointers=[ _SpanPointerDescription( pointer_kind="aws.dynamodb.item", @@ -139,16 +168,18 @@ class SpanPointersCase(NamedTuple): }, ], }, + dd_botocore_span_pointers=True, span_pointers=[], ), ], ids=lambda test_case: test_case.name, ) - def test_calculate_span_pointers(self, test_case: SpanPointersCase): + def test_calculate_span_pointers(self, test_case: SpanPointersCase) -> None: assert ( calculate_span_pointers( test_case.event_source, test_case.event, + botocore_add_span_pointers=test_case.dd_botocore_span_pointers, ) == test_case.span_pointers )