Refactoring functions

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka_consumer/__init__.py

@@ -1,9 +1,9 @@
-from aws_lambda_powertools.utilities.kafka_consumer.consumer_record import ConsumerRecord
+from aws_lambda_powertools.utilities.kafka_consumer.consumer_records import ConsumerRecords
 from aws_lambda_powertools.utilities.kafka_consumer.kafka_consumer import kafka_consumer
 from aws_lambda_powertools.utilities.kafka_consumer.schema_config import SchemaConfig
 
 __all__ = [
     "kafka_consumer",
-    "ConsumerRecord",
+    "ConsumerRecords",
     "SchemaConfig",
 ]

aws_lambda_powertools/utilities/kafka_consumer/consumer_record.py renamed to aws_lambda_powertools/utilities/kafka_consumer/consumer_records.py

@@ -4,10 +4,8 @@
 
 from aws_lambda_powertools.utilities.data_classes.common import CaseInsensitiveDict
 from aws_lambda_powertools.utilities.data_classes.kafka_event import KafkaEvent, KafkaEventBase
-from aws_lambda_powertools.utilities.kafka_consumer.functions import (
-    deserialize_avro,
-    deserialize_protobuf_with_compiled_classes,
-)
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.deserializer import get_deserializer
+from aws_lambda_powertools.utilities.kafka_consumer.serialization.serialization import serialize_to_output_type
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -25,32 +23,43 @@ def __init__(self, data: dict[str, Any], deserialize: SchemaConfig | None = None
         self.deserialize = deserialize
 
     @property
-    def key(self) -> str | None:
+    def key(self) -> Any:
         key = self.get("key")
-        if key and self.deserialize.key_schema_type:
-            if self.deserialize.value_schema_type == "AVRO":
-                return deserialize_avro(key, self.deserialize.value_schema_str)
-            elif self.deserialize.value_schema_type == "PROTOBUF":
-                return deserialize_protobuf_with_compiled_classes(key, self.deserialize.value_schema_str)
-            elif self.deserialize.value_schema_type == "JSON":
-                return self._json_deserializer(key)
-            else:
-                raise ValueError("Invalid value_schema_type")
+        if key and (self.deserialize and self.deserialize.key_schema_type):
+            deserializer = get_deserializer(
+                self.deserialize.value_schema_type,
+                self.deserialize.value_schema_str,
+            )
+            deserialized_key = deserializer.deserialize(key)
+
+            if self.deserialize.key_output_serializer:
+                return serialize_to_output_type(
+                    deserialized_key,
+                    self.deserialize.key_output_serializer,
+                )
+
+            return deserialized_key
 
         return key
 
     @property
-    def value(self) -> str:
+    def value(self) -> Any:
         value = self["value"]
-        if self.deserialize.value_schema_type:
-            if self.deserialize.value_schema_type == "AVRO":
-                return deserialize_avro(value, self.deserialize.value_schema_str)
-            elif self.deserialize.value_schema_type == "PROTOBUF":
-                return deserialize_protobuf_with_compiled_classes(value, self.deserialize.value_schema_str)
-            elif self.deserialize.value_schema_type == "JSON":
-                return self._json_deserializer(value)
-            else:
-                raise ValueError("Invalid value_schema_type")
+        if value and (self.deserialize and self.deserialize.value_schema_type):
+            deserializer = get_deserializer(
+                self.deserialize.value_schema_type,
+                self.deserialize.value_schema_str,
+            )
+            deserialized_value = deserializer.deserialize(value)
+
+            if self.deserialize.value_output_serializer:
+                return serialize_to_output_type(
+                    deserialized_value,
+                    self.deserialize.value_output_serializer,
+                )
+
+            return deserialized_value
+
         return value
 
     @property
@@ -81,7 +90,7 @@ def original_headers(self) -> dict[str, bytes]:
         return self["headers"]
 
 
-class ConsumerRecord(KafkaEvent):
+class ConsumerRecords(KafkaEvent):
     """Self-managed or MSK Apache Kafka event trigger
     Documentation:
     --------------

aws_lambda_powertools/utilities/kafka_consumer/deserializer/__init__.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import io
+from typing import Any
+
+from avro.io import BinaryDecoder, DatumReader
+from avro.schema import parse as parse_schema
+
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
+from aws_lambda_powertools.utilities.kafka_consumer.exceptions import (
+    KafkaConsumerAvroMissingSchemaError,
+    KafkaConsumerDeserializationError,
+)
+
+
+class AvroDeserializer(DeserializerBase):
+    def __init__(self, schema_str: str):
+        if not schema_str:
+            raise KafkaConsumerAvroMissingSchemaError("Schema string must be provided for Avro deserialization")
+        self.parsed_schema = parse_schema(schema_str)
+        self.reader = DatumReader(self.parsed_schema)
+
+    def deserialize(self, data: bytes | str) -> dict[str, Any]:
+        try:
+            value = self._decode_input(data)
+            bytes_reader = io.BytesIO(value)
+            decoder = BinaryDecoder(bytes_reader)
+            return self.reader.read(decoder)
+        except (TypeError, ValueError) as e:
+            raise KafkaConsumerDeserializationError(
+                f"Avro deserialization error: {type(e).__name__}: {str(e)}",
+            ) from e

aws_lambda_powertools/utilities/kafka_consumer/deserializer/avro.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import base64
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class DeserializerBase(ABC):
+    @abstractmethod
+    def deserialize(self, data: bytes | str) -> dict[str, Any]:
+        pass
+
+    def _decode_input(self, data: bytes | str) -> bytes:
+        if isinstance(data, str):
+            return base64.b64decode(data)
+        elif isinstance(data, bytes):
+            return data
+        else:
+            try:
+                return base64.b64decode(data)
+            except Exception as e:
+                raise TypeError(
+                    f"Expected bytes or base64-encoded string, got {type(data).__name__}. Error: {str(e)}",
+                ) from e

aws_lambda_powertools/utilities/kafka_consumer/deserializer/base.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
+
+
+def get_deserializer(schema_type: str, schema_value: Any) -> DeserializerBase:
+    if schema_type == "AVRO":
+        # Import here to avoid dependency if not used
+        from aws_lambda_powertools.utilities.kafka_consumer.deserializer.avro import AvroDeserializer
+
+        return AvroDeserializer(schema_value)
+    elif schema_type == "PROTOBUF":
+        # Import here to avoid dependency if not used
+        from aws_lambda_powertools.utilities.kafka_consumer.deserializer.protobuf import ProtobufDeserializer
+
+        return ProtobufDeserializer(schema_value)
+    elif schema_type == "JSON":
+        # Import here to avoid dependency if not used
+        from aws_lambda_powertools.utilities.kafka_consumer.deserializer.json import JsonDeserializer
+
+        return JsonDeserializer()
+    else:
+        raise ValueError(f"Invalid schema_type: {schema_type}")

aws_lambda_powertools/utilities/kafka_consumer/deserializer/deserializer.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import json
+
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
+from aws_lambda_powertools.utilities.kafka_consumer.exceptions import KafkaConsumerDeserializationError
+
+
+class JsonDeserializer(DeserializerBase):
+    def deserialize(self, data: bytes | str) -> dict:
+        try:
+            value = self._decode_input(data)
+            return json.loads(value.decode("utf-8"))
+        except Exception as e:
+            raise KafkaConsumerDeserializationError(
+                f"JSON deserialization error: {type(e).__name__}: {str(e)}",
+            ) from e

aws_lambda_powertools/utilities/kafka_consumer/deserializer/json.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+from google.protobuf.json_format import MessageToDict
+
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
+from aws_lambda_powertools.utilities.kafka_consumer.exceptions import (
+    KafkaConsumerDeserializationError,
+)
+
+
+class ProtobufDeserializer(DeserializerBase):
+    def __init__(self, message_class: Any):
+        self.message_class = message_class
+
+    def deserialize(self, data: bytes | str) -> dict:
+        try:
+            value = self._decode_input(data)
+            message = self.message_class()
+            message.ParseFromString(value)
+            return MessageToDict(message, preserving_proto_field_name=True)
+        except Exception as e:
+            raise KafkaConsumerDeserializationError(
+                f"Protocol Buffer deserialization error: {type(e).__name__}: {str(e)}",
+            ) from e