feat(event_source): add Kinesis Firehose Data Transformation data class (#3029)

* support kinesis response

* fix lint, address Leandro suggestions

* remove deleted const

* fix Literal import in 3.7

* change to use data-classes

* fix mypy

* fix typo, make asdict a function

* address Troy/Leandro suggestions

* remove 6MB comment

* fix comments

* address Heitor's suggestion

* data class default optimization

* remove slot for static check

* fix doc, example

* rename r->record

* Addressing Heitor's feedback

* Addressing Heitor's feedback

* Addressing Heitor's feedback

* add result warning, add asdict test, metadata test

* refactor: initial refactoring

* chore: branding

Signed-off-by: heitorlessa <[email protected]>

* refactor: use classvar and tuple for perf

Signed-off-by: heitorlessa <[email protected]>

* chore: fix rebase issue

Signed-off-by: heitorlessa <[email protected]>

* chore: fix mypy tuple exactness type

Signed-off-by: heitorlessa <[email protected]>

* remove Ok in example response,add failure example

* chore: clean up docs example

Signed-off-by: heitorlessa <[email protected]>

* chore: lower cognitive overhead; add example docstring

Signed-off-by: heitorlessa <[email protected]>

* add drop example

* docs: give info upfront, name examples

* docs: improve transforming records example

* docs: improve dropping records example

* docs: improve exception example

---------

Signed-off-by: heitorlessa <[email protected]>
Co-authored-by: Leandro Damascena <[email protected]>
Co-authored-by: heitorlessa <[email protected]>

Author: 3 people

aws_lambda_powertools/shared/__init__.py

@@ -0,0 +1 @@
+"""Internal shared functions. Do not rely on it besides internal usage."""

aws_lambda_powertools/utilities/data_classes/__init__.py

@@ -14,7 +14,12 @@
 from .event_bridge_event import EventBridgeEvent
 from .event_source import event_source
 from .kafka_event import KafkaEvent
-from .kinesis_firehose_event import KinesisFirehoseEvent
+from .kinesis_firehose_event import (
+    KinesisFirehoseDataTransformationRecord,
+    KinesisFirehoseDataTransformationRecordMetadata,
+    KinesisFirehoseDataTransformationResponse,
+    KinesisFirehoseEvent,
+)
 from .kinesis_stream_event import KinesisStreamEvent
 from .lambda_function_url_event import LambdaFunctionUrlEvent
 from .s3_event import S3Event, S3EventBridgeNotificationEvent
@@ -39,6 +44,9 @@
     "KafkaEvent",
     "KinesisFirehoseEvent",
     "KinesisStreamEvent",
+    "KinesisFirehoseDataTransformationResponse",
+    "KinesisFirehoseDataTransformationRecord",
+    "KinesisFirehoseDataTransformationRecordMetadata",
     "LambdaFunctionUrlEvent",
     "S3Event",
     "S3EventBridgeNotificationEvent",

aws_lambda_powertools/utilities/data_classes/kinesis_firehose_event.py

@@ -1,9 +1,179 @@
 import base64
-from typing import Iterator, Optional
+import json
+import warnings
+from dataclasses import dataclass, field
+from typing import Any, Callable, ClassVar, Dict, Iterator, List, Optional, Tuple
+
+from typing_extensions import Literal
 
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 
 
+@dataclass(repr=False, order=False, frozen=True)
+class KinesisFirehoseDataTransformationRecordMetadata:
+    """
+    Metadata in Firehose Data Transform Record.
+
+    Parameters
+    ----------
+    partition_keys: Dict[str, str]
+        A dict of partition keys/value in string format, e.g. `{"year":"2023","month":"09"}`
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
+    """
+
+    partition_keys: Dict[str, str] = field(default_factory=lambda: {})
+
+    def asdict(self) -> Dict:
+        if self.partition_keys is not None:
+            return {"partitionKeys": self.partition_keys}
+        return {}
+
+
+@dataclass(repr=False, order=False)
+class KinesisFirehoseDataTransformationRecord:
+    """Record in Kinesis Data Firehose response object.
+
+    Parameters
+    ----------
+    record_id: str
+        uniquely identifies this record within the current batch
+    result: Literal["Ok", "Dropped", "ProcessingFailed"]
+        record data transformation status, whether it succeeded, should be dropped, or failed.
+    data: str
+        base64-encoded payload, by default empty string.
+
+        Use `data_from_text` or `data_from_json` methods to convert data if needed.
+
+    metadata: Optional[KinesisFirehoseDataTransformationRecordMetadata]
+        Metadata associated with this record; can contain partition keys.
+
+        See: https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
+    json_serializer: Callable
+        function to serialize `obj` to a JSON formatted `str`, by default json.dumps
+    json_deserializer: Callable
+        function to deserialize `str`, `bytes`, bytearray` containing a JSON document to a Python `obj`,
+        by default json.loads
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html
+    """
+
+    _valid_result_types: ClassVar[Tuple[str, str, str]] = ("Ok", "Dropped", "ProcessingFailed")
+
+    record_id: str
+    result: Literal["Ok", "Dropped", "ProcessingFailed"] = "Ok"
+    data: str = ""
+    metadata: Optional[KinesisFirehoseDataTransformationRecordMetadata] = None
+    json_serializer: Callable = json.dumps
+    json_deserializer: Callable = json.loads
+    _json_data: Optional[Any] = None
+
+    def asdict(self) -> Dict:
+        if self.result not in self._valid_result_types:
+            warnings.warn(
+                stacklevel=1,
+                message=f'The result "{self.result}" is not valid, Choose from "Ok", "Dropped", "ProcessingFailed"',
+            )
+
+        record: Dict[str, Any] = {
+            "recordId": self.record_id,
+            "result": self.result,
+            "data": self.data,
+        }
+        if self.metadata:
+            record["metadata"] = self.metadata.asdict()
+        return record
+
+    @property
+    def data_as_bytes(self) -> bytes:
+        """Decoded base64-encoded data as bytes"""
+        if not self.data:
+            return b""
+        return base64.b64decode(self.data)
+
+    @property
+    def data_as_text(self) -> str:
+        """Decoded base64-encoded data as text"""
+        if not self.data:
+            return ""
+        return self.data_as_bytes.decode("utf-8")
+
+    @property
+    def data_as_json(self) -> Dict:
+        """Decoded base64-encoded data loaded to json"""
+        if not self.data:
+            return {}
+        if self._json_data is None:
+            self._json_data = self.json_deserializer(self.data_as_text)
+        return self._json_data
+
+
+@dataclass(repr=False, order=False)
+class KinesisFirehoseDataTransformationResponse:
+    """Kinesis Data Firehose response object
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html
+
+    Parameters
+    ----------
+    records : List[KinesisFirehoseResponseRecord]
+        records of Kinesis Data Firehose response object,
+        optional parameter at start. can be added later using `add_record` function.
+
+    Examples
+    --------
+
+    **Transforming data records**
+
+    ```python
+    from aws_lambda_powertools.utilities.data_classes import (
+        KinesisFirehoseDataTransformationRecord,
+        KinesisFirehoseDataTransformationResponse,
+        KinesisFirehoseEvent,
+    )
+    from aws_lambda_powertools.utilities.serialization import base64_from_json
+    from aws_lambda_powertools.utilities.typing import LambdaContext
+
+
+    def lambda_handler(event: dict, context: LambdaContext):
+        firehose_event = KinesisFirehoseEvent(event)
+        result = KinesisFirehoseDataTransformationResponse()
+
+        for record in firehose_event.records:
+            payload = record.data_as_text  # base64 decoded data as str
+
+            ## generate data to return
+            transformed_data = {"tool_used": "powertools_dataclass", "original_payload": payload}
+            processed_record = KinesisFirehoseDataTransformationRecord(
+                record_id=record.record_id,
+                data=base64_from_json(transformed_data),
+            )
+
+            result.add_record(processed_record)
+
+        # return transformed records
+        return result.asdict()
+    ```
+    """
+
+    records: List[KinesisFirehoseDataTransformationRecord] = field(default_factory=list)
+
+    def add_record(self, record: KinesisFirehoseDataTransformationRecord):
+        self.records.append(record)
+
+    def asdict(self) -> Dict:
+        if not self.records:
+            raise ValueError("Amazon Kinesis Data Firehose doesn't accept empty response")
+
+        return {"records": [record.asdict() for record in self.records]}
+
+
 class KinesisFirehoseRecordMetadata(DictWrapper):
     @property
     def _metadata(self) -> dict:
@@ -77,6 +247,32 @@ def data_as_json(self) -> dict:
             self._json_data = self._json_deserializer(self.data_as_text)
         return self._json_data
 
+    def build_data_transformation_response(
+        self,
+        result: Literal["Ok", "Dropped", "ProcessingFailed"] = "Ok",
+        data: str = "",
+        metadata: Optional[KinesisFirehoseDataTransformationRecordMetadata] = None,
+    ) -> KinesisFirehoseDataTransformationRecord:
+        """Create a KinesisFirehoseResponseRecord directly using the record_id and given values
+
+        Parameters
+        ----------
+        result : Literal["Ok", "Dropped", "ProcessingFailed"]
+            processing result, supported value: Ok, Dropped, ProcessingFailed
+        data : str, optional
+            data blob, base64-encoded, optional at init. Allows pass in base64-encoded data directly or
+            use either function like `data_from_text`, `data_from_json` to populate data
+        metadata: KinesisFirehoseResponseRecordMetadata, optional
+            Metadata associated with this record; can contain partition keys
+            - https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
+        """
+        return KinesisFirehoseDataTransformationRecord(
+            record_id=self.record_id,
+            result=result,
+            data=data,
+            metadata=metadata,
+        )
+
 
 class KinesisFirehoseEvent(DictWrapper):
     """Kinesis Data Firehose event

aws_lambda_powertools/utilities/serialization.py

@@ -0,0 +1,59 @@
+"""Standalone functions to serialize/deserialize common data structures"""
+import base64
+import json
+from typing import Any, Callable
+
+
+def base64_encode(data: str) -> str:
+    """Encode a string and returns Base64-encoded encoded value.
+
+    Parameters
+    ----------
+    data: str
+        The string to encode.
+
+    Returns
+    -------
+    str
+        The Base64-encoded encoded value.
+    """
+    return base64.b64encode(data.encode()).decode("utf-8")
+
+
+def base64_decode(data: str) -> str:
+    """Decodes a Base64-encoded string and returns the decoded value.
+
+    Parameters
+    ----------
+    data: str
+        The Base64-encoded string to decode.
+
+    Returns
+    -------
+    str
+        The decoded string value.
+    """
+    return base64.b64decode(data).decode("utf-8")
+
+
+def base64_from_str(data: str) -> str:
+    """Encode str as base64 string"""
+    return base64.b64encode(data.encode()).decode("utf-8")
+
+
+def base64_from_json(data: Any, json_serializer: Callable[..., str] = json.dumps) -> str:
+    """Encode JSON serializable data as base64 string
+
+    Parameters
+    ----------
+    data: Any
+        JSON serializable (dict, list, boolean, etc.)
+    json_serializer: Callable
+        function to serialize `obj` to a JSON formatted `str`, by default json.dumps
+
+    Returns
+    -------
+    str:
+        JSON string as base64 string
+    """
+    return base64_from_str(data=json_serializer(data))

docs/utilities/data_classes.md

@@ -975,18 +975,39 @@ or plain text, depending on the original payload.
 
 ### Kinesis Firehose delivery stream
 
-Kinesis Firehose Data Transformation can use a Lambda Function to modify the records
-inline, and re-emit them back to the Delivery Stream.
+When using Kinesis Firehose, you can use a Lambda function to [perform data transformation](https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html){target="_blank"}. For each transformed record, you can choose to either:
 
-Similar to Kinesis Data Streams, the events contain base64 encoded data. You can use the helper
-function to access the data either as json or plain text, depending on the original payload.
+* **A)** Put them back to the delivery stream (default)
+* **B)** Drop them so consumers don't receive them (e.g., data validation)
+* **C)** Indicate a record failed data transformation and should be retried
 
-=== "app.py"
+To do that, you can use `KinesisFirehoseDataTransformationResponse` class along with helper functions to make it easier to decode and encode base64 data in the stream.
 
-    ```python
+=== "Transforming streaming records"
+
+    ```python hl_lines="2-3 12 28"
     --8<-- "examples/event_sources/src/kinesis_firehose_delivery_stream.py"
     ```
 
+    1. **Ingesting JSON payloads?** <br><br> Use `record.data_as_json` to easily deserialize them.
+    2. For your convenience, `base64_from_json` serializes a dict to JSON, then encode as base64 data.
+
+=== "Dropping invalid records"
+
+    ```python hl_lines="5-6 16 34"
+    --8<-- "examples/event_sources/src/kinesis_firehose_response_drop.py"
+    ```
+
+    1. This exception would be generated from `record.data_as_json` if invalid payload.
+
+=== "Indicating a processing failure"
+
+    ```python hl_lines="2-3 33"
+    --8<-- "examples/event_sources/src/kinesis_firehose_response_exception.py"
+    ```
+
+    1. This record will now be sent to your [S3 bucket in the `processing-failed` folder](https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html#data-transformation-failure-handling){target="_blank"}.
+
 ### Lambda Function URL
 
 === "app.py"

examples/event_sources/src/kinesis_firehose_delivery_stream.py

@@ -1,28 +1,28 @@
-import base64
-import json
-
 from aws_lambda_powertools.utilities.data_classes import (
+    KinesisFirehoseDataTransformationResponse,
     KinesisFirehoseEvent,
     event_source,
 )
+from aws_lambda_powertools.utilities.serialization import base64_from_json
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
 @event_source(data_class=KinesisFirehoseEvent)
 def lambda_handler(event: KinesisFirehoseEvent, context: LambdaContext):
-    result = []
+    result = KinesisFirehoseDataTransformationResponse()
 
     for record in event.records:
-        # if data was delivered as json; caches loaded value
-        data = record.data_as_json
+        # get original data using data_as_text property
+        data = record.data_as_text  # (1)!
+
+        ## generate data to return
+        transformed_data = {"new_data": "transformed data using Powertools", "original_payload": data}
 
-        processed_record = {
-            "recordId": record.record_id,
-            "data": base64.b64encode(json.dumps(data).encode("utf-8")),
-            "result": "Ok",
-        }
+        processed_record = record.build_data_transformation_response(
+            data=base64_from_json(transformed_data),  # (2)!
+        )
 
-        result.append(processed_record)
+        result.add_record(processed_record)
 
     # return transformed records
-    return {"records": result}
+    return result.asdict()