support kinesis response

Author: roger-zhangg

aws_lambda_powertools/utilities/data_classes/__init__.py

@@ -15,6 +15,17 @@
 from .event_source import event_source
 from .kafka_event import KafkaEvent
 from .kinesis_firehose_event import KinesisFirehoseEvent
+from .kinesis_firehose_response import (
+    FirehoseStateDropped,
+    FirehoseStateFailed,
+    FirehoseStateOk,
+    KinesisFirehoseResponce,
+    KinesisFirehoseResponceFactory,
+    KinesisFirehoseResponceRecord,
+    KinesisFirehoseResponceRecordFactory,
+    KinesisFirehoseResponseRecordMetadata,
+    KinesisFirehoseResponseRecordMetadataFactory,
+)
 from .kinesis_stream_event import KinesisStreamEvent
 from .lambda_function_url_event import LambdaFunctionUrlEvent
 from .s3_event import S3Event, S3EventBridgeNotificationEvent
@@ -37,6 +48,15 @@
     "KafkaEvent",
     "KinesisFirehoseEvent",
     "KinesisStreamEvent",
+    "KinesisFirehoseResponce",
+    "KinesisFirehoseResponceRecord",
+    "KinesisFirehoseResponseRecordMetadata",
+    "FirehoseStateOk",
+    "FirehoseStateDropped",
+    "FirehoseStateFailed",
+    "KinesisFirehoseResponceFactory",
+    "KinesisFirehoseResponceRecordFactory",
+    "KinesisFirehoseResponseRecordMetadataFactory",
     "LambdaFunctionUrlEvent",
     "S3Event",
     "S3EventBridgeNotificationEvent",

aws_lambda_powertools/utilities/data_classes/kinesis_firehose_response.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import base64
+from typing import Callable, Iterator, List, Optional, Union
+
+from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
+
+FirehoseStateOk = "Ok"
+FirehoseStateDropped = "Dropped"
+FirehoseStateFailed = "ProcessingFailed"
+
+
+class KinesisFirehoseResponseRecordMetadata(DictWrapper):
+    """
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
+    """
+
+    @property
+    def _metadata(self) -> Optional[dict]:
+        """Optional: metadata associated with this record; present only when Kinesis Stream is source"""
+        return self["metadata"]  # could raise KeyError
+
+    @property
+    def partition_keys(self) -> Optional[dict[str, str]]:
+        """Kinesis stream partition key; present only when Kinesis Stream is source"""
+        return self._metadata["partitionKeys"]
+
+
+def KinesisFirehoseResponseRecordMetadataFactory(
+    partition_keys: dict[str, str],
+    json_deserializer: Optional[Callable] = None,
+) -> KinesisFirehoseResponseRecordMetadata:
+    data = {
+        "metadata": {
+            "partitionKeys": partition_keys,
+        },
+    }
+    return KinesisFirehoseResponseRecordMetadata(data=data, json_deserializer=json_deserializer)
+
+
+class KinesisFirehoseResponceRecord(DictWrapper):
+    """Record in Kinesis Data Firehose event
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html
+    """
+
+    @property
+    def record_id(self) -> str:
+        """Record ID; uniquely identifies this record within the current batch"""
+        return self["recordId"]
+
+    @property
+    def result(self) -> Union[FirehoseStateOk, FirehoseStateDropped, FirehoseStateFailed]:
+        """processing result, supported value: Ok, Dropped, ProcessingFailed"""
+        return self["result"]
+
+    @property
+    def data(self) -> str:
+        """The data blob, base64-encoded"""
+        return self["data"]
+
+    @property
+    def metadata(self) -> Optional[KinesisFirehoseResponseRecordMetadata]:
+        """Optional: metadata associated with this record; present only when Kinesis Stream is source"""
+        return KinesisFirehoseResponseRecordMetadata(self._data) if self.get("metadata") else None
+
+    @property
+    def data_as_bytes(self) -> bytes:
+        """Decoded base64-encoded data as bytes"""
+        return base64.b64decode(self.data)
+
+    @property
+    def data_as_text(self) -> str:
+        """Decoded base64-encoded data as text"""
+        return self.data_as_bytes.decode("utf-8")
+
+    @property
+    def data_as_json(self) -> dict:
+        """Decoded base64-encoded data loaded to json"""
+        if self._json_data is None:
+            self._json_data = self._json_deserializer(self.data_as_text)
+        return self._json_data
+
+
+def KinesisFirehoseResponceRecordFactory(
+    record_id: str,
+    result: Union[FirehoseStateOk, FirehoseStateDropped, FirehoseStateFailed],
+    data: str,
+    metadata: Optional[KinesisFirehoseResponseRecordMetadata] = None,
+    json_deserializer: Optional[Callable] = None,
+) -> KinesisFirehoseResponceRecord:
+    pass_data = {
+        "recordId": record_id,
+        "result": result,
+        "data": base64.b64encode(data.encode("utf-8")).decode("utf-8"),
+    }
+    if metadata:
+        data["metadata"] = metadata
+    return KinesisFirehoseResponceRecord(data=pass_data, json_deserializer=json_deserializer)
+
+
+class KinesisFirehoseResponce(DictWrapper):
+    """Kinesis Data Firehose event
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html
+    """
+
+    @property
+    def records(self) -> Iterator[KinesisFirehoseResponceRecord]:
+        for record in self["records"]:
+            yield KinesisFirehoseResponceRecord(data=record, json_deserializer=self._json_deserializer)
+
+
+def KinesisFirehoseResponceFactory(
+    records: List[KinesisFirehoseResponceRecord],
+    json_deserializer: Optional[Callable] = None,
+) -> KinesisFirehoseResponce:
+    pass_data = {"records": records}
+    return KinesisFirehoseResponce(data=pass_data, json_deserializer=json_deserializer)

