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.

Resolving patch:
diff --git a/aws_lambda_powertools/utilities/parameters/base.py b/aws_lambda_powertools/utilities/parameters/base.py
--- a/aws_lambda_powertools/utilities/parameters/base.py
+++ b/aws_lambda_powertools/utilities/parameters/base.py
@@ -7,7 +7,7 @@
 from abc import ABC, abstractmethod
 from collections import namedtuple
 from datetime import datetime, timedelta
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Tuple, Union
 
 from .exceptions import GetParameterError, TransformParameterError
 
@@ -31,6 +31,9 @@ def __init__(self):
 
         self.store = {}
 
+    def _has_not_expired(self, key: Tuple[str, Optional[str]]) -> bool:
+        return key in self.store and self.store[key].ttl >= datetime.now()
+
     def get(
         self, name: str, max_age: int = DEFAULT_MAX_AGE_SECS, transform: Optional[str] = None, **sdk_options
     ) -> Union[str, list, dict, bytes]:
@@ -70,24 +73,26 @@ def get(
         # an acceptable tradeoff.
         key = (name, transform)
 
-        if key not in self.store or self.store[key].ttl < datetime.now():
-            try:
-                value = self._get(name, **sdk_options)
-            # Encapsulate all errors into a generic GetParameterError
-            except Exception as exc:
-                raise GetParameterError(str(exc))
+        if self._has_not_expired(key):
+            return self.store[key].value
+
+        try:
+            value = self._get(name, **sdk_options)
+        # Encapsulate all errors into a generic GetParameterError
+        except Exception as exc:
+            raise GetParameterError(str(exc))
 
-            if transform is not None:
-                value = transform_value(value, transform)
+        if transform is not None:
+            value = transform_value(value, transform)
 
-            self.store[key] = ExpirableValue(value, datetime.now() + timedelta(seconds=max_age),)
+        self.store[key] = ExpirableValue(value, datetime.now() + timedelta(seconds=max_age),)
 
-        return self.store[key].value
+        return value
 
     @abstractmethod
     def _get(self, name: str, **sdk_options) -> str:
         """
-        Retrieve paramater value from the underlying parameter store
+        Retrieve parameter value from the underlying parameter store
         """
         raise NotImplementedError()
 
@@ -129,29 +134,22 @@ def get_multiple(
 
         key = (path, transform)
 
-        if key not in self.store or self.store[key].ttl < datetime.now():
-            try:
-                values = self._get_multiple(path, **sdk_options)
-            # Encapsulate all errors into a generic GetParameterError
-            except Exception as exc:
-                raise GetParameterError(str(exc))
+        if self._has_not_expired(key):
+            return self.store[key].value
 
-            if transform is not None:
-                new_values = {}
-                for key, value in values.items():
-                    try:
-                        new_values[key] = transform_value(value, transform)
-                    except Exception as exc:
-                        if raise_on_transform_error:
-                            raise exc
-                        else:
-                            new_values[key] = None
+        try:
+            values: Dict[str, Union[str, bytes, dict, None]] = self._get_multiple(path, **sdk_options)
+        # Encapsulate all errors into a generic GetParameterError
+        except Exception as exc:
+            raise GetParameterError(str(exc))
 
-                values = new_values
+        if transform is not None:
+            for (key, value) in values.items():
+                values[key] = transform_value(value, transform, raise_on_transform_error)
 
-            self.store[key] = ExpirableValue(values, datetime.now() + timedelta(seconds=max_age),)
+        self.store[key] = ExpirableValue(values, datetime.now() + timedelta(seconds=max_age),)
 
-        return self.store[key].value
+        return values
 
     @abstractmethod
     def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
@@ -161,16 +159,19 @@ def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
         raise NotImplementedError()
 
 
-def transform_value(value: str, transform: str) -> Union[dict, bytes]:
+def transform_value(value: str, transform: str, raise_on_transform_error: bool = True) -> Union[dict, bytes, None]:
     """
     Apply a transform to a value
 
     Parameters
     ---------
     value: str
-        Parameter alue to transform
+        Parameter value to transform
     transform: str
         Type of transform, supported values are "json" and "binary"
+    raise_on_transform_error: bool, optional
+        Raises an exception if any transform fails, otherwise this will
+        return a None value for each transform that failed
 
     Raises
     ------
@@ -187,4 +188,6 @@ def transform_value(value: str, transform: str) -> Union[dict, bytes]:
             raise ValueError(f"Invalid transform type '{transform}'")
 
     except Exception as exc:
-        raise TransformParameterError(str(exc))
+        if raise_on_transform_error:
+            raise TransformParameterError(str(exc))
+        return None
diff --git a/aws_lambda_powertools/utilities/parameters/secrets.py b/aws_lambda_powertools/utilities/parameters/secrets.py
--- a/aws_lambda_powertools/utilities/parameters/secrets.py
+++ b/aws_lambda_powertools/utilities/parameters/secrets.py
@@ -77,7 +77,7 @@ def _get(self, name: str, **sdk_options) -> str:
         ----------
         name: str
             Name of the parameter
-        sdk_options: dict
+        sdk_options: dict, optional
             Dictionary of options that will be passed to the Secrets Manager get_secret_value API call
         """
 
diff --git a/aws_lambda_powertools/utilities/parameters/ssm.py b/aws_lambda_powertools/utilities/parameters/ssm.py
--- a/aws_lambda_powertools/utilities/parameters/ssm.py
+++ b/aws_lambda_powertools/utilities/parameters/ssm.py
@@ -8,7 +8,7 @@
 import boto3
 from botocore.config import Config
 
-from .base import DEFAULT_PROVIDERS, BaseProvider
+from .base import DEFAULT_MAX_AGE_SECS, DEFAULT_PROVIDERS, BaseProvider
 
 
 class SSMProvider(BaseProvider):
@@ -86,6 +86,46 @@ def __init__(
 
         super().__init__()
 
+    def get(
+        self,
+        name: str,
+        max_age: int = DEFAULT_MAX_AGE_SECS,
+        transform: Optional[str] = None,
+        decrypt: bool = False,
+        **sdk_options
+    ) -> Union[str, list, dict, bytes]:
+        """
+        Retrieve a parameter value or return the cached value
+
+        Parameters
+        ----------
+        name: str
+            Parameter name
+        max_age: int
+            Maximum age of the cached value
+        transform: str
+            Optional transformation of the parameter value. Supported values
+            are "json" for JSON strings and "binary" for base 64 encoded
+            values.
+        decrypt: bool, optional
+            If the parameter value should be decrypted
+        sdk_options: dict, optional
+            Arguments that will be passed directly to the underlying API call
+
+        Raises
+        ------
+        GetParameterError
+            When the parameter provider fails to retrieve a parameter value for
+            a given name.
+        TransformParameterError
+            When the parameter provider fails to transform a parameter value.
+        """
+
+        # Add to `decrypt` sdk_options to we can have an explicit option for this
+        sdk_options["decrypt"] = decrypt
+
+        return super().get(name, max_age, transform, **sdk_options)
+
     def _get(self, name: str, decrypt: bool = False, **sdk_options) -> str:
         """
         Retrieve a parameter value from AWS Systems Manager Parameter Store
@@ -144,7 +184,9 @@ def _get_multiple(self, path: str, decrypt: bool = False, recursive: bool = Fals
         return parameters
 
 
-def get_parameter(name: str, transform: Optional[str] = None, **sdk_options) -> Union[str, list, dict, bytes]:
+def get_parameter(
+    name: str, transform: Optional[str] = None, decrypt: bool = False, **sdk_options
+) -> Union[str, list, dict, bytes]:
     """
     Retrieve a parameter value from AWS Systems Manager (SSM) Parameter Store
 
@@ -154,6 +196,8 @@ def get_parameter(name: str, transform: Optional[str] = None, **sdk_options) ->
         Name of the parameter
     transform: str, optional
         Transforms the content from a JSON object ('json') or base64 binary string ('binary')
+    decrypt: bool, optional
+        If the parameter values should be decrypted
     sdk_options: dict, optional
         Dictionary of options that will be passed to the Parameter Store get_parameter API call
 
@@ -190,7 +234,10 @@ def get_parameter(name: str, transform: Optional[str] = None, **sdk_options) ->
     if "ssm" not in DEFAULT_PROVIDERS:
         DEFAULT_PROVIDERS["ssm"] = SSMProvider()
 
-    return DEFAULT_PROVIDERS["ssm"].get(name, transform=transform)
+    # Add to `decrypt` sdk_options to we can have an explicit option for this
+    sdk_options["decrypt"] = decrypt
+
+    return DEFAULT_PROVIDERS["ssm"].get(name, transform=transform, **sdk_options)
 
 
 def get_parameters(
@@ -205,10 +252,10 @@ def get_parameters(
         Path to retrieve the parameters
     transform: str, optional
         Transforms the content from a JSON object ('json') or base64 binary string ('binary')
-    decrypt: bool, optional
-        If the parameter values should be decrypted
     recursive: bool, optional
         If this should retrieve the parameter values recursively or not, defaults to True
+    decrypt: bool, optional
+        If the parameter values should be decrypted
     sdk_options: dict, optional
         Dictionary of options that will be passed to the Parameter Store get_parameters_by_path API call
 
@@ -245,4 +292,7 @@ def get_parameters(
     if "ssm" not in DEFAULT_PROVIDERS:
         DEFAULT_PROVIDERS["ssm"] = SSMProvider()
 
-    return DEFAULT_PROVIDERS["ssm"].get_multiple(path, transform=transform, recursive=recursive, decrypt=decrypt)
+    sdk_options["recursive"] = recursive
+    sdk_options["decrypt"] = decrypt
+
+    return DEFAULT_PROVIDERS["ssm"].get_multiple(path, transform=transform, **sdk_options)


Test patch:
diff --git a/tests/functional/test_utilities_parameters.py b/tests/functional/test_utilities_parameters.py
--- a/tests/functional/test_utilities_parameters.py
+++ b/tests/functional/test_utilities_parameters.py
@@ -1310,6 +1310,7 @@ def test_get_parameter_new(monkeypatch, mock_name, mock_value):
     class TestProvider(BaseProvider):
         def _get(self, name: str, **kwargs) -> str:
             assert name == mock_name
+            assert not kwargs["decrypt"]
             return mock_value
 
         def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
@@ -1355,6 +1356,8 @@ def _get(self, name: str, **kwargs) -> str:
 
         def _get_multiple(self, path: str, **kwargs) -> Dict[str, str]:
             assert path == mock_name
+            assert kwargs["recursive"]
+            assert not kwargs["decrypt"]
             return mock_value
 
     monkeypatch.setattr(parameters.ssm, "DEFAULT_PROVIDERS", {})
@@ -1468,3 +1471,13 @@ def test_transform_value_wrong(mock_value):
         parameters.base.transform_value(mock_value, "INCORRECT")
 
     assert "Invalid transform type" in str(excinfo)
+
+
+def test_transform_value_ignore_error(mock_value):
+    """
+    Test transform_value() does not raise errors when raise_on_transform_error is False
+    """
+
+    value = parameters.base.transform_value(mock_value, "INCORRECT", raise_on_transform_error=False)
+
+    assert value is None