Internal refactoring

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka_consumer/consumer_records.py

@@ -3,7 +3,7 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from aws_lambda_powertools.utilities.data_classes.common import CaseInsensitiveDict
+from aws_lambda_powertools.utilities.data_classes.common import CaseInsensitiveDict, DictWrapper
 from aws_lambda_powertools.utilities.data_classes.kafka_event import KafkaEventBase, KafkaEventRecordBase
 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
@@ -14,6 +14,18 @@
     from aws_lambda_powertools.utilities.kafka_consumer.schema_config import SchemaConfig
 
 
+class ConsumerRecordSchemaMetadata(DictWrapper):
+    @property
+    def data_format(self) -> str | None:
+        """The data format of the Kafka record."""
+        return self.get("dataFormat", None)
+
+    @property
+    def schema_id(self) -> str | None:
+        """The schema id of the Kafka record."""
+        return self.get("schemaId", None)
+
+
 class ConsumerRecordRecords(KafkaEventRecordBase):
     """
     A Kafka Consumer Record
@@ -26,42 +38,54 @@ def __init__(self, data: dict[str, Any], schema_config: SchemaConfig | None = No
     @cached_property
     def key(self) -> Any:
         key = self.get("key")
-        if key and (self.schema_config and self.schema_config.key_schema_type):
-            deserializer = get_deserializer(
-                self.schema_config.key_schema_type,
-                self.schema_config.key_schema_str,
-            )
-            deserialized_key = deserializer.deserialize(key)
 
-            if self.schema_config.key_output_serializer:
-                return serialize_to_output_type(
-                    deserialized_key,
-                    self.schema_config.key_output_serializer,
-                )
+        # Return None if key doesn't exist
+        if not key:
+            return None
+
+        # Determine schema type and schema string
+        schema_type = None
+        schema_str = None
+        output_serializer = None
+
+        if self.schema_config and self.schema_config.key_schema_type:
+            schema_type = self.schema_config.key_schema_type
+            schema_str = self.schema_config.key_schema_str
+            output_serializer = self.schema_config.key_output_serializer
 
-            return deserialized_key
+        # Always use get_deserializer if None it will default to DEFAULT
+        deserializer = get_deserializer(schema_type, schema_str)
+        deserialized_value = deserializer.deserialize(key)
 
-        return key  # MISSING DESERIALIZER
+        # Apply output serializer if specified
+        if output_serializer:
+            return serialize_to_output_type(deserialized_value, output_serializer)
+
+        return deserialized_value
 
     @cached_property
     def value(self) -> Any:
         value = self["value"]
-        if value and (self.schema_config and self.schema_config.value_schema_type):
-            deserializer = get_deserializer(
-                self.schema_config.value_schema_type,
-                self.schema_config.value_schema_str,
-            )
-            deserialized_value = deserializer.deserialize(value)
 
-            if self.schema_config.value_output_serializer:
-                return serialize_to_output_type(
-                    deserialized_value,
-                    self.schema_config.value_output_serializer,
-                )
+        # Determine schema type and schema string
+        schema_type = None
+        schema_str = None
+        output_serializer = None
+
+        if self.schema_config and self.schema_config.value_schema_type:
+            schema_type = self.schema_config.value_schema_type
+            schema_str = self.schema_config.value_schema_str
+            output_serializer = self.schema_config.value_output_serializer
 
-            return deserialized_value
+        # Always use get_deserializer if None it will default to DEFAULT
+        deserializer = get_deserializer(schema_type, schema_str)
+        deserialized_value = deserializer.deserialize(value)
 
-        return value  # MISSING DESERIALIZER
+        # Apply output serializer if specified
+        if output_serializer:
+            return serialize_to_output_type(deserialized_value, output_serializer)
+
+        return deserialized_value
 
     @property
     def original_value(self) -> str:
@@ -90,6 +114,22 @@ def headers(self) -> dict[str, bytes]:
         """Decodes the headers as a single dictionary."""
         return CaseInsensitiveDict((k, bytes(v)) for chunk in self.original_headers for k, v in chunk.items())
 
+    @property
+    def key_schema_metadata(self) -> ConsumerRecordSchemaMetadata | None:
+        """The metadata of the Key Kafka record."""
+        return (
+            None if self.get("keySchemaMetadata") is None else ConsumerRecordSchemaMetadata(self["keySchemaMetadata"])
+        )
+
+    @property
+    def value_schema_metadata(self) -> ConsumerRecordSchemaMetadata | None:
+        """The metadata of the Value Kafka record."""
+        return (
+            None
+            if self.get("valueSchemaMetadata") is None
+            else ConsumerRecordSchemaMetadata(self["valueSchemaMetadata"])
+        )
+
 
 class ConsumerRecords(KafkaEventBase):
     """Self-managed or MSK Apache Kafka event trigger