examples/event_sources/src/kinesis_firehose_delivery_stream.py

@@ -1,28 +1,30 @@
-import base64
 import json
 
 from aws_lambda_powertools.utilities.data_classes import (
+    FirehoseStateOk,
     KinesisFirehoseEvent,
+    KinesisFirehoseResponce,
+    KinesisFirehoseResponceRecordFactory,
     event_source,
 )
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 
 @event_source(data_class=KinesisFirehoseEvent)
 def lambda_handler(event: KinesisFirehoseEvent, context: LambdaContext):
-    result = []
+    result = KinesisFirehoseResponce({})
 
     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",
-        }
+        processed_record = KinesisFirehoseResponceRecordFactory(
+            record_id=record.record_id,
+            result=FirehoseStateOk,
+            data=(json.dumps(data)),
+        )
 
-        result.append(processed_record)
+        result.add_record(processed_record)
 
     # return transformed records
-    return {"records": result}
+    return result

tests/unit/data_classes/test_kinesis_firehose_response.py

@@ -0,0 +1,36 @@
+from aws_lambda_powertools.utilities.data_classes import (
+    FirehoseStateOk,
+    KinesisFirehoseEvent,
+    KinesisFirehoseResponceFactory,
+    KinesisFirehoseResponceRecordFactory,
+)
+from tests.functional.utils import load_event
+
+
+def test_kinesis_firehose_response():
+    raw_event = load_event("kinesisFirehoseKinesisEvent.json")
+    parsed_event = KinesisFirehoseEvent(raw_event)
+
+    result = []
+    for record in parsed_event.records:
+        # if data was delivered as json; caches loaded value
+        data = record.data_as_text
+
+        processed_record = KinesisFirehoseResponceRecordFactory(
+            record_id=record.record_id,
+            result=FirehoseStateOk,
+            data=(data),
+        )
+
+        result.append(processed_record)
+    response = KinesisFirehoseResponceFactory(result)
+
+    res_records = list(response.records)
+    assert len(res_records) == 2
+    record_01, record_02 = res_records[:]
+    record01_raw = raw_event["records"][0]
+    assert record_01.result == FirehoseStateOk
+    assert record_01.record_id == record01_raw["recordId"]
+    assert record_01.data_as_bytes == b"Hello World"
+    assert record_01.data_as_text == "Hello World"
+    assert record_01.data == record01_raw["data"]