diff --git a/ddtrace/contrib/aws_lambda/patch.py b/ddtrace/contrib/aws_lambda/patch.py index bd1a9f476ee..1253593fb0f 100644 --- a/ddtrace/contrib/aws_lambda/patch.py +++ b/ddtrace/contrib/aws_lambda/patch.py @@ -100,8 +100,24 @@ def __call__(self, func, args, kwargs): finally: self._after() + def _set_context(self, args, kwargs): + """Sets the context attribute.""" + # The context is the second argument in a handler + # signature and it is always sent. + # + # note: AWS Lambda context is an object, the event is a dict. + # `get_remaining_time_in_millis` is guaranteed to be + # present in the context. + _context = get_argument_value(args, kwargs, 1, "context") + if hasattr(_context, "get_remaining_time_in_millis"): + self.context = _context + else: + # Handler was possibly manually wrapped, and the first + # argument is the `datadog-lambda` decorator object. + self.context = get_argument_value(args, kwargs, 2, "context") + def _before(self, args, kwargs): - self.context = get_argument_value(args, kwargs, -1, "context") + self._set_context(args, kwargs) self.timeoutChannel = TimeoutChannel(self.context) self.timeoutChannel._start() diff --git a/releasenotes/notes/fix-aws-lambda-get-argument-value-position-01ffab721cd2e6b7.yaml b/releasenotes/notes/fix-aws-lambda-get-argument-value-position-01ffab721cd2e6b7.yaml new file mode 100644 index 00000000000..c847082d01b --- /dev/null +++ b/releasenotes/notes/fix-aws-lambda-get-argument-value-position-01ffab721cd2e6b7.yaml @@ -0,0 +1,4 @@ +--- +fixes: + - | + aws_lambda: Fix AttributeError raised when extracting context from arguments. diff --git a/tests/contrib/aws_lambda/handlers.py b/tests/contrib/aws_lambda/handlers.py index 224340d8e47..060088f9893 100644 --- a/tests/contrib/aws_lambda/handlers.py +++ b/tests/contrib/aws_lambda/handlers.py @@ -6,16 +6,16 @@ @datadog_lambda_wrapper -def manually_wrapped_handler(event, context): +def manually_wrapped_handler(event, context, success=True): @tracer.wrap("result-trace") def get_result(): - return {"success": True} + return {"success": success} return get_result() -def handler(event, context): - return {"success": True} +def handler(event, context, success=True): + return {"success": success} def timeout_handler(event, context):