Repository: aws-powertools/powertools-lambda-python

Issue:
@tracer.capture_lambda_handler does not support functions that return generator expressions
<!--- Provide a general summary of the issue in the Title above -->
<!--- How has this issue affected you? What are you trying to accomplish? -->
I have an async lambda listening to S3 bucket notifications and i wanted to track the performance of downloading a large json file. Original this was a regular function that worked but later when refactoring it to yield the file pointer it failed when deployed. Nothing was found during unit testing.

**What were you trying to accomplish?**
Add a `@tracer.capture_method`  to a function that `yields` a result.

## Expected Behavior
<!--- If you're describing a bug, tell us what should happen -->
<!--- If you're suggesting a change/improvement, tell us how it should work -->
Either fail the same way during testing when `POWERTOOLS_TRACE_DISABLED='true'` or support generator expression

## Current Behavior
<!--- If describing a bug, tell us what happens instead of the expected behavior -->
<!--- If suggesting a change/improvement, explain the difference from current behavior -->
Either results in a `RuntimeError("generator didn't yield")` or a `Task timed out `

## Possible Solution
<!--- Not obligatory, but suggest a fix/reason for the bug, -->
<!--- or ideas how to implement the addition or change -->
Detect and support generator expression, or have a better emulation during testing which returns an error when a return
returns a generator expression.

## Steps to Reproduce (for bugs)
<!--- Provide a link to a live example, or an unambiguous set of steps to -->
<!--- reproduce this bug. Include code to reproduce, if relevant -->

```python3
import contextlib
import json
import tempfile
from typing import IO

import boto3
from aws_lambda_powertools import Tracer

bucket_name = "some_bucket"
key = "file.json"
tracer = Tracer()
s3 = boto3.client("s3")


@contextlib.contextmanager
def yield_without_capture() -> str:
    yield "[]"


@tracer.capture_method
@contextlib.contextmanager
def yield_with_capture() -> str:
    yield "[]"


@tracer.capture_method
@contextlib.contextmanager
def boto_yield_with_capture() -> IO:
    with tempfile.TemporaryFile() as fp:
        s3.download_fileobj(bucket_name, key, fp)
        print("download complete")  # Never prints
        fp.seek(0)
        yield fp


@contextlib.contextmanager
def boto_yield_without_capture() -> IO:
    with tempfile.TemporaryFile() as fp:
        s3.download_fileobj(bucket_name, key, fp)
        fp.seek(0)
        print("download complete")
        yield fp


@tracer.capture_lambda_handler
def lambda_handler(event: dict, _):
    if event["case"] == "yield_without_capture":
        with yield_without_capture() as yielded_value:
            result = yielded_value
    if event["case"] == "yield_with_capture":
        with yield_with_capture() as yielded_value:
            result = yielded_value
    if event["case"] == "boto_yield_without_capture":
        with boto_yield_without_capture() as fp:
            result = fp.read().decode("UTF-8")
    if event["case"] == "boto_yield_with_capture":
        with boto_yield_with_capture() as fp:
            result = fp.read().decode("UTF-8")  # "Task timed out after 10.01 seconds"

    return {"statusCode": 200, "result": json.loads(result)}

```

1. case `yield_without_capture` returns the expected value
2. case `yield_with_capture`  returns a `RuntimeError("generator didn't yield")`
3. case `boto_yield_without_capture` returns the expected value
4. case `boto_yield_with_capture` returns a `"Task timed out after 900.01 seconds"`

## Environment

* **Powertools version used**: 1.1.2
* **Packaging format (Layers, PyPi)**: PyPi
* **AWS Lambda function runtime:** Python 3.8



Discussion:
This should be fixed, version 1.2.0 is released which adds support for generators (including context managers). Please give it a try and either close the issue or let us know here if any issues persist.
@cakepietoast  thanks i will try it out and close this issue if this no longer has issues.
@cakepietoast  this is working better now, but i now have an issue with the context manager.

`boto_yield_without_capture` case keeps the file open until both with statements exit.

`boto_yield_with_capture` prematurely closes the file and fails to read from it

Example output: 

