@tracer.capture_lambda_handler does not support functions that return generator expressions

[michaelbrewer]

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

Either fail the same way during testing when POWERTOOLS_TRACE_DISABLED='true' or support generator expression

Current Behavior

Either results in a RuntimeError("generator didn't yield") or a Task timed out

Possible Solution

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)

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

[to-mc]

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.

[michaelbrewer]

@cakepietoast thanks i will try it out and close this issue if this no longer has issues.

[michaelbrewer]

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

{
  "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"
  ]
}

[michaelbrewer]

@cakepietoast @heitorlessa - i have a potential patch for this.

[michaelbrewer]

@cakepietoast @heitorlessa - this PR might resolve this: #124

[to-mc]

Thanks for spotting this @michaelbrewer! Your PR looks good to me, I'll get a patch version out soon with your changes.

[to-mc]

1.3.1 is published with your changes, please let us know if this issue can be closed.