chore: lazy load some imports (#581)

* lazy load metrics

also hashlib right away

* black

* update invocation metric to be error metric

* lazyload `base64`

* patch right call

Author: duncanista

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.

datadog_lambda/tracing.py

@@ -2,10 +2,8 @@
 # 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
 import traceback
 import ujson as json
 from datetime import datetime, timezone
@@ -259,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)
@@ -373,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")
@@ -387,6 +389,8 @@ 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)
@@ -551,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

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:

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
@@ -23,11 +22,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 +242,11 @@ 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_errors_metric
+
+                submit_errors_metric(context)
+
             if self.span:
                 self.span.set_traceback()
             raise
@@ -274,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()
@@ -284,7 +285,12 @@ 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 +389,8 @@ 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

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)