```json5
{
  "errorMessage": "read of closed file",
  "errorType": "ValueError",
  "stackTrace": [
    "  File \"/var/task/aws_lambda_powertools/tracing/tracer.py\", line 267, in decorate\n    response = lambda_handler(event, context)\n",
    "  File \"/var/task/tmp.py\", line 58, in lambda_handler\n    result = fp.read().decode(\"UTF-8\")  # ValueError(read of closed file)\n"
  ]
}
```
@cakepietoast  @heitorlessa  - i have a potential patch for this.

Resolving patch:
diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py
--- a/aws_lambda_powertools/tracing/tracer.py
+++ b/aws_lambda_powertools/tracing/tracer.py
@@ -487,14 +487,13 @@ def decorate(*args, **kwargs):
                     logger.debug(f"Calling method: {method_name}")
                     with method(*args, **kwargs) as return_val:
                         result = return_val
+                        yield result
                     self._add_response_as_metadata(function_name=method_name, data=result, subsegment=subsegment)
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
                     self._add_full_exception_as_metadata(function_name=method_name, error=err, subsegment=subsegment)
                     raise
 
-                yield result
-
         return decorate
 
     def _decorate_sync_function(self, method: Callable = None):


Test patch:
diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py
--- a/tests/unit/test_tracing.py
+++ b/tests/unit/test_tracing.py
@@ -383,6 +383,72 @@ def handler(event, context):
     assert "test result" in result
 
 
+def test_tracer_yield_from_context_manager_exception_metadata(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a context manager
+    # and the method raises an exception
+    @tracer.capture_method
+    @contextlib.contextmanager
+    def yield_with_capture():
+        yield "partial"
+        raise ValueError("test")
+
+    with pytest.raises(ValueError):
+        with yield_with_capture() as partial_val:
+            assert partial_val == "partial"
+
+    # THEN we should add the exception using method name as key plus error
+    # and their service name as the namespace
+    put_metadata_mock_args = in_subsegment_mock.put_metadata.call_args[1]
+    assert put_metadata_mock_args["key"] == "yield_with_capture error"
+    assert isinstance(put_metadata_mock_args["value"], ValueError)
+    assert put_metadata_mock_args["namespace"] == "booking"
+
+
+def test_tracer_yield_from_nested_context_manager(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a context manager nesting another context manager
+    class NestedContextManager(object):
+        def __enter__(self):
+            self._value = {"result": "test result"}
+            return self._value
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            self._value["result"] = "exit was called before yielding"
+
+    @tracer.capture_method
+    @contextlib.contextmanager
+    def yield_with_capture():
+        with NestedContextManager() as nested_context:
+            yield nested_context
+
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        response = []
+        with yield_with_capture() as yielded_value:
+            response.append(yielded_value["result"])
+
+        return response
+
+    result = handler({}, {})
+
+    # THEN we should have a subsegment named after the method name
+    # and add its response as trace metadata
+    handler_trace, yield_function_trace = in_subsegment_mock.in_subsegment.call_args_list
+
+    assert "test result" in in_subsegment_mock.put_metadata.call_args[1]["value"]
+    assert in_subsegment_mock.in_subsegment.call_count == 2
+    assert handler_trace == mocker.call(name="## handler")
+    assert yield_function_trace == mocker.call(name="## yield_with_capture")
+    assert "test result" in result
+
+
 def test_tracer_yield_from_generator(mocker, provider_stub, in_subsegment_mock):
     # GIVEN tracer is initialized
     provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
@@ -411,3 +477,28 @@ def handler(event, context):
     assert handler_trace == mocker.call(name="## handler")
     assert generator_fn_trace == mocker.call(name="## generator_fn")
     assert "test result" in result
+
+
+def test_tracer_yield_from_generator_exception_metadata(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a generator function
+    # and the method raises an exception
+    @tracer.capture_method
+    def generator_fn():
+        yield "partial"
+        raise ValueError("test")
+
+    with pytest.raises(ValueError):
+        gen = generator_fn()
+        list(gen)
+
+    # THEN we should add the exception using method name as key plus error
+    # and their service name as the namespace
+    put_metadata_mock_args = in_subsegment_mock.put_metadata.call_args[1]
+    assert put_metadata_mock_args["key"] == "generator_fn error"
+    assert put_metadata_mock_args["namespace"] == "booking"
+    assert isinstance(put_metadata_mock_args["value"], ValueError)
+    assert str(put_metadata_mock_args["value"]) == "test"