Adding support for Kafka Consumer - first commit

Author: leandrodamascena

aws_lambda_powertools/utilities/data_classes/kafka_event.py

@@ -10,7 +10,7 @@
     from collections.abc import Iterator
 
 
-class KafkaEventRecord(DictWrapper):
+class KafkaEventBase(DictWrapper):
     @property
     def topic(self) -> str:
         """The Kafka topic."""
@@ -36,6 +36,8 @@ def timestamp_type(self) -> str:
         """The Kafka record timestamp type."""
         return self["timestampType"]
 
+
+class KafkaEventRecord(KafkaEventBase):
     @property
     def key(self) -> str | None:
         """

aws_lambda_powertools/utilities/kafka_consumer/consumer_record.py

@@ -1,15 +1,122 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from aws_lambda_powertools.utilities.data_classes.kafka_event import KafkaEventRecord
+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,
+)
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
 
-class ConsumerRecord(KafkaEventRecord):
+    from aws_lambda_powertools.utilities.kafka_consumer.schema_config import SchemaConfig
+
+
+class ConsumerRecordRecords(KafkaEventBase):
     """
     A Kafka Consumer Record
     """
 
-    def __init__(self, data: dict[str, Any], json_deserializer=None):
-        super().__init__(data, json_deserializer=json_deserializer)
-        self._json_deserializer = json_deserializer
+    def __init__(self, data: dict[str, Any], deserialize: SchemaConfig | None = None):
+        super().__init__(data)
+        self.deserialize = deserialize
+
+    @property
+    def key(self) -> str | None:
+        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")
+
+        return key
+
+    @property
+    def value(self) -> str:
+        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")
+        return value
+
+    @property
+    def original_value(self) -> str:
+        """The original (base64 encoded) Kafka record value."""
+        return self["value"]
+
+    @property
+    def original_key(self) -> str | None:
+        """
+        The original (base64 encoded) Kafka record key.
+
+        This key is optional; if not provided,
+        a round-robin algorithm will be used to determine
+        the partition for the message.
+        """
+
+        return self.get("key")
+
+    @property
+    def headers(self) -> list[dict[str, list[int]]]:
+        """The raw Kafka record headers."""
+        return CaseInsensitiveDict((k, bytes(v)) for chunk in self.headers for k, v in chunk.items())
+
+    @property
+    def original_headers(self) -> dict[str, bytes]:
+        """Decodes the headers as a single dictionary."""
+        return self["headers"]
+
+
+class ConsumerRecord(KafkaEvent):
+    """Self-managed or MSK Apache Kafka event trigger
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/lambda/latest/dg/with-kafka.html
+    - https://docs.aws.amazon.com/lambda/latest/dg/with-msk.html
+    """
+
+    def __init__(self, data: dict[str, Any], deserialize: SchemaConfig | None = None):
+        super().__init__(data)
+        self._records: Iterator[ConsumerRecordRecords] | None = None
+        self.deserialize = deserialize
+
+    @property
+    def records(self) -> Iterator[ConsumerRecordRecords]:
+        """The Kafka records."""
+        for chunk in self["records"].values():
+            for record in chunk:
+                yield ConsumerRecordRecords(data=record, deserialize=self.deserialize)
+
+    @property
+    def record(self) -> ConsumerRecordRecords:
+        """
+        Returns the next Kafka record using an iterator.
+
+        Returns
+        -------
+        ConsumerRecordRecords
+            The next Kafka record.
+
+        Raises
+        ------
+        StopIteration
+            If there are no more records available.
+
+        """
+        if self._records is None:
+            self._records = self.records
+        return next(self._records)

aws_lambda_powertools/utilities/kafka_consumer/exceptions.py

@@ -0,0 +1,16 @@
+class KafkaConsumerAvroSchemaMismatchError(Exception):
+    """
+    Avro schema mismatch
+    """
+
+
+class KafkaConsumerDeserializationError(Exception):
+    """
+    Avro schema impossible to deserialize
+    """
+
+
+class KafkaConsumerAvroMissingSchemaError(Exception):
+    """
+    Avro schema mismatch
+    """

aws_lambda_powertools/utilities/kafka_consumer/functions.py

@@ -1,38 +1,150 @@
 from __future__ import annotations
 
+import base64
 import io
+from typing import Any
 
-from avro.errors import SchemaResolutionException
 from avro.io import BinaryDecoder, DatumReader
+from avro.schema import parse as parse_schema
+from google.protobuf.json_format import MessageToDict
 
+from aws_lambda_powertools.utilities.kafka_consumer.exceptions import (
+    KafkaConsumerAvroMissingSchemaError,
+    KafkaConsumerAvroSchemaMismatchError,
+    KafkaConsumerDeserializationError,
+)
 
-def deserialize_avro(avro_bytes, reader_schema: str | None = None):
+
+def deserialize_avro(avro_bytes: bytes | str, value_schema_str: str) -> dict:
     """
-    Deserialize Avro binary data to Python objects
+    Deserialize Avro binary data to Python dictionary objects.
+
+    This function handles the deserialization of Avro-formatted binary data
+    using a specified schema string. It supports both raw bytes and
+    base64-encoded string inputs.
 
     Parameters
     ----------
-    avro_bytes: bytes
-        Avro binary data
-    reader_schema: str, Optional
-        Schema to use for reading
+    avro_bytes : bytes or str
+        Avro binary data, either as raw bytes or base64-encoded string.
+        If a string is provided, it will be treated as base64-encoded.
+    value_schema_str : str
+        Avro schema definition in JSON string format to use for reading.
+        Must be a valid Avro schema definition.
 
     Returns
     -------
