diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 92d0d55b..19ef8c04 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -345,6 +345,20 @@ def extract_context_custom_extractor(extractor, event, lambda_context): return None, None, None +def is_authorizer_response(response) -> bool: + try: + return ( + response is not None + and response["principalId"] + and response["policyDocument"] + ) + except KeyError: + pass + except Exception as e: + logger.debug("unknown error while checking is_authorizer_response %s", e) + return False + + def get_injected_authorizer_data(event, is_http_api) -> dict: try: authorizer_headers = event.get("requestContext", {}).get("authorizer") diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 9e246760..51ffccec 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -37,6 +37,7 @@ create_function_execution_span, create_inferred_span, InferredSpanInfo, + is_authorizer_response, ) from datadog_lambda.trigger import ( extract_trigger_tags, @@ -280,12 +281,7 @@ def _after(self, event, context): if should_use_extension: flush_extension() - if ( - self.encode_authorizer_context - and self.response - and self.response.get("principalId") - and self.response.get("policyDocument") - ): + if self.encode_authorizer_context and is_authorizer_response(self.response): self._inject_authorizer_span_headers( event.get("requestContext", {}).get("requestId") ) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 4efadd96..bde397c2 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -544,3 +544,17 @@ def lambda_handler(event, context): self.assertEquals(inject_data[TraceHeader.TRACE_ID], "456") self.assertEquals(inject_data[TraceHeader.SAMPLING_PRIORITY], "1") self.assertEquals(result["context"]["scope"], "still here") + + @patch("traceback.print_exc") + def test_different_return_type_no_error(self, MockPrintExc): + TEST_RESULTS = ["a str to return", 42, {"value": 42}, ["A", 42], None] + mock_context = get_mock_context() + for test_result in TEST_RESULTS: + + @datadog_lambda_wrapper + def return_type_test(event, context): + return test_result + + result = return_type_test({}, mock_context) + self.assertEquals(result, test_result) + self.assertFalse(MockPrintExc.called)