Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions aws_lambda_powertools/utilities/kafka/consumer_records.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,17 +41,20 @@ def key(self) -> Any:
schema_type = None
schema_value = None
output_serializer = None
key_schema_wire_format = None

if self.schema_config and self.schema_config.key_schema_type:
schema_type = self.schema_config.key_schema_type
schema_value = self.schema_config.key_schema
output_serializer = self.schema_config.key_output_serializer
key_schema_wire_format = self.schema_config.key_schema_wire_format

# Always use get_deserializer if None it will default to DEFAULT
deserializer = get_deserializer(
schema_type=schema_type,
schema_value=schema_value,
field_metadata=self.key_schema_metadata,
wire_format=key_schema_wire_format,
)
deserialized_value = deserializer.deserialize(key)

Expand All @@ -69,19 +72,22 @@ def value(self) -> Any:
schema_type = None
schema_value = None
output_serializer = None
value_schema_wire_format = None

logger.debug("Deserializing value field")

if self.schema_config and self.schema_config.value_schema_type:
schema_type = self.schema_config.value_schema_type
schema_value = self.schema_config.value_schema
output_serializer = self.schema_config.value_output_serializer
value_schema_wire_format = self.schema_config.value_schema_wire_format

# Always use get_deserializer if None it will default to DEFAULT
deserializer = get_deserializer(
schema_type=schema_type,
schema_value=schema_value,
field_metadata=self.value_schema_metadata,
wire_format=value_schema_wire_format,
)
deserialized_value = deserializer.deserialize(value)

Expand Down
37 changes: 35 additions & 2 deletions aws_lambda_powertools/utilities/kafka/deserializer/avro.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import io
import logging
from typing import Any
from typing import Any, Literal

from avro.io import BinaryDecoder, DatumReader
from avro.schema import parse as parse_schema
Expand All @@ -16,6 +16,9 @@

logger = logging.getLogger(__name__)

_CONFLUENT_HEADER_SIZE = 5
_CONFLUENT_MAGIC_BYTE = 0x00


class AvroDeserializer(DeserializerBase):
"""
Expand All @@ -25,16 +28,43 @@ class AvroDeserializer(DeserializerBase):
a provided Avro schema definition.
"""

def __init__(self, schema_str: str, field_metadata: dict[str, Any] | None = None):
def __init__(
self,
schema_str: str,
field_metadata: dict[str, Any] | None = None,
wire_format: Literal["CONFLUENT"] | None = None,
):
try:
self.parsed_schema = parse_schema(schema_str)
self.reader = DatumReader(self.parsed_schema)
self.field_metatada = field_metadata
self.wire_format = wire_format
except Exception as e:
raise KafkaConsumerAvroSchemaParserError(
f"Invalid Avro schema. Please ensure the provided avro schema is valid: {type(e).__name__}: {str(e)}",
) from e

def _strip_wire_format_header(self, value: bytes) -> bytes:
if self.wire_format is None:
return value

if self.wire_format != "CONFLUENT":
raise KafkaConsumerDeserializationError(f"Unsupported Avro wire format: {self.wire_format}")

if len(value) < _CONFLUENT_HEADER_SIZE:
raise KafkaConsumerDeserializationError(
"Invalid Confluent wire format: payload must contain a 5-byte header",
)

if value[0] != _CONFLUENT_MAGIC_BYTE:
raise KafkaConsumerDeserializationError(
"Invalid Confluent wire format: expected magic byte 0x00",
)

schema_id = int.from_bytes(value[1:_CONFLUENT_HEADER_SIZE], byteorder="big")
logger.debug("Deserializing Confluent payload with schema ID %s", schema_id)
return value[_CONFLUENT_HEADER_SIZE:]

