fix(aws_lambda): update how context is extracted (#5461)

Handlers are allowed to have more than the two required positional
arguments.

```python
def handler(event, context, extra_arg, optional_arg=False):
    ...
    return ...
```

Our current `context` extraction checks for the last element. The patch
didn't take into account that more arguments can be provided besides the
two required. This PR fixes that.
It is guaranteed that the `(event, context)` arguments will be in the
first positions.

## Testing strategy

- Updates tests to have extra arguments
- Tested with an actual AWS Lambda too.


## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] PR description includes explicit acknowledgement/acceptance of the
performance implications of this PR as reported in the benchmarks PR
comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

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()

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.

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):