-    dict
-        Deserialized Python object
+    Any
+        Deserialized Python dictionary representing the Avro data.
 
     Raises
     ------
-    ValueError
-        If reader_schema schema is None or if deserialization fails
+    KafkaConsumerAvroMissingSchemaError
+        If the schema is not provided
+    KafkaConsumerAvroSchemaMismatchError
+        If there's a schema mismatch
+    KafkaConsumerDeserializationError
+        If deserialization fails due to data corruption.
+    TypeError
+        If avro_bytes is neither bytes nor a base64-encoded string.
+
+    Examples
+    --------
+    >>> schema_str = '{"type": "record", "name": "User", "fields": [{"name": "name", "type": "string"}]}'
+    >>> encoded_data = base64.b64encode(b'some-avro-binary-data')
+    >>> user_dict = deserialize_avro(encoded_data, schema_str)
     """
+    if not value_schema_str:
+        raise KafkaConsumerAvroMissingSchemaError("Schema string must be provided for Avro deserialization")
+
     try:
-        reader = DatumReader(reader_schema)
+        # Parse the provided schema
+        parsed_schema = parse_schema(value_schema_str)
+        reader = DatumReader(parsed_schema)
+
+        # Handle different input types
+        if isinstance(avro_bytes, str):
+            # Assume base64 encoded string
+            value = base64.b64decode(avro_bytes)
+        elif isinstance(avro_bytes, bytes):
+            # Already raw bytes
+            value = avro_bytes
+        else:
+            # Try base64 decoding as a fallback
+            try:
+                value = base64.b64decode(avro_bytes)
+            except Exception as e:
+                raise TypeError(
+                    f"Expected bytes or base64-encoded string, got {type(avro_bytes).__name__}. Error: {str(e)}",
+                ) from e
 
-        decoder = BinaryDecoder(io.BytesIO(avro_bytes))
+        # Create binary decoder and read data
+        bytes_reader = io.BytesIO(value)
+        decoder = BinaryDecoder(bytes_reader)
         return reader.read(decoder)
-    except SchemaResolutionException as e:
-        raise ValueError(f"Schema mismatch: {e}") from e
+
+    except KafkaConsumerAvroSchemaMismatchError as e:
+        raise ValueError(
+            f"Schema mismatch detected: Message schema doesn't match expected schema. "
+            f"Details: {str(e)}. Verify schema registry configuration and message format.",
+        ) from e
+    except KafkaConsumerDeserializationError as e:
+        raise ValueError(
+            f"Deserialization failed: Unable to decode message data using Avro schema. "
+            f"Error: {str(e)}. Check for data corruption or schema evolution issues.",
+        ) from e
+
+
+def deserialize_protobuf_with_compiled_classes(
+    protobuf_bytes: bytes | str,
+    message_class: Any,
+) -> dict[str, Any]:
+    """
+    A deserialize that works with pre-compiled protobuf classes.
+
+    Parameters
+    ----------
+    protobuf_bytes : Union[bytes, str]
+        Protocol Buffer binary data, either as raw bytes or base64-encoded string.
+    message_class : Any
+        The pre-compiled Protocol Buffer message class.
+
+    Returns
+    -------
+    Dict[str, Any]
+        Deserialized Python dictionary representing the Protocol Buffer data.
+
+    Example
+    -------
+    >>> from my_proto_package.user_pb2 import User
+    >>> user_dict = deserialize_protobuf_with_compiled_classes(encoded_data, User)
+    """
+
+    try:
+        # Handle different input types for the binary data
+        if isinstance(protobuf_bytes, str):
+            # Assume base64 encoded string
+            value = base64.b64decode(protobuf_bytes)
+        elif isinstance(protobuf_bytes, bytes):
+            # Already raw bytes
+            value = protobuf_bytes
+        else:
+            # Try base64 decoding as a fallback
+            try:
+                value = base64.b64decode(protobuf_bytes)
+            except Exception as e:
+                raise TypeError(
+                    f"Expected bytes or base64-encoded string, got {type(protobuf_bytes).__name__}. Error: {str(e)}",
+                ) from e
+
+        # Create message instance and deserialize
+        message = message_class()
+        message.ParseFromString(value)
+
+        # Convert to dictionary
+        return MessageToDict(message, preserving_proto_field_name=True)
+
     except Exception as e:
-        raise ValueError(f"Failed to deserialize Avro data: {e}") from e
+        raise KafkaConsumerDeserializationError(
+            f"Protocol Buffer deserialization error: {type(e).__name__}: {str(e)}",
+        ) from e

aws_lambda_powertools/utilities/kafka_consumer/kafka_consumer.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
+from aws_lambda_powertools.utilities.kafka_consumer.consumer_record import ConsumerRecord
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_lambda_powertools.utilities.typing import LambdaContext
+
+
+@lambda_handler_decorator
+def kafka_consumer(
+    handler: Callable[[Any, LambdaContext], Any],
+    event: dict[str, Any],
+    context: LambdaContext,
+    schema_registry_config: Any | None = None,
+):
+    """Middleware to create an instance of the passed in event source data class
+
+    Parameters
+    ----------
+    handler: Callable
+        Lambda's handler
+    event: dict[str, Any]
+        Lambda's Event
+    context: LambdaContext
+        Lambda's Context
+    data_class: type[DictWrapper]
+        Data class type to instantiate
+
+    Example
+    --------
+
+    **Sample usage**
+
+        from aws_lambda_powertools.utilities.data_classes import S3Event, event_source
+
+        @event_source(data_class=S3Event)
+        def handler(event: S3Event, context):
+             return {"key": event.object_key}
+    """
+    return handler(ConsumerRecord(event, schema_registry_config), context)