Handle empty context properly for Avro deserializer (#2035)

* update

* lint

* address feedback

Author: fangnx

src/confluent_kafka/schema_registry/__init__.py

@@ -36,7 +36,7 @@
   SchemaReference,
   ServerConfig
 )
-from ..serialization import SerializationError, MessageField
+from ..serialization import SerializationError, MessageField, SerializationContext
 
 _KEY_SCHEMA_ID = "__key_schema_id"
 _VALUE_SCHEMA_ID = "__value_schema_id"
@@ -128,7 +128,7 @@ def reference_subject_name_strategy(ctx, schema_ref: SchemaReference) -> Optiona
     return schema_ref.name if schema_ref is not None else None
 
 
-def header_schema_id_serializer(payload: bytes, ctx, schema_id) -> bytes:
+def header_schema_id_serializer(payload: bytes, ctx: Optional[SerializationContext], schema_id) -> bytes:
     """
     Serializes the schema guid into the header.
 
@@ -141,6 +141,9 @@ def header_schema_id_serializer(payload: bytes, ctx, schema_id) -> bytes:
     Returns:
         bytes: The payload
     """
+    if ctx is None:
+        raise SerializationError("SerializationContext is required for header_schema_id_serializer")
+
     headers = ctx.headers
     if headers is None:
         raise SerializationError("Missing headers")
@@ -171,7 +174,7 @@ def prefix_schema_id_serializer(payload: bytes, ctx, schema_id) -> bytes:
     return schema_id.id_to_bytes() + payload
 
 
-def dual_schema_id_deserializer(payload: bytes, ctx, schema_id) -> io.BytesIO:
+def dual_schema_id_deserializer(payload: bytes, ctx: Optional[SerializationContext], schema_id) -> io.BytesIO:
     """
     Deserializes the schema id by first checking the header, then the payload prefix.
 
@@ -184,22 +187,28 @@ def dual_schema_id_deserializer(payload: bytes, ctx, schema_id) -> io.BytesIO:
     Returns:
         bytes: The payload
     """
-    headers = ctx.headers
-    header_key = _KEY_SCHEMA_ID if ctx.field == MessageField.KEY else _VALUE_SCHEMA_ID
-    if headers is not None:
-        header_value = None
-        if isinstance(headers, list):
-            # look for header_key in headers
-            for header in headers:
-                if header[0] == header_key:
-                    header_value = header[1]
-                    break
-        elif isinstance(headers, dict):
-            header_value = headers.get(header_key, None)
-        if header_value is not None:
-            schema_id.from_bytes(io.BytesIO(header_value))
-            return io.BytesIO(payload)
-    return schema_id.from_bytes(io.BytesIO(payload))
+    # Look for schema ID in headers
+    header_value = None
+
+    if ctx is not None:
+        headers = ctx.headers
+        if headers is not None:
+            header_key = _KEY_SCHEMA_ID if ctx.field == MessageField.KEY else _VALUE_SCHEMA_ID
+            if isinstance(headers, list):
+                # look for header_key in headers
+                for header in headers:
+                    if header[0] == header_key:
+                        header_value = header[1]
+                        break
+            elif isinstance(headers, dict):
+                header_value = headers.get(header_key, None)
+
+    # Parse schema ID from determined source and return appropriate payload
+    if header_value is not None:
+        schema_id.from_bytes(io.BytesIO(header_value))
+        return io.BytesIO(payload)  # Return full payload when schema ID is in header
+    else:
+        return schema_id.from_bytes(io.BytesIO(payload))  # Parse from payload, return remainder
 
 
 def prefix_schema_id_deserializer(payload: bytes, ctx, schema_id) -> io.BytesIO:

src/confluent_kafka/schema_registry/_async/schema_registry_client.py

@@ -763,8 +763,9 @@ async def get_schema_by_guid(
         if schema is not None:
             return schema
 
+        query = {}
         if fmt is not None:
-            query = {'format': fmt}
+            query['format'] = fmt
         response = await self._rest_client.get('schemas/guids/{}'.format(guid), query)
 
         registered_schema = RegisteredSchema.from_dict(response)

src/confluent_kafka/schema_registry/_sync/schema_registry_client.py

@@ -763,8 +763,9 @@ def get_schema_by_guid(
         if schema is not None:
             return schema
 
+        query = {}
         if fmt is not None:
-            query = {'format': fmt}
+            query['format'] = fmt
         response = self._rest_client.get('schemas/guids/{}'.format(guid), query)
 
         registered_schema = RegisteredSchema.from_dict(response)

tests/schema_registry/test_schema_id.py

@@ -16,7 +16,14 @@
 # limitations under the License.
 #
 import io
+import pytest
+
 from confluent_kafka.schema_registry.serde import SchemaId
+from confluent_kafka.schema_registry import (
+    dual_schema_id_deserializer,
+    header_schema_id_serializer,
+    SerializationError
+)
 
 
 def test_schema_guid():
@@ -71,3 +78,26 @@ def test_schema_id_with_message_indexes():
     assert indexes == [1, 2, 3]
     output = schema_id.id_to_bytes()
     assert output == input
+
+
+def test_dual_schema_id_deserializer_handles_none_context():
+    """
+    Ensures dual_schema_id_deserializer handles None SerializationContext properly.
+    """
+    schema_id = SchemaId("AVRO")
+    test_data = b'\x00\x00\x00\x00\x01'  # Valid schema ID format
+
+    result = dual_schema_id_deserializer(test_data, ctx=None, schema_id=schema_id)
+
+    # Verify it returns BytesIO and parsed the schema ID
+    assert isinstance(result, io.BytesIO)
+    assert schema_id.id == 1
+
+
+def test_header_schema_id_serializer_handles_none_context():
+    """
+    Ensures header_schema_id_serializer handles None SerializationContext properly.
+    """
+    # schema_id won't be used since function raises error early when ctx=None
+    with pytest.raises(SerializationError, match="SerializationContext is required"):
+        header_schema_id_serializer(b"test_payload", ctx=None, schema_id=None)