From 05c1aa57e334a7a415a93c6d450b594788517d33 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 28 Mar 2025 00:08:27 -0400 Subject: [PATCH 1/5] lazy load metrics also hashlib right away --- datadog_lambda/tracing.py | 2 +- datadog_lambda/wrapper.py | 17 ++++++++++------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 0fae76dd..43636789 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0. # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import hashlib import logging import os import base64 @@ -387,6 +386,7 @@ def extract_context_from_kinesis_event(event, lambda_context): def _deterministic_sha256_hash(s: str, part: str) -> int: + import hashlib sha256_hash = hashlib.sha256(s.encode()).hexdigest() # First two chars is '0b'. zfill to ensure 256 bits, but we only care about the first 128 bits binary_hash = bin(int(sha256_hash, 16))[2:].zfill(256) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 5641bd15..1ae92285 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -23,11 +23,6 @@ XraySubsegment, Headers, ) -from datadog_lambda.metric import ( - flush_stats, - submit_invocations_metric, - submit_errors_metric, -) from datadog_lambda.module_name import modify_module_name from datadog_lambda.patch import patch_all from datadog_lambda.span_pointers import calculate_span_pointers @@ -248,7 +243,10 @@ def __call__(self, event, context, **kwargs): self.response = self.func(event, context, **kwargs) return self.response except Exception: - submit_errors_metric(context) + if not should_use_extension: + from datadog_lambda.metric import submit_invocations_metric + submit_invocations_metric(context) + if self.span: self.span.set_traceback() raise @@ -284,7 +282,11 @@ def _before(self, event, context): try: self.response = None set_cold_start(init_timestamp_ns) - submit_invocations_metric(context) + + if not should_use_extension: + from datadog_lambda.metric import submit_invocations_metric + submit_invocations_metric(context) + self.trigger_tags = extract_trigger_tags(event, context) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( @@ -383,6 +385,7 @@ def _after(self, event, context): logger.debug("Failed to create cold start spans. %s", e) if not self.flush_to_log or should_use_extension: + from datadog_lambda.metric import flush_stats flush_stats(context) if should_use_extension and self.local_testing_mode: # when testing locally, the extension does not know when an From 246f2b4afa8f539d8ce7dfc1f7eb0bc6da0a86bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 28 Mar 2025 00:08:46 -0400 Subject: [PATCH 2/5] black --- datadog_lambda/tracing.py | 1 + datadog_lambda/wrapper.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 43636789..a787bb12 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -387,6 +387,7 @@ def extract_context_from_kinesis_event(event, lambda_context): def _deterministic_sha256_hash(s: str, part: str) -> int: import hashlib + sha256_hash = hashlib.sha256(s.encode()).hexdigest() # First two chars is '0b'. zfill to ensure 256 bits, but we only care about the first 128 bits binary_hash = bin(int(sha256_hash, 16))[2:].zfill(256) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 1ae92285..0b17860d 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -245,6 +245,7 @@ def __call__(self, event, context, **kwargs): except Exception: if not should_use_extension: from datadog_lambda.metric import submit_invocations_metric + submit_invocations_metric(context) if self.span: @@ -282,9 +283,10 @@ def _before(self, event, context): try: self.response = None set_cold_start(init_timestamp_ns) - + if not should_use_extension: from datadog_lambda.metric import submit_invocations_metric + submit_invocations_metric(context) self.trigger_tags = extract_trigger_tags(event, context) @@ -386,6 +388,7 @@ def _after(self, event, context): if not self.flush_to_log or should_use_extension: from datadog_lambda.metric import flush_stats + flush_stats(context) if should_use_extension and self.local_testing_mode: # when testing locally, the extension does not know when an From b922ee62ca67eacdd03a4e8a2274aa2a8478a617 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 28 Mar 2025 12:13:32 -0400 Subject: [PATCH 3/5] update invocation metric to be error metric --- datadog_lambda/wrapper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 0b17860d..5a87f07f 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -244,9 +244,9 @@ def __call__(self, event, context, **kwargs): return self.response except Exception: if not should_use_extension: - from datadog_lambda.metric import submit_invocations_metric + from datadog_lambda.metric import submit_errors_metric - submit_invocations_metric(context) + submit_errors_metric(context) if self.span: self.span.set_traceback() From 96d3df0b620e3d69260bfa3bed14c50c0edaaef9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 28 Mar 2025 12:22:38 -0400 Subject: [PATCH 4/5] lazyload `base64` --- datadog_lambda/api.py | 2 +- datadog_lambda/tracing.py | 7 ++++++- datadog_lambda/trigger.py | 3 ++- datadog_lambda/wrapper.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index c539ea05..fd3e2c17 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -1,6 +1,5 @@ import os import logging -import base64 logger = logging.getLogger(__name__) KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" @@ -9,6 +8,7 @@ def decrypt_kms_api_key(kms_client, ciphertext): from botocore.exceptions import ClientError + import base64 """ Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS. diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index a787bb12..9a27673c 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -4,7 +4,6 @@ # Copyright 2019 Datadog, Inc. import logging import os -import base64 import traceback import ujson as json from datetime import datetime, timezone @@ -258,6 +257,8 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): dd_json_data = None dd_json_data_type = dd_payload.get("Type") or dd_payload.get("dataType") if dd_json_data_type == "Binary": + import base64 + dd_json_data = dd_payload.get("binaryValue") or dd_payload.get("Value") if dd_json_data: dd_json_data = base64.b64decode(dd_json_data) @@ -372,6 +373,8 @@ def extract_context_from_kinesis_event(event, lambda_context): return extract_context_from_lambda_context(lambda_context) data = kinesis.get("data") if data: + import base64 + b64_bytes = data.encode("ascii") str_bytes = base64.b64decode(b64_bytes) data_str = str_bytes.decode("ascii") @@ -552,6 +555,8 @@ def get_injected_authorizer_data(event, is_http_api) -> dict: if not dd_data_raw: return None + import base64 + injected_data = json.loads(base64.b64decode(dd_data_raw)) # Lambda authorizer's results can be cached. But the payload will still have the injected diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index 708138bf..8090e36e 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -3,7 +3,6 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import base64 import gzip import ujson as json from io import BytesIO, BufferedReader @@ -242,6 +241,8 @@ def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> s # e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz if source.event_type == EventTypes.CLOUDWATCH_LOGS: + import base64 + with gzip.GzipFile( fileobj=BytesIO(base64.b64decode(event.get("awslogs", {}).get("data"))) ) as decompress_stream: diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 5a87f07f..8c1914e3 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0. # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import base64 import os import logging import traceback @@ -273,6 +272,9 @@ def _inject_authorizer_span_headers(self, request_id): injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns if request_id is not None: injected_headers[Headers.Authorizing_Request_Id] = request_id + + import base64 + datadog_data = base64.b64encode( json.dumps(injected_headers, escape_forward_slashes=False).encode() ).decode() From 6b0eed81b57f88789119e3791b112b445a476835 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 28 Mar 2025 12:27:53 -0400 Subject: [PATCH 5/5] patch right call --- tests/test_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index f47285e6..4b243036 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -470,7 +470,7 @@ def lambda_handler(event, context): self.mock_write_metric_point_to_stdout.assert_not_called() def test_only_one_wrapper_in_use(self): - patcher = patch("datadog_lambda.wrapper.submit_invocations_metric") + patcher = patch("datadog_lambda.metric.submit_invocations_metric") self.mock_submit_invocations_metric = patcher.start() self.addCleanup(patcher.stop)