From a9378d8a14de5b4e129a4eee4f27dc8d28139b14 Mon Sep 17 00:00:00 2001 From: Amin Farjadi Date: Thu, 13 Aug 2026 12:07:05 +0100 Subject: [PATCH 1/4] feat: add optional args to 'SchemaConfig' for ignoring prefix --- .../utilities/kafka/consumer_records.py | 6 + .../utilities/kafka/deserializer/avro.py | 10 +- .../kafka/deserializer/deserializer.py | 24 +++- .../utilities/kafka/schema_config.py | 20 ++++ .../_avro/test_kafka_consumer_with_avro.py | 108 ++++++++++++++++++ 5 files changed, 162 insertions(+), 6 deletions(-) diff --git a/aws_lambda_powertools/utilities/kafka/consumer_records.py b/aws_lambda_powertools/utilities/kafka/consumer_records.py index 1fa6afba15c..9fbccefa245 100644 --- a/aws_lambda_powertools/utilities/kafka/consumer_records.py +++ b/aws_lambda_powertools/utilities/kafka/consumer_records.py @@ -41,17 +41,20 @@ def key(self) -> Any: schema_type = None schema_value = None output_serializer = None + schema_id_prefix_length = 0 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 + schema_id_prefix_length = self.schema_config.key_schema_id_prefix_length # 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, + schema_id_prefix_length=schema_id_prefix_length, ) deserialized_value = deserializer.deserialize(key) @@ -69,6 +72,7 @@ def value(self) -> Any: schema_type = None schema_value = None output_serializer = None + schema_id_prefix_length = 0 logger.debug("Deserializing value field") @@ -76,12 +80,14 @@ def value(self) -> Any: schema_type = self.schema_config.value_schema_type schema_value = self.schema_config.value_schema output_serializer = self.schema_config.value_output_serializer + schema_id_prefix_length = self.schema_config.value_schema_id_prefix_length # 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, + schema_id_prefix_length=schema_id_prefix_length, ) deserialized_value = deserializer.deserialize(value) diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/avro.py b/aws_lambda_powertools/utilities/kafka/deserializer/avro.py index d3b96da9d34..d70c638bbf2 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/avro.py +++ b/aws_lambda_powertools/utilities/kafka/deserializer/avro.py @@ -25,11 +25,17 @@ 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, + schema_id_prefix_length: int = 0, + ): try: self.parsed_schema = parse_schema(schema_str) self.reader = DatumReader(self.parsed_schema) self.field_metatada = field_metadata + self.schema_id_prefix_length = schema_id_prefix_length except Exception as e: raise KafkaConsumerAvroSchemaParserError( f"Invalid Avro schema. Please ensure the provided avro schema is valid: {type(e).__name__}: {str(e)}", @@ -75,6 +81,8 @@ def deserialize(self, data: bytes | str) -> object: try: value = self._decode_input(data) + if self.schema_id_prefix_length: + value = value[self.schema_id_prefix_length :] bytes_reader = io.BytesIO(value) decoder = BinaryDecoder(bytes_reader) return self.reader.read(decoder) diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py b/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py index c1443c83b00..7c0f9f26fa4 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py +++ b/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py @@ -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], + schema_id_prefix_length: int = 0, +) -> str: schema_metadata = None if field_metadata: @@ -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}_{schema_id_prefix_length}" -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, + schema_id_prefix_length: int = 0, +) -> DeserializerBase: """ Factory function to get the appropriate deserializer based on schema type. @@ -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, schema_id_prefix_length) # Check if we already have this deserializer in cache if cache_key in _deserializer_cache: @@ -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, + schema_id_prefix_length=schema_id_prefix_length, + ) elif schema_type == "PROTOBUF": # Import here to avoid dependency if not used from aws_lambda_powertools.utilities.kafka.deserializer.protobuf import ProtobufDeserializer diff --git a/aws_lambda_powertools/utilities/kafka/schema_config.py b/aws_lambda_powertools/utilities/kafka/schema_config.py index 96eed96984f..aa2c9376ef1 100644 --- a/aws_lambda_powertools/utilities/kafka/schema_config.py +++ b/aws_lambda_powertools/utilities/kafka/schema_config.py @@ -20,12 +20,19 @@ class SchemaConfig: Schema definition for message values. Required when value_schema_type is 'AVRO' or 'PROTOBUF'. value_output_serializer : Any, optional Custom output serializer for message values. Supports Pydantic classes, Dataclasses and Custom Class + value_schema_id_prefix_length : int, default=0 + Number of leading bytes to skip on the value payload before Avro deserialization. Use this + when the payload was produced by a schema-registry-aware serializer (e.g. Confluent's + 5-byte magic byte + schema ID prefix) but you are supplying the Avro schema offline rather + than relying on the ESM Schema Registry integration. Only applied for AVRO values. key_schema_type : {'AVRO', 'PROTOBUF', 'JSON', None}, default=None Schema type for message keys. key_schema : str, optional 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 + key_schema_id_prefix_length : int, default=0 + Same as ``value_schema_id_prefix_length`` but for the record key. Raises ------ @@ -60,13 +67,17 @@ def __init__( value_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None, value_schema: str | None = None, value_output_serializer: Any | None = None, + value_schema_id_prefix_length: int = 0, key_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None, key_schema: str | None = None, key_output_serializer: Any | None = None, + key_schema_id_prefix_length: int = 0, ): # 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_prefix_length(value_schema_id_prefix_length, "value") + self._validate_prefix_length(key_schema_id_prefix_length, "key") self.value_schema_type = value_schema_type self.value_schema = value_schema @@ -74,6 +85,8 @@ def __init__( self.key_schema_type = key_schema_type self.key_schema = key_schema self.key_output_serializer = key_output_serializer + self.value_schema_id_prefix_length = value_schema_id_prefix_length + self.key_schema_id_prefix_length = key_schema_id_prefix_length 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.""" @@ -81,3 +94,10 @@ def _validate_schema_requirements(self, schema_type: str | None, schema: str | N raise KafkaConsumerMissingSchemaError( f"{prefix}_schema must be provided when {prefix}_schema_type is {schema_type}", ) + + def _validate_prefix_length(self, prefix_length: int, prefix: str) -> None: + """Validate that a schema-id prefix length is a non-negative integer.""" + if not isinstance(prefix_length, int) or isinstance(prefix_length, bool) or prefix_length < 0: + raise ValueError( + f"{prefix}_schema_id_prefix_length must be a non-negative integer, got {prefix_length!r}", + ) diff --git a/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py b/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py index f22171c37af..de4ea2983c0 100644 --- a/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py +++ b/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py @@ -67,6 +67,23 @@ def avro_encoded_key(avro_key_schema): return base64.b64encode(bytes_writer.getvalue()).decode("utf-8") +SCHEMA_ID_PREFIX = b"\x00\x00\x00\x00\x01" + + +def _prepend_prefix_to_base64(encoded: str, prefix: bytes = SCHEMA_ID_PREFIX) -> str: + return base64.b64encode(prefix + base64.b64decode(encoded)).decode("utf-8") + + +@pytest.fixture +def avro_encoded_value_with_prefix(avro_encoded_value): + return _prepend_prefix_to_base64(avro_encoded_value) + + +@pytest.fixture +def avro_encoded_key_with_prefix(avro_encoded_key): + return _prepend_prefix_to_base64(avro_encoded_key) + + @pytest.fixture def kafka_event_with_avro_data(avro_encoded_value, avro_encoded_key): return { @@ -312,6 +329,97 @@ def test_kafka_consumer_without_avro_key_schema(): assert "key_schema" in str(excinfo.value) +def test_kafka_consumer_avro_produces_wrong_output_without_prefix_length_setting( + kafka_event_with_avro_data, + avro_encoded_value_with_prefix, + avro_value_schema, + lambda_context, +): + # GIVEN An Avro payload that has been serialized by a Confluent-style producer, + # so it carries a 5-byte "magic byte + schema ID" prefix in front of the Avro body + event = deepcopy(kafka_event_with_avro_data) + event["records"]["my-topic-1"][0]["value"] = avro_encoded_value_with_prefix + + # AND a SchemaConfig without the new offset parameter (today's behaviour) + schema_config = SchemaConfig(value_schema_type="AVRO", value_schema=avro_value_schema) + + @kafka_consumer(schema_config=schema_config) + def handler(event: ConsumerRecords, context): + return event.record.value + + # WHEN/THEN The deserializer cannot know it should skip the leading bytes. + # Depending on the prefix content, this either raises or silently returns corrupted data. + # Both outcomes are broken; the fix must let callers opt into skipping the prefix. + try: + result = handler(event, lambda_context) + except KafkaConsumerDeserializationError: + return + + assert result != {"name": "John Doe", "age": 30} + + +def test_kafka_consumer_avro_with_value_schema_id_prefix_length( + kafka_event_with_avro_data, + avro_encoded_value_with_prefix, + avro_value_schema, + lambda_context, +): + # GIVEN An Avro payload with a 5-byte magic-byte + schema-ID prefix + event = deepcopy(kafka_event_with_avro_data) + event["records"]["my-topic-1"][0]["value"] = avro_encoded_value_with_prefix + + # AND a SchemaConfig instructed to skip the first 5 bytes before Avro decoding + schema_config = SchemaConfig( + value_schema_type="AVRO", + value_schema=avro_value_schema, + value_schema_id_prefix_length=5, + ) + + @kafka_consumer(schema_config=schema_config) + def handler(event: ConsumerRecords, context): + return event.record.value + + # WHEN The handler processes the event + result = handler(event, lambda_context) + + # THEN The Avro body should be decoded correctly after the prefix is stripped + assert result["name"] == "John Doe" + assert result["age"] == 30 + + +def test_kafka_consumer_avro_with_key_schema_id_prefix_length( + kafka_event_with_avro_data, + avro_encoded_key_with_prefix, + avro_value_schema, + avro_key_schema, + lambda_context, +): + # GIVEN A Kafka event whose key is Avro data behind a 5-byte prefix + event = deepcopy(kafka_event_with_avro_data) + event["records"]["my-topic-1"][0]["key"] = avro_encoded_key_with_prefix + + # AND a SchemaConfig that only strips the prefix on the key side + schema_config = SchemaConfig( + value_schema_type="AVRO", + value_schema=avro_value_schema, + key_schema_type="AVRO", + key_schema=avro_key_schema, + key_schema_id_prefix_length=5, + ) + + @kafka_consumer(schema_config=schema_config) + def handler(event: ConsumerRecords, context): + record = next(event.records) + return {"key": record.key, "value": record.value} + + # WHEN The handler processes the event + result = handler(event, lambda_context) + + # THEN Both key (prefixed, offset applied) and value (plain) deserialize correctly + assert result["key"] == {"user_id": "user-123"} + assert result["value"] == {"name": "John Doe", "age": 30} + + def test_kafka_consumer_avro_with_wrong_json_schema( kafka_event_with_avro_data, lambda_context, From 30090053236979015eeb8fc5a4da1802e25f0950 Mon Sep 17 00:00:00 2001 From: Amin Farjadi Date: Thu, 13 Aug 2026 12:24:06 +0100 Subject: [PATCH 2/4] chore: add documentation --- docs/utilities/kafka.md | 44 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/docs/utilities/kafka.md b/docs/utilities/kafka.md index 5bbab7e3062..33ad6b68108 100644 --- a/docs/utilities/kafka.md +++ b/docs/utilities/kafka.md @@ -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 * Proper error handling for deserialization issues ## Terminology @@ -255,6 +256,49 @@ 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 your Kafka producer serializes messages with a schema-registry-aware Avro serializer (e.g. Confluent's `KafkaAvroSerializer` or the AWS Glue Avro serializer), each payload carries a short wire-format prefix in front of the Avro body: + +* **Confluent**: `1-byte magic byte (0x00) + 4-byte big-endian schema ID` — 5 bytes total. +* **AWS Glue**: `1-byte header version + 1-byte compression + 16-byte UUID schema-version ID` — 18 bytes total. + +When you enable the ESM Schema Registry integration, Lambda strips those bytes for you and populates `value_schema_metadata.schemaId`. But when you rely on an **offline Avro schema** (checked into your Lambda) and do **not** use the ESM Schema Registry integration, those prefix bytes reach your function and would otherwise corrupt Avro deserialization. + +Use `value_schema_id_prefix_length` (and/or `key_schema_id_prefix_length`) on `SchemaConfig` to tell Powertools how many leading bytes to skip after base64 decoding, before running the Avro decoder. + +???+ info "When do I need this?" + Only when you are supplying the Avro schema yourself **and** the producer prepended a schema-registry wire-format wrapper. If the ESM Schema Registry integration is on, leave these parameters at their default (`0`). + +=== "Offline Avro schema with a Confluent-style 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_id_prefix_length=5, # 1-byte magic byte + 4-byte schema ID + ) + + + @kafka_consumer(schema_config=schema_config) + def lambda_handler(event: ConsumerRecords, context: LambdaContext): + for record in event.records: + # record.value is the fully-deserialized Avro payload + # with the 5-byte wire-format prefix transparently stripped. + ... + ``` + +The offset is applied symmetrically for keys — set `key_schema_id_prefix_length` when a schema-registry-aware serializer was used for record keys too. + +???+ warning "Scope" + `value_schema_id_prefix_length` / `key_schema_id_prefix_length` only affect the **Avro** deserializer. Protobuf already infers the wire-format from `value_schema_metadata.schemaId` when the ESM Schema Registry integration is on, and JSON-with-registry is not supported today. + ### 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. From e9c5f6711f4314c49fef20f0f15b60ac3c659edc Mon Sep 17 00:00:00 2001 From: Amin Farjadi Date: Fri, 4 Sep 2026 18:34:44 +0100 Subject: [PATCH 3/4] address comment on issue --- .../utilities/kafka/consumer_records.py | 9 ++-- .../utilities/kafka/deserializer/avro.py | 13 +++--- .../kafka/deserializer/deserializer.py | 12 +++--- .../utilities/kafka/schema_config.py | 39 ++++++++--------- docs/utilities/kafka.md | 24 +++++------ .../_avro/test_kafka_consumer_with_avro.py | 42 +------------------ 6 files changed, 49 insertions(+), 90 deletions(-) diff --git a/aws_lambda_powertools/utilities/kafka/consumer_records.py b/aws_lambda_powertools/utilities/kafka/consumer_records.py index 9fbccefa245..84a526ecfec 100644 --- a/aws_lambda_powertools/utilities/kafka/consumer_records.py +++ b/aws_lambda_powertools/utilities/kafka/consumer_records.py @@ -41,20 +41,17 @@ def key(self) -> Any: schema_type = None schema_value = None output_serializer = None - schema_id_prefix_length = 0 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 - schema_id_prefix_length = self.schema_config.key_schema_id_prefix_length # 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, - schema_id_prefix_length=schema_id_prefix_length, ) deserialized_value = deserializer.deserialize(key) @@ -72,7 +69,7 @@ def value(self) -> Any: schema_type = None schema_value = None output_serializer = None - schema_id_prefix_length = 0 + value_schema_wire_format = None logger.debug("Deserializing value field") @@ -80,14 +77,14 @@ def value(self) -> Any: schema_type = self.schema_config.value_schema_type schema_value = self.schema_config.value_schema output_serializer = self.schema_config.value_output_serializer - schema_id_prefix_length = self.schema_config.value_schema_id_prefix_length + 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, - schema_id_prefix_length=schema_id_prefix_length, + wire_format=value_schema_wire_format, ) deserialized_value = deserializer.deserialize(value) diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/avro.py b/aws_lambda_powertools/utilities/kafka/deserializer/avro.py index d70c638bbf2..44c7ba4c644 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/avro.py +++ b/aws_lambda_powertools/utilities/kafka/deserializer/avro.py @@ -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 @@ -29,13 +29,13 @@ def __init__( self, schema_str: str, field_metadata: dict[str, Any] | None = None, - schema_id_prefix_length: int = 0, + value_schema_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.schema_id_prefix_length = schema_id_prefix_length + self.value_schema_wire_format = value_schema_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)}", @@ -81,8 +81,11 @@ def deserialize(self, data: bytes | str) -> object: try: value = self._decode_input(data) - if self.schema_id_prefix_length: - value = value[self.schema_id_prefix_length :] + if self.value_schema_wire_format == "CONFLUENT": + # removing the first 5 bytes from payload: + # 1B magic byte 0x00 + # 4B big-endian schema ID + value = value[5:] bytes_reader = io.BytesIO(value) decoder = BinaryDecoder(bytes_reader) return self.reader.read(decoder) diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py b/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py index 7c0f9f26fa4..373407a6244 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py +++ b/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py @@ -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 @@ -17,7 +17,7 @@ def _get_cache_key( schema_type: str | object, schema_value: Any, field_metadata: dict[str, Any], - schema_id_prefix_length: int = 0, + wire_format: Literal["CONFLUENT"] | None, ) -> str: schema_metadata = None @@ -35,14 +35,14 @@ def _get_cache_key( # For objects like Protobuf, use the object id schema_hash = f"{str(id(schema_value))}_{schema_metadata}" - return f"{schema_type}_{schema_hash}_{schema_id_prefix_length}" + return f"{schema_type}_{schema_hash}_{wire_format}" def get_deserializer( schema_type: str | object, schema_value: Any, field_metadata: Any, - schema_id_prefix_length: int = 0, + wire_format: Literal["CONFLUENT"] | None = None, ) -> DeserializerBase: """ Factory function to get the appropriate deserializer based on schema type. @@ -91,7 +91,7 @@ def get_deserializer( """ # Generate a cache key based on schema type and value - cache_key = _get_cache_key(schema_type, schema_value, field_metadata, schema_id_prefix_length) + 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: @@ -106,7 +106,7 @@ def get_deserializer( deserializer = AvroDeserializer( schema_str=schema_value, field_metadata=field_metadata, - schema_id_prefix_length=schema_id_prefix_length, + value_schema_wire_format=wire_format, ) elif schema_type == "PROTOBUF": # Import here to avoid dependency if not used diff --git a/aws_lambda_powertools/utilities/kafka/schema_config.py b/aws_lambda_powertools/utilities/kafka/schema_config.py index aa2c9376ef1..96e17288953 100644 --- a/aws_lambda_powertools/utilities/kafka/schema_config.py +++ b/aws_lambda_powertools/utilities/kafka/schema_config.py @@ -20,19 +20,16 @@ class SchemaConfig: Schema definition for message values. Required when value_schema_type is 'AVRO' or 'PROTOBUF'. value_output_serializer : Any, optional Custom output serializer for message values. Supports Pydantic classes, Dataclasses and Custom Class - value_schema_id_prefix_length : int, default=0 - Number of leading bytes to skip on the value payload before Avro deserialization. Use this - when the payload was produced by a schema-registry-aware serializer (e.g. Confluent's - 5-byte magic byte + schema ID prefix) but you are supplying the Avro schema offline rather - than relying on the ESM Schema Registry integration. Only applied for AVRO values. + value_schema_wire_format : {'CONFLUENT', None}, default=None + Set this when the payload was produced by a Confluent's schema-registry-aware serializer (KafkaAvroSerializer) + but you are supplying the Avro schema offline rather than relying on the ESM Schema Registry integration. + Only applied for AVRO values. key_schema_type : {'AVRO', 'PROTOBUF', 'JSON', None}, default=None Schema type for message keys. key_schema : str, optional 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 - key_schema_id_prefix_length : int, default=0 - Same as ``value_schema_id_prefix_length`` but for the record key. Raises ------ @@ -67,17 +64,15 @@ def __init__( value_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None, value_schema: str | None = None, value_output_serializer: Any | None = None, - value_schema_id_prefix_length: int = 0, + value_schema_wire_format: Literal["CONFLUENT"] | None = None, key_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None, key_schema: str | None = None, key_output_serializer: Any | None = None, - key_schema_id_prefix_length: int = 0, ): # 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_prefix_length(value_schema_id_prefix_length, "value") - self._validate_prefix_length(key_schema_id_prefix_length, "key") + self._validate_wire_format(value_schema_wire_format, value_schema_type) self.value_schema_type = value_schema_type self.value_schema = value_schema @@ -85,8 +80,7 @@ def __init__( self.key_schema_type = key_schema_type self.key_schema = key_schema self.key_output_serializer = key_output_serializer - self.value_schema_id_prefix_length = value_schema_id_prefix_length - self.key_schema_id_prefix_length = key_schema_id_prefix_length + self.value_schema_wire_format = value_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.""" @@ -95,9 +89,16 @@ def _validate_schema_requirements(self, schema_type: str | None, schema: str | N f"{prefix}_schema must be provided when {prefix}_schema_type is {schema_type}", ) - def _validate_prefix_length(self, prefix_length: int, prefix: str) -> None: - """Validate that a schema-id prefix length is a non-negative integer.""" - if not isinstance(prefix_length, int) or isinstance(prefix_length, bool) or prefix_length < 0: - raise ValueError( - f"{prefix}_schema_id_prefix_length must be a non-negative integer, got {prefix_length!r}", - ) + def _validate_wire_format(self, wire_format: str | None, schema_type: str | None) -> None: + """Validate the wire format for value payload.""" + + if wire_format is None: + return + + if wire_format != "CONFLUENT": + raise ValueError("Only 'CONFLUENT' wire format is supported.") + + if schema_type != "AVRO": + raise ValueError("Wire format is supported for only for 'AVRO' schema.") + + return None diff --git a/docs/utilities/kafka.md b/docs/utilities/kafka.md index 33ad6b68108..9c43fc54071 100644 --- a/docs/utilities/kafka.md +++ b/docs/utilities/kafka.md @@ -29,7 +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 +* Support for offline Avro schemas with schema-registry wire-format prefixes (Confluent only) * Proper error handling for deserialization issues ## Terminology @@ -258,19 +258,17 @@ Each Kafka record contains important metadata that you can access alongside the ### Using an offline Avro schema with a schema-registry wire-format prefix -When your Kafka producer serializes messages with a schema-registry-aware Avro serializer (e.g. Confluent's `KafkaAvroSerializer` or the AWS Glue Avro serializer), each payload carries a short wire-format prefix in front of the Avro body: +When Confluent serializes messages with its schema-registry-aware Avro serializer (i.e. `KafkaAvroSerializer`), each payload carries a short wire-format prefix in front of the Avro body. +Said prefix is 5 bytes long, consisting of 1B magic byte (0x00) and 4B big-endian schema ID. -* **Confluent**: `1-byte magic byte (0x00) + 4-byte big-endian schema ID` — 5 bytes total. -* **AWS Glue**: `1-byte header version + 1-byte compression + 16-byte UUID schema-version ID` — 18 bytes total. +When the ESM Schema Registry integration is enabled, Lambda strips those bytes automatically and populates `value_schema_metadata.schemaId`. But when an **offline Avro schema** is used (checked into your Lambda) and do **not** use the ESM Schema Registry integration, those prefix bytes reach the function and would otherwise corrupt Avro deserialization. -When you enable the ESM Schema Registry integration, Lambda strips those bytes for you and populates `value_schema_metadata.schemaId`. But when you rely on an **offline Avro schema** (checked into your Lambda) and do **not** use the ESM Schema Registry integration, those prefix bytes reach your function and would otherwise corrupt Avro deserialization. - -Use `value_schema_id_prefix_length` (and/or `key_schema_id_prefix_length`) on `SchemaConfig` to tell Powertools how many leading bytes to skip after base64 decoding, before running the Avro decoder. +By setting the `value_schema_id_wire_format` argument on `SchemaConfig` to `"CONFLUENT"`, Powertools with strip the leading 5 bytes of the payload before running the Avro decoder. ???+ info "When do I need this?" - Only when you are supplying the Avro schema yourself **and** the producer prepended a schema-registry wire-format wrapper. If the ESM Schema Registry integration is on, leave these parameters at their default (`0`). + Only when you are supplying the Avro schema yourself **and** the producer is Confluent. If the ESM Schema Registry integration is on, leave this parameter at its default (`None`). -=== "Offline Avro schema with a Confluent-style prefix" +=== "Offline Avro schema with a Confluent prefix" ```python hl_lines="10" from aws_lambda_powertools.utilities.kafka import SchemaConfig, kafka_consumer @@ -282,7 +280,7 @@ Use `value_schema_id_prefix_length` (and/or `key_schema_id_prefix_length`) on `S schema_config = SchemaConfig( value_schema_type="AVRO", value_schema=AVRO_SCHEMA, - value_schema_id_prefix_length=5, # 1-byte magic byte + 4-byte schema ID + value_schema_wire_format="CONFLUENT" ) @@ -290,14 +288,12 @@ Use `value_schema_id_prefix_length` (and/or `key_schema_id_prefix_length`) on `S def lambda_handler(event: ConsumerRecords, context: LambdaContext): for record in event.records: # record.value is the fully-deserialized Avro payload - # with the 5-byte wire-format prefix transparently stripped. + # with the 5-byte wire-format **prefix** stripped. ... ``` -The offset is applied symmetrically for keys — set `key_schema_id_prefix_length` when a schema-registry-aware serializer was used for record keys too. - ???+ warning "Scope" - `value_schema_id_prefix_length` / `key_schema_id_prefix_length` only affect the **Avro** deserializer. Protobuf already infers the wire-format from `value_schema_metadata.schemaId` when the ESM Schema Registry integration is on, and JSON-with-registry is not supported today. + `value_schema_id_wire_format` only affects the **Avro** deserializer, just for value payloads. This implementation is easily extensible to key payloads as well if there is demand. ### Custom output serializers diff --git a/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py b/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py index de4ea2983c0..6b7b6ab41c3 100644 --- a/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py +++ b/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py @@ -79,11 +79,6 @@ def avro_encoded_value_with_prefix(avro_encoded_value): return _prepend_prefix_to_base64(avro_encoded_value) -@pytest.fixture -def avro_encoded_key_with_prefix(avro_encoded_key): - return _prepend_prefix_to_base64(avro_encoded_key) - - @pytest.fixture def kafka_event_with_avro_data(avro_encoded_value, avro_encoded_key): return { @@ -358,7 +353,7 @@ def handler(event: ConsumerRecords, context): assert result != {"name": "John Doe", "age": 30} -def test_kafka_consumer_avro_with_value_schema_id_prefix_length( +def test_kafka_consumer_avro_with_value_wire_format( kafka_event_with_avro_data, avro_encoded_value_with_prefix, avro_value_schema, @@ -372,7 +367,7 @@ def test_kafka_consumer_avro_with_value_schema_id_prefix_length( schema_config = SchemaConfig( value_schema_type="AVRO", value_schema=avro_value_schema, - value_schema_id_prefix_length=5, + value_schema_wire_format="CONFLUENT", ) @kafka_consumer(schema_config=schema_config) @@ -387,39 +382,6 @@ def handler(event: ConsumerRecords, context): assert result["age"] == 30 -def test_kafka_consumer_avro_with_key_schema_id_prefix_length( - kafka_event_with_avro_data, - avro_encoded_key_with_prefix, - avro_value_schema, - avro_key_schema, - lambda_context, -): - # GIVEN A Kafka event whose key is Avro data behind a 5-byte prefix - event = deepcopy(kafka_event_with_avro_data) - event["records"]["my-topic-1"][0]["key"] = avro_encoded_key_with_prefix - - # AND a SchemaConfig that only strips the prefix on the key side - schema_config = SchemaConfig( - value_schema_type="AVRO", - value_schema=avro_value_schema, - key_schema_type="AVRO", - key_schema=avro_key_schema, - key_schema_id_prefix_length=5, - ) - - @kafka_consumer(schema_config=schema_config) - def handler(event: ConsumerRecords, context): - record = next(event.records) - return {"key": record.key, "value": record.value} - - # WHEN The handler processes the event - result = handler(event, lambda_context) - - # THEN Both key (prefixed, offset applied) and value (plain) deserialize correctly - assert result["key"] == {"user_id": "user-123"} - assert result["value"] == {"name": "John Doe", "age": 30} - - def test_kafka_consumer_avro_with_wrong_json_schema( kafka_event_with_avro_data, lambda_context, From e10e47010cffc00c73a33665aee06ac2013b2b8c Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Sat, 12 Sep 2026 00:14:07 +0100 Subject: [PATCH 4/4] fix(kafka): validate Confluent Avro wire format --- .../utilities/kafka/consumer_records.py | 3 + .../utilities/kafka/deserializer/avro.py | 36 +++- .../kafka/deserializer/deserializer.py | 2 +- .../utilities/kafka/schema_config.py | 29 ++-- docs/utilities/kafka.md | 18 +- .../_avro/test_kafka_consumer_with_avro.py | 163 +++++++++++++++--- 6 files changed, 201 insertions(+), 50 deletions(-) diff --git a/aws_lambda_powertools/utilities/kafka/consumer_records.py b/aws_lambda_powertools/utilities/kafka/consumer_records.py index 84a526ecfec..724161828ee 100644 --- a/aws_lambda_powertools/utilities/kafka/consumer_records.py +++ b/aws_lambda_powertools/utilities/kafka/consumer_records.py @@ -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) diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/avro.py b/aws_lambda_powertools/utilities/kafka/deserializer/avro.py index 44c7ba4c644..e0ca94568e2 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/avro.py +++ b/aws_lambda_powertools/utilities/kafka/deserializer/avro.py @@ -16,6 +16,9 @@ logger = logging.getLogger(__name__) +_CONFLUENT_HEADER_SIZE = 5 +_CONFLUENT_MAGIC_BYTE = 0x00 + class AvroDeserializer(DeserializerBase): """ @@ -29,18 +32,39 @@ def __init__( self, schema_str: str, field_metadata: dict[str, Any] | None = None, - value_schema_wire_format: Literal["CONFLUENT"] | 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.value_schema_wire_format = value_schema_wire_format + 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. @@ -81,14 +105,12 @@ def deserialize(self, data: bytes | str) -> object: try: value = self._decode_input(data) - if self.value_schema_wire_format == "CONFLUENT": - # removing the first 5 bytes from payload: - # 1B magic byte 0x00 - # 4B big-endian schema ID - value = value[5:] + 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)}", diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py b/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py index 373407a6244..706e7da3d25 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py +++ b/aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py @@ -106,7 +106,7 @@ def get_deserializer( deserializer = AvroDeserializer( schema_str=schema_value, field_metadata=field_metadata, - value_schema_wire_format=wire_format, + wire_format=wire_format, ) elif schema_type == "PROTOBUF": # Import here to avoid dependency if not used diff --git a/aws_lambda_powertools/utilities/kafka/schema_config.py b/aws_lambda_powertools/utilities/kafka/schema_config.py index 96e17288953..d901f19d2e9 100644 --- a/aws_lambda_powertools/utilities/kafka/schema_config.py +++ b/aws_lambda_powertools/utilities/kafka/schema_config.py @@ -20,16 +20,20 @@ class SchemaConfig: Schema definition for message values. Required when value_schema_type is 'AVRO' or 'PROTOBUF'. value_output_serializer : Any, optional Custom output serializer for message values. Supports Pydantic classes, Dataclasses and Custom Class - value_schema_wire_format : {'CONFLUENT', None}, default=None - Set this when the payload was produced by a Confluent's schema-registry-aware serializer (KafkaAvroSerializer) - but you are supplying the Avro schema offline rather than relying on the ESM Schema Registry integration. - Only applied for AVRO values. key_schema_type : {'AVRO', 'PROTOBUF', 'JSON', None}, default=None Schema type for message keys. key_schema : str, optional 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 ------ @@ -64,15 +68,17 @@ def __init__( value_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None, value_schema: str | None = None, value_output_serializer: Any | None = None, - value_schema_wire_format: Literal["CONFLUENT"] | None = None, 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) + 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 @@ -81,6 +87,7 @@ def __init__( 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.""" @@ -89,16 +96,14 @@ def _validate_schema_requirements(self, schema_type: str | None, schema: str | N 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) -> None: - """Validate the wire format for value payload.""" + 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("Only 'CONFLUENT' wire format is supported.") + raise ValueError(f"{prefix}_schema_wire_format must be 'CONFLUENT'.") if schema_type != "AVRO": - raise ValueError("Wire format is supported for only for 'AVRO' schema.") - - return None + raise ValueError(f"{prefix}_schema_wire_format is supported only when {prefix}_schema_type is 'AVRO'.") diff --git a/docs/utilities/kafka.md b/docs/utilities/kafka.md index 9c43fc54071..5fafd65df26 100644 --- a/docs/utilities/kafka.md +++ b/docs/utilities/kafka.md @@ -258,15 +258,15 @@ Each Kafka record contains important metadata that you can access alongside the ### Using an offline Avro schema with a schema-registry wire-format prefix -When Confluent serializes messages with its schema-registry-aware Avro serializer (i.e. `KafkaAvroSerializer`), each payload carries a short wire-format prefix in front of the Avro body. -Said prefix is 5 bytes long, consisting of 1B magic byte (0x00) and 4B big-endian schema ID. +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 automatically and populates `value_schema_metadata.schemaId`. But when an **offline Avro schema** is used (checked into your Lambda) and do **not** use the ESM Schema Registry integration, those prefix bytes reach the function and would otherwise corrupt Avro deserialization. +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. -By setting the `value_schema_id_wire_format` argument on `SchemaConfig` to `"CONFLUENT"`, Powertools with strip the leading 5 bytes of the payload before running the Avro decoder. +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?" - Only when you are supplying the Avro schema yourself **and** the producer is Confluent. If the ESM Schema Registry integration is on, leave this parameter at its default (`None`). + 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" @@ -280,20 +280,20 @@ By setting the `value_schema_id_wire_format` argument on `SchemaConfig` to `"CON schema_config = SchemaConfig( value_schema_type="AVRO", value_schema=AVRO_SCHEMA, - value_schema_wire_format="CONFLUENT" + 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 fully-deserialized Avro payload - # with the 5-byte wire-format **prefix** stripped. + # record.value is the deserialized Avro payload + # with the validated 5-byte wire-format header removed. ... ``` ???+ warning "Scope" - `value_schema_id_wire_format` only affects the **Avro** deserializer, just for value payloads. This implementation is easily extensible to key payloads as well if there is demand. + `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 diff --git a/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py b/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py index 6b7b6ab41c3..bdf839e31ce 100644 --- a/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py +++ b/tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py @@ -8,6 +8,8 @@ from avro.schema import parse as parse_schema from aws_lambda_powertools.utilities.kafka.consumer_records import ConsumerRecords +from aws_lambda_powertools.utilities.kafka.deserializer import deserializer as deserializer_factory +from aws_lambda_powertools.utilities.kafka.deserializer.avro import AvroDeserializer from aws_lambda_powertools.utilities.kafka.exceptions import ( KafkaConsumerAvroSchemaParserError, KafkaConsumerDeserializationError, @@ -79,6 +81,11 @@ def avro_encoded_value_with_prefix(avro_encoded_value): return _prepend_prefix_to_base64(avro_encoded_value) +@pytest.fixture +def avro_encoded_key_with_prefix(avro_encoded_key): + return _prepend_prefix_to_base64(avro_encoded_key) + + @pytest.fixture def kafka_event_with_avro_data(avro_encoded_value, avro_encoded_key): return { @@ -324,46 +331,76 @@ def test_kafka_consumer_without_avro_key_schema(): assert "key_schema" in str(excinfo.value) -def test_kafka_consumer_avro_produces_wrong_output_without_prefix_length_setting( +def test_kafka_consumer_avro_with_value_wire_format( kafka_event_with_avro_data, avro_encoded_value_with_prefix, avro_value_schema, lambda_context, ): - # GIVEN An Avro payload that has been serialized by a Confluent-style producer, - # so it carries a 5-byte "magic byte + schema ID" prefix in front of the Avro body + # GIVEN An Avro payload with a 5-byte magic-byte + schema-ID prefix event = deepcopy(kafka_event_with_avro_data) event["records"]["my-topic-1"][0]["value"] = avro_encoded_value_with_prefix - # AND a SchemaConfig without the new offset parameter (today's behaviour) - schema_config = SchemaConfig(value_schema_type="AVRO", value_schema=avro_value_schema) + # AND a SchemaConfig instructed to validate and remove the Confluent header + schema_config = SchemaConfig( + value_schema_type="AVRO", + value_schema=avro_value_schema, + value_schema_wire_format="CONFLUENT", + ) @kafka_consumer(schema_config=schema_config) def handler(event: ConsumerRecords, context): return event.record.value - # WHEN/THEN The deserializer cannot know it should skip the leading bytes. - # Depending on the prefix content, this either raises or silently returns corrupted data. - # Both outcomes are broken; the fix must let callers opt into skipping the prefix. - try: - result = handler(event, lambda_context) - except KafkaConsumerDeserializationError: - return + # WHEN The handler processes the event + result = handler(event, lambda_context) - assert result != {"name": "John Doe", "age": 30} + # THEN The Avro body should be decoded correctly after the prefix is stripped + assert result["name"] == "John Doe" + assert result["age"] == 30 -def test_kafka_consumer_avro_with_value_wire_format( +def test_kafka_consumer_avro_with_key_and_value_wire_format( kafka_event_with_avro_data, + avro_encoded_key_with_prefix, avro_encoded_value_with_prefix, + avro_key_schema, avro_value_schema, lambda_context, ): - # GIVEN An Avro payload with a 5-byte magic-byte + schema-ID prefix + # GIVEN Confluent-framed Avro key and value payloads event = deepcopy(kafka_event_with_avro_data) + event["records"]["my-topic-1"][0]["key"] = avro_encoded_key_with_prefix event["records"]["my-topic-1"][0]["value"] = avro_encoded_value_with_prefix + schema_config = SchemaConfig( + value_schema_type="AVRO", + value_schema=avro_value_schema, + key_schema_type="AVRO", + key_schema=avro_key_schema, + value_schema_wire_format="CONFLUENT", + key_schema_wire_format="CONFLUENT", + ) + + @kafka_consumer(schema_config=schema_config) + def handler(event: ConsumerRecords, context): + record = event.record + return record.key, record.value + + # WHEN the handler processes both payloads + key, value = handler(event, lambda_context) - # AND a SchemaConfig instructed to skip the first 5 bytes before Avro decoding + # THEN both headers are removed before Avro deserialization + assert key == {"user_id": "user-123"} + assert value == {"name": "John Doe", "age": 30} + + +def test_kafka_consumer_rejects_short_confluent_header( + kafka_event_with_avro_data, + avro_value_schema, + lambda_context, +): + event = deepcopy(kafka_event_with_avro_data) + event["records"]["my-topic-1"][0]["value"] = base64.b64encode(b"\x00\x00\x00\x00").decode("utf-8") schema_config = SchemaConfig( value_schema_type="AVRO", value_schema=avro_value_schema, @@ -374,12 +411,96 @@ def test_kafka_consumer_avro_with_value_wire_format( def handler(event: ConsumerRecords, context): return event.record.value - # WHEN The handler processes the event - result = handler(event, lambda_context) + with pytest.raises(KafkaConsumerDeserializationError, match="payload must contain a 5-byte header"): + handler(event, lambda_context) - # THEN The Avro body should be decoded correctly after the prefix is stripped - assert result["name"] == "John Doe" - assert result["age"] == 30 + +def test_kafka_consumer_rejects_invalid_confluent_magic_byte( + kafka_event_with_avro_data, + avro_encoded_value, + avro_value_schema, + lambda_context, +): + event = deepcopy(kafka_event_with_avro_data) + invalid_prefix = b"\x01\x00\x00\x00\x01" + event["records"]["my-topic-1"][0]["value"] = _prepend_prefix_to_base64( + avro_encoded_value, + prefix=invalid_prefix, + ) + schema_config = SchemaConfig( + value_schema_type="AVRO", + value_schema=avro_value_schema, + value_schema_wire_format="CONFLUENT", + ) + + @kafka_consumer(schema_config=schema_config) + def handler(event: ConsumerRecords, context): + return event.record.value + + with pytest.raises(KafkaConsumerDeserializationError, match="expected magic byte 0x00"): + handler(event, lambda_context) + + +def test_schema_config_preserves_existing_positional_arguments(avro_value_schema, avro_key_schema): + config = SchemaConfig("AVRO", avro_value_schema, None, "AVRO", avro_key_schema, None) + + assert config.value_schema_type == "AVRO" + assert config.value_schema == avro_value_schema + assert config.key_schema_type == "AVRO" + assert config.key_schema == avro_key_schema + assert config.value_schema_wire_format is None + assert config.key_schema_wire_format is None + + +@pytest.mark.parametrize("prefix", ["value", "key"]) +def test_schema_config_rejects_wire_format_for_non_avro_schema(prefix): + kwargs = { + f"{prefix}_schema_type": "JSON", + f"{prefix}_schema_wire_format": "CONFLUENT", + } + + with pytest.raises(ValueError, match=rf"{prefix}_schema_wire_format is supported only"): + SchemaConfig(**kwargs) + + +def test_schema_config_rejects_unknown_wire_format(avro_value_schema): + with pytest.raises(ValueError, match="value_schema_wire_format must be 'CONFLUENT'"): + SchemaConfig( + value_schema_type="AVRO", + value_schema=avro_value_schema, + value_schema_wire_format="GLUE", # type: ignore[arg-type] + ) + + +def test_avro_deserializer_rejects_unknown_wire_format(avro_value_schema, avro_encoded_value): + deserializer = AvroDeserializer( + avro_value_schema, + wire_format="GLUE", # type: ignore[arg-type] + ) + + with pytest.raises(KafkaConsumerDeserializationError, match="Unsupported Avro wire format: GLUE"): + deserializer.deserialize(avro_encoded_value) + + +def test_avro_deserializer_cache_includes_wire_format(monkeypatch, avro_value_schema): + monkeypatch.setattr(deserializer_factory, "_deserializer_cache", {}) + + plain = deserializer_factory.get_deserializer("AVRO", avro_value_schema, {}) + confluent = deserializer_factory.get_deserializer( + "AVRO", + avro_value_schema, + {}, + wire_format="CONFLUENT", + ) + + assert plain is not confluent + assert plain is deserializer_factory.get_deserializer("AVRO", avro_value_schema, {}) + assert confluent is deserializer_factory.get_deserializer( + "AVRO", + avro_value_schema, + {}, + wire_format="CONFLUENT", + ) def test_kafka_consumer_avro_with_wrong_json_schema(