def deserialize(self, data: bytes | str) -> object:
"""
Deserialize Avro binary data to a Python dictionary.
Expand Down Expand Up @@ -75,9 +105,12 @@ def deserialize(self, data: bytes | str) -> object:

try:
value = self._decode_input(data)
value = self._strip_wire_format_header(value)
bytes_reader = io.BytesIO(value)
decoder = BinaryDecoder(bytes_reader)
return self.reader.read(decoder)
except KafkaConsumerDeserializationError:
raise
except Exception as e:
raise KafkaConsumerDeserializationError(
f"Error trying to deserialize avro data - {type(e).__name__}: {str(e)}",
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

import hashlib
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Literal

from aws_lambda_powertools.utilities.kafka.deserializer.default import DefaultDeserializer
from aws_lambda_powertools.utilities.kafka.deserializer.json import JsonDeserializer
Expand All @@ -13,7 +13,12 @@
_deserializer_cache: dict[str, DeserializerBase] = {}


def _get_cache_key(schema_type: str | object, schema_value: Any, field_metadata: dict[str, Any]) -> str:
def _get_cache_key(
schema_type: str | object,
schema_value: Any,
field_metadata: dict[str, Any],
wire_format: Literal["CONFLUENT"] | None,
) -> str:
schema_metadata = None

if field_metadata:
Expand All @@ -30,10 +35,15 @@ def _get_cache_key(schema_type: str | object, schema_value: Any, field_metadata:
# For objects like Protobuf, use the object id
schema_hash = f"{str(id(schema_value))}_{schema_metadata}"

return f"{schema_type}_{schema_hash}"
return f"{schema_type}_{schema_hash}_{wire_format}"


def get_deserializer(schema_type: str | object, schema_value: Any, field_metadata: Any) -> DeserializerBase:
def get_deserializer(
schema_type: str | object,
schema_value: Any,
field_metadata: Any,
wire_format: Literal["CONFLUENT"] | None = None,
) -> DeserializerBase:
"""
Factory function to get the appropriate deserializer based on schema type.