aws_lambda_powertools/utilities/kafka_consumer/deserializer/deserializer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import hashlib
 from typing import TYPE_CHECKING, Any
 
 from aws_lambda_powertools.utilities.kafka_consumer.deserializer.default import DefaultDeserializer
@@ -8,6 +9,23 @@
 if TYPE_CHECKING:
     from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
 
+# Cache for deserializers
+_deserializer_cache: dict[str, DeserializerBase] = {}
+
+
+def _get_cache_key(schema_type: str | object, schema_value: Any) -> str:
+    if schema_value is None:
+        return str(schema_type)
+
+    if isinstance(schema_value, str):
+        # For string schemas like Avro, hash the content
+        schema_hash = hashlib.md5(schema_value.encode("utf-8")).hexdigest()
+    else:
+        # For objects like Protobuf, use the object id
+        schema_hash = str(id(schema_value))
+
+    return f"{schema_type}_{schema_hash}"
+
 
 def get_deserializer(schema_type: str | object, schema_value: Any) -> DeserializerBase:
     """
@@ -55,18 +73,35 @@ def get_deserializer(schema_type: str | object, schema_value: Any) -> Deserializ
     >>> # Get a no-op deserializer for raw data
     >>> no_op_deserializer = get_deserializer("RAW", None)
     """
+
+    # Generate a cache key based on schema type and value
+    cache_key = _get_cache_key(schema_type, schema_value)
+
+    # Check if we already have this deserializer in cache
+    if cache_key in _deserializer_cache:
+        return _deserializer_cache[cache_key]
+
+    deserializer: 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)
+        deserializer = 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)
+        deserializer = ProtobufDeserializer(schema_value)
     elif schema_type == "JSON":
-        return JsonDeserializer()
+        deserializer = JsonDeserializer()
+
+    else:
+        # Default to no-op deserializer
+        deserializer = DefaultDeserializer()
+
+    # Store in cache for future use
+    _deserializer_cache[cache_key] = deserializer
 
     # Default to default deserializer that is base64 decode + bytes decoded
-    return DefaultDeserializer()
+    return deserializer

tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py

@@ -11,6 +11,7 @@
 from aws_lambda_powertools.utilities.kafka_consumer.exceptions import (
     KafkaConsumerAvroSchemaParserError,
     KafkaConsumerDeserializationError,
+    KafkaConsumerMissingSchemaError,
 )
 from aws_lambda_powertools.utilities.kafka_consumer.kafka_consumer import kafka_consumer
 from aws_lambda_powertools.utilities.kafka_consumer.schema_config import SchemaConfig
@@ -292,3 +293,17 @@ def lambda_handler(event: ConsumerRecords, context):
     assert key_value_result["value_type"] == "UserValueDataClass"
     assert key_value_result["value_name"] == "John Doe"
     assert key_value_result["value_age"] == 30
+
+
+def test_kafka_consumer_without_avro_value_schema():
+    """Test error handling when Avro data is invalid."""
+
+    with pytest.raises(KafkaConsumerMissingSchemaError):
+        SchemaConfig(value_schema_type="AVRO", value_schema=None)
+
+
+def test_kafka_consumer_without_avro_key_schema():
+    """Test error handling when Avro data is invalid."""
+
+    with pytest.raises(KafkaConsumerMissingSchemaError):
+        SchemaConfig(key_schema_type="AVRO", key_schema=None)

tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py

@@ -7,6 +7,7 @@
 from aws_lambda_powertools.utilities.kafka_consumer.consumer_records import ConsumerRecords
 from aws_lambda_powertools.utilities.kafka_consumer.exceptions import (
     KafkaConsumerDeserializationError,
+    KafkaConsumerMissingSchemaError,
 )
 from aws_lambda_powertools.utilities.kafka_consumer.kafka_consumer import kafka_consumer
 from aws_lambda_powertools.utilities.kafka_consumer.schema_config import SchemaConfig
@@ -316,3 +317,17 @@ def handler(event: ConsumerRecords, context):
     assert processed_records[0]["age"] == 30
     assert processed_records[1]["name"] == "Jane Smith"
     assert processed_records[1]["age"] == 25
+
+
+def test_kafka_consumer_without_protobuf_value_schema():
+    """Test error handling when Avro data is invalid."""
+
+    with pytest.raises(KafkaConsumerMissingSchemaError):
+        SchemaConfig(value_schema_type="PROTOBUF", value_schema=None)
+
+
+def test_kafka_consumer_without_protobuf_key_schema():
+    """Test error handling when Avro data is invalid."""
+
+    with pytest.raises(KafkaConsumerMissingSchemaError):
+        SchemaConfig(key_schema_type="PROTOBUF", key_schema=None)

tests/functional/kafka_consumer/required_dependencies/test_kafka_consumer.py

@@ -213,23 +213,62 @@ def handler(event: ConsumerRecords, context):
     assert any(r["name"] == "Bob Johnson" and r["age"] == 40 for r in processed_records)
 
 
-def test_kafka_consumer_without_schema_config(kafka_event_with_json_data, lambda_context):
+def test_kafka_consumer_default_deserializer_value(kafka_event_with_json_data, lambda_context):
     """Test Kafka consumer when no schema config is provided."""
 
-    # Create dict to capture results
-    result_data = {}
+    base64_data = base64.b64encode(b"data")
+    kafka_event_with_json_data = deepcopy(kafka_event_with_json_data)
+    kafka_event_with_json_data["records"]["my-topic-1"][0]["value"] = base64_data
 
     @kafka_consumer()
     def handler(event: ConsumerRecords, context):
         # Capture the results to verify
         record = next(event.records)
         # Should get raw base64-encoded data with no deserialization
-        result_data["value_type"] = type(record.value).__name__
-        return {"processed": True}
+        return record.value
 
     # Call the handler
     result = handler(kafka_event_with_json_data, lambda_context)
 
     # Verify the results
-    assert result == {"processed": True}
-    assert result_data["value_type"] == "str"  # Raw base64 string
+    assert result == "data"
+
+
+def test_kafka_consumer_default_deserializer_key(kafka_event_with_json_data, lambda_context):
+    """Test Kafka consumer when no schema config is provided."""
+
+    base64_data = base64.b64encode(b"data")
+    kafka_event_with_json_data = deepcopy(kafka_event_with_json_data)
+    kafka_event_with_json_data["records"]["my-topic-1"][0]["key"] = base64_data
+
+    @kafka_consumer()
+    def handler(event: ConsumerRecords, context):
+        # Capture the results to verify
+        record = next(event.records)
+        # Should get raw base64-encoded data with no deserialization
+        return record.key
+
+    # Call the handler
+    result = handler(kafka_event_with_json_data, lambda_context)
+
+    # Verify the results
+    assert result == "data"
+
+
+def test_kafka_consumer_default_deserializer_key_is_none(kafka_event_with_json_data, lambda_context):
+    """Test Kafka consumer when no schema config is provided."""
+
+    kafka_event_with_json_data["records"]["my-topic-1"][0]["key"] = None
+
+    @kafka_consumer()
+    def handler(event: ConsumerRecords, context):
+        # Capture the results to verify
+        record = next(event.records)
+        # Should get raw base64-encoded data with no deserialization
+        return record.key
+
+    # Call the handler
+    result = handler(kafka_event_with_json_data, lambda_context)
+
+    # Verify the results
+    assert result is None