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

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
@@ -37,6 +42,9 @@
     "KafkaEvent",
     "KinesisFirehoseEvent",
     "KinesisStreamEvent",
+    "KinesisFirehoseDataTransformationResponse",
+    "KinesisFirehoseDataTransformationRecord",
+    "KinesisFirehoseDataTransformationRecordMetadata",
     "LambdaFunctionUrlEvent",
     "S3Event",
     "S3EventBridgeNotificationEvent",

aws_lambda_powertools/utilities/data_classes/kinesis_firehose_event.py

@@ -1,9 +1,125 @@
 import base64
-from typing import Iterator, Optional
+import json
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Iterator, List, Optional
+
+from typing_extensions import Literal
 
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 
 
+@dataclass(repr=False, order=False, slots=True, frozen=True)
+class KinesisFirehoseDataTransformationRecordMetadata:
+    """
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
+    """
+
+    partition_keys: Optional[Dict[str, str]]
+
+    def asdict(self) -> Dict:
+        if self.partition_keys is not None:
+            return {"partitionKeys": self.partition_keys}
+        return {}
+
+
+@dataclass(repr=False, order=False, slots=True)
+class KinesisFirehoseDataTransformationRecord:
+    """Record in Kinesis Data Firehose response object
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html
+    """
+
+    # Record ID; uniquely identifies this record within the current batch"""
+    record_id: str
+    # Processing result, supported value: Ok, Dropped, ProcessingFailed"""
+    result: Literal["Ok", "Dropped", "ProcessingFailed"] = "Ok"
+    # 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"""
+    data: Optional[str] = None
+    # Optional: Metadata associated with this record; can contain partition keys
+    # See - https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
+    metadata: Optional[KinesisFirehoseDataTransformationRecordMetadata] = None
+    _json_data: Optional[Any] = None
+    json_serializer: Callable = json.dumps
+    json_deserializer: Callable = json.loads
+
+    def data_from_byte(self, data: bytes):
+        """Populate data field using a byte like data"""
+        self.data = base64.b64encode(data).decode("utf-8")
+
+    def data_from_text(self, data: str):
+        """Populate data field using a string like data"""
+        self.data_from_byte(data.encode("utf-8"))
+
+    def data_from_json(self, data: Any):
+        """Populate data field using any structure that could be converted to json"""
+        self.data_from_text(data=self.json_serializer(data))
+
+    def asdict(self) -> Dict:
+        r: Dict[str, Any] = {
+            "recordId": self.record_id,
+            "result": self.result,
+            "data": self.data,
+        }
+        if self.metadata:
+            r["metadata"] = self.metadata.asdict()
+        return r
+
+    @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, slots=True)
+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.
+    """
+
+    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("Kinesis Firehose doesn't accept empty response")
+
+        return {"records": [r.asdict() for r in self.records]}
+
+
 class KinesisFirehoseRecordMetadata(DictWrapper):
     @property
     def _metadata(self) -> dict:
@@ -77,6 +193,31 @@ 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: Optional[str] = None,
+        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

docs/utilities/data_classes.md

@@ -981,12 +981,23 @@ inline, and re-emit them back to the Delivery Stream.
 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.
 
+When constructing response to Firehose, You can utilize the `KinesisFirehoseResponse` class shown
+in the example below.
+
 === "app.py"
 
     ```python
     --8<-- "examples/event_sources/src/kinesis_firehose_delivery_stream.py"
     ```
 
+You can also construct response without using `event_source` wrapper. Shown in the example below.
+
+=== "app.py"
+
+    ```python
+    --8<-- "examples/event_sources/src/kinesis_firehose_response.py"
+    ```
+
 ### Lambda Function URL
 
 === "app.py"

examples/event_sources/src/kinesis_firehose_delivery_stream.py

@@ -1,7 +1,5 @@
-import base64
-import json
-
 from aws_lambda_powertools.utilities.data_classes import (
+    KinesisFirehoseDataTransformationResponse,
     KinesisFirehoseEvent,
     event_source,
 )
@@ -10,19 +8,20 @@
 
 @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
 
-        processed_record = {
-            "recordId": record.record_id,
-            "data": base64.b64encode(json.dumps(data).encode("utf-8")),
-            "result": "Ok",
-        }
+        ## do all kind of stuff with data
+        ## generate data to return
+        new_data = {"tool_used": "powertools_dataclass", "original_payload": data}
+
+        processed_record = record.build_data_transformation_response(result="Ok")
+        processed_record.data_from_json(data=new_data)
 
-        result.append(processed_record)
+        result.add_record(processed_record)
 
     # return transformed records
-    return {"records": result}
+    return result.asdict()

examples/event_sources/src/kinesis_firehose_response.py

@@ -0,0 +1,24 @@
+import base64
+
+from aws_lambda_powertools.utilities.data_classes import (
+    KinesisFirehoseDataTransformationRecord,
+    KinesisFirehoseDataTransformationResponse,
+)
+
+
+def lambda_handler(event, context):
+    result = KinesisFirehoseDataTransformationResponse()
+
+    for record in event["records"]:
+        print(record["recordId"])
+        payload = base64.b64decode(record["data"]).decode("utf-8")
+        ## do all kind of stuff with payload
+        ## generate data to return
+        new_data = {"tool_used": "powertools_dataclass", "original_payload": payload}
+
+        processed_record = KinesisFirehoseDataTransformationRecord(record_id=record["recordId"], result="Ok")
+        processed_record.data_from_text(data=new_data)
+        result.add_record(processed_record)
+
+    # return transformed records
+    return result.asdict()

tests/unit/data_classes/test_kinesis_firehose_response.py

@@ -0,0 +1,69 @@
+from aws_lambda_powertools.utilities.data_classes import (
+    KinesisFirehoseDataTransformationRecord,
+    KinesisFirehoseDataTransformationRecordMetadata,
+    KinesisFirehoseDataTransformationResponse,
+    KinesisFirehoseEvent,
+)
+from tests.functional.utils import load_event
+
+
+def test_kinesis_firehose_response():
+    raw_event = load_event("kinesisFirehoseKinesisEvent.json")
+    parsed_event = KinesisFirehoseEvent(data=raw_event)
+
+    response = KinesisFirehoseDataTransformationResponse()
+    for record in parsed_event.records:
+        # if data was delivered as json; caches loaded value
+        data = record.data_as_text
+
+        metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": 2023})
+        processed_record = KinesisFirehoseDataTransformationRecord(
+            record_id=record.record_id,
+            result="Ok",
+            metadata=metadata_partition,
+        )
+        processed_record.data_from_text(data=data)
+        response.add_record(record=processed_record)
+    response_dict = response.asdict()
+
+    res_records = list(response_dict["records"])
+    assert len(res_records) == 2
+    record_01, record_02 = res_records[:]
+    record01_raw = raw_event["records"][0]
+    assert record_01["result"] == "Ok"
+    assert record_01["recordId"] == record01_raw["recordId"]
+    assert record_01["data"] == record01_raw["data"]
+    assert record_01["metadata"]["partitionKeys"]["year"] == 2023
+
+    assert response.records[0].data_as_bytes == b"Hello World"
+    assert response.records[0].data_as_text == "Hello World"
+
+
+def test_kinesis_firehose_create_response():
+    raw_event = load_event("kinesisFirehoseKinesisEvent.json")
+    parsed_event = KinesisFirehoseEvent(data=raw_event)
+
+    response = KinesisFirehoseDataTransformationResponse()
+    for record in parsed_event.records:
+        # if data was delivered as json; caches loaded value
+        data = record.data_as_text
+        metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": 2023})
+        processed_record = record.build_data_transformation_response(
+            result="Ok",
+            metadata=metadata_partition,
+        )
+        processed_record.data_from_text(data=data)
+        response.add_record(record=processed_record)
+    response_dict = response.asdict()
+
+    res_records = list(response_dict["records"])
+    assert len(res_records) == 2
+    record_01, record_02 = res_records[:]
+    record01_raw = raw_event["records"][0]
+    assert record_01["result"] == "Ok"
+    assert record_01["recordId"] == record01_raw["recordId"]
+    assert record_01["data"] == record01_raw["data"]
+    assert record_01["metadata"]["partitionKeys"]["year"] == 2023
+
+    assert response.records[0].data_as_bytes == b"Hello World"
+    assert response.records[0].data_as_text == "Hello World"