Expand Down Expand Up @@ -81,7 +91,7 @@ def get_deserializer(schema_type: str | object, schema_value: Any, field_metadat
"""

# Generate a cache key based on schema type and value
cache_key = _get_cache_key(schema_type, schema_value, field_metadata)
cache_key = _get_cache_key(schema_type, schema_value, field_metadata, wire_format)

# Check if we already have this deserializer in cache
if cache_key in _deserializer_cache:
Expand All @@ -93,7 +103,11 @@ def get_deserializer(schema_type: str | object, schema_value: Any, field_metadat
# Import here to avoid dependency if not used
from aws_lambda_powertools.utilities.kafka.deserializer.avro import AvroDeserializer

deserializer = AvroDeserializer(schema_str=schema_value, field_metadata=field_metadata)
deserializer = AvroDeserializer(
schema_str=schema_value,
field_metadata=field_metadata,
wire_format=wire_format,
)
elif schema_type == "PROTOBUF":
# Import here to avoid dependency if not used
from aws_lambda_powertools.utilities.kafka.deserializer.protobuf import ProtobufDeserializer
Expand Down
26 changes: 26 additions & 0 deletions aws_lambda_powertools/utilities/kafka/schema_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,14 @@ class SchemaConfig:
Schema definition for message keys. Required when key_schema_type is 'AVRO' or 'PROTOBUF'.
key_output_serializer : Any, optional
Custom serializer for message keys. Supports Pydantic classes, Dataclasses and Custom Class
value_schema_wire_format : {'CONFLUENT', None}, default=None
Set this when a Confluent schema-registry-aware serializer produced the value payload
but you are supplying the Avro schema offline rather than using the ESM Schema Registry integration.
Only applies to AVRO values.
key_schema_wire_format : {'CONFLUENT', None}, default=None
Set this when a Confluent schema-registry-aware serializer produced the key payload
but you are supplying the Avro schema offline rather than using the ESM Schema Registry integration.
Only applies to AVRO keys.

Raises
------
Expand Down Expand Up @@ -63,21 +71,39 @@ def __init__(
key_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None,
key_schema: str | None = None,
key_output_serializer: Any | None = None,
value_schema_wire_format: Literal["CONFLUENT"] | None = None,
key_schema_wire_format: Literal["CONFLUENT"] | None = None,
):
# Validate schema requirements
self._validate_schema_requirements(value_schema_type, value_schema, "value")
self._validate_schema_requirements(key_schema_type, key_schema, "key")
self._validate_wire_format(value_schema_wire_format, value_schema_type, "value")
self._validate_wire_format(key_schema_wire_format, key_schema_type, "key")

self.value_schema_type = value_schema_type
self.value_schema = value_schema
self.value_output_serializer = value_output_serializer
self.key_schema_type = key_schema_type
self.key_schema = key_schema
self.key_output_serializer = key_output_serializer
self.value_schema_wire_format = value_schema_wire_format
self.key_schema_wire_format = key_schema_wire_format

def _validate_schema_requirements(self, schema_type: str | None, schema: str | None, prefix: str) -> None:
"""Validate that schema is provided when required by schema_type."""
if schema_type in ["AVRO", "PROTOBUF"] and schema is None:
raise KafkaConsumerMissingSchemaError(
f"{prefix}_schema must be provided when {prefix}_schema_type is {schema_type}",
)

def _validate_wire_format(self, wire_format: str | None, schema_type: str | None, prefix: str) -> None:
"""Validate the wire format for a key or value payload."""

if wire_format is None:
return

if wire_format != "CONFLUENT":
raise ValueError(f"{prefix}_schema_wire_format must be 'CONFLUENT'.")

if schema_type != "AVRO":
raise ValueError(f"{prefix}_schema_wire_format is supported only when {prefix}_schema_type is 'AVRO'.")
40 changes: 40 additions & 0 deletions docs/utilities/kafka.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ flowchart LR
* Support for key and value deserialization
* Support for custom output serializers (e.g., dataclasses, Pydantic models)
* Support for ESM with and without Schema Registry integration
* Support for offline Avro schemas with schema-registry wire-format prefixes (Confluent only)
* Proper error handling for deserialization issues

## Terminology
Expand Down Expand Up @@ -255,6 +256,45 @@ Each Kafka record contains important metadata that you can access alongside the
| `value_schema_metadata` | Metadata about the value schema like `schemaId` and `dataFormat` | Data format and schemaId propagated when integrating with Schema Registry |
| `key_schema_metadata` | Metadata about the key schema like `schemaId` and `dataFormat` | Data format and schemaId propagated when integrating with Schema Registry |

### Using an offline Avro schema with a schema-registry wire-format prefix

When Confluent serializes messages with its schema-registry-aware Avro serializer (for example, `KafkaAvroSerializer`), each payload carries a wire-format header before the Avro body.
The header is 5 bytes long: 1-byte magic byte (`0x00`) followed by a 4-byte big-endian schema ID.

When the ESM Schema Registry integration is enabled, Lambda strips those bytes and populates the record's schema metadata. When you use an **offline Avro schema** without the ESM Schema Registry integration, the header reaches the function and prevents plain Avro deserialization.

Set `value_schema_wire_format` or `key_schema_wire_format` on `SchemaConfig` to `"CONFLUENT"`. Powertools validates the magic byte and strips the 5-byte header before running the Avro decoder.

???+ info "When do I need this?"
Use this option when you supply the Avro schema and the producer uses the Confluent wire format. If ESM Schema Registry integration has already removed the header, leave the option as `None`.

=== "Offline Avro schema with a Confluent prefix"

```python hl_lines="10"
from aws_lambda_powertools.utilities.kafka import SchemaConfig, kafka_consumer
from aws_lambda_powertools.utilities.kafka.consumer_records import ConsumerRecords
from aws_lambda_powertools.utilities.typing import LambdaContext

AVRO_SCHEMA = open("user.avsc").read()

schema_config = SchemaConfig(
value_schema_type="AVRO",
value_schema=AVRO_SCHEMA,
value_schema_wire_format="CONFLUENT",
)


@kafka_consumer(schema_config=schema_config)
def lambda_handler(event: ConsumerRecords, context: LambdaContext):
for record in event.records:
# record.value is the deserialized Avro payload
# with the validated 5-byte wire-format header removed.
...
```

???+ warning "Scope"
`value_schema_wire_format` and `key_schema_wire_format` apply only to **Avro** payloads. Leave them as `None` when ESM Schema Registry integration has already removed the wire-format header.

### Custom output serializers

Transform deserialized data into your preferred object types using output serializers. This can help you integrate Kafka data with your domain models and application architecture, providing type hints, validation, and structured data access.
Expand Down
Loading