Fix handling of named Avro schemas (#1928)

Author: rayokota

src/confluent_kafka/schema_registry/avro.py

@@ -18,15 +18,17 @@
 import decimal
 import re
 from collections import defaultdict
+from copy import deepcopy
 from io import BytesIO
 from json import loads
 from struct import pack, unpack
 from typing import Dict, Union, Optional, Set, Callable
 
-from fastavro import (parse_schema,
-                      schemaless_reader,
+from fastavro import (schemaless_reader,
                       schemaless_writer,
+                      repository,
                       validate)
+from fastavro.schema import load_schema
 
 from . import (_MAGIC_BYTE,
                Schema,
@@ -104,7 +106,8 @@ def _resolve_named_schema(
         for ref in schema.references:
             referenced_schema = schema_registry_client.get_version(ref.subject, ref.version, True)
             ref_named_schemas = _resolve_named_schema(referenced_schema.schema, schema_registry_client)
-            parsed_schema = parse_schema(loads(referenced_schema.schema.schema_str), named_schemas=ref_named_schemas)
+            parsed_schema = parse_schema_with_repo(
+                referenced_schema.schema.schema_str, named_schemas=ref_named_schemas)
             named_schemas.update(ref_named_schemas)
             named_schemas[ref.name] = parsed_schema
     return named_schemas
@@ -378,8 +381,8 @@ def _get_parsed_schema(self, schema: Schema) -> AvroSchema:
 
         named_schemas = _resolve_named_schema(schema, self._registry)
         prepared_schema = _schema_loads(schema.schema_str)
-        parsed_schema = parse_schema(
-            loads(prepared_schema.schema_str), named_schemas=named_schemas, expand=True)
+        parsed_schema = parse_schema_with_repo(
+            prepared_schema.schema_str, named_schemas=named_schemas)
 
         self._parsed_schemas.set(schema, parsed_schema)
         return parsed_schema
@@ -606,13 +609,28 @@ def _get_parsed_schema(self, schema: Schema) -> AvroSchema:
 
         named_schemas = _resolve_named_schema(schema, self._registry)
         prepared_schema = _schema_loads(schema.schema_str)
-        parsed_schema = parse_schema(
-            loads(prepared_schema.schema_str), named_schemas=named_schemas, expand=True)
+        parsed_schema = parse_schema_with_repo(
+            prepared_schema.schema_str, named_schemas=named_schemas)
 
         self._parsed_schemas.set(schema, parsed_schema)
         return parsed_schema
 
 
+class LocalSchemaRepository(repository.AbstractSchemaRepository):
+    def __init__(self, schemas):
+        self.schemas = schemas
+
+    def load(self, subject):
+        return self.schemas.get(subject)
+
+
+def parse_schema_with_repo(schema_str: str, named_schemas: Dict[str, AvroSchema]) -> AvroSchema:
+    copy = deepcopy(named_schemas)
+    copy["$root"] = loads(schema_str)
+    repo = LocalSchemaRepository(copy)
+    return load_schema("$root", repo=repo)
+
+
 def transform(
     ctx: RuleContext, schema: AvroSchema, message: AvroMessage,
     field_transform: FieldTransform

tests/schema_registry/test_avro_serdes.py

@@ -216,6 +216,37 @@ def test_avro_serialize_references():
     assert obj == obj2
 
 
+def test_avro_serialize_union():
+    conf = {'url': _BASE_URL}
+    client = SchemaRegistryClient.new_client(conf)
+    ser_conf = {'auto.register.schemas': False, 'use.latest.version': True}
+
+    obj = {
+        'First': {'stringField': 'hi'},
+        'Second': {'stringField': 'hi'},
+    }
+    schema = ['null', {
+        'type': 'record',
+        'name': 'A',
+        'namespace': 'test',
+        'fields': [
+            {'name': 'First', 'type': {'type': 'record', 'name': 'B', 'fields': [
+                {'name': 'stringField', 'type': 'string'},
+            ]}},
+            {'name': 'Second', 'type': 'B'}
+        ]
+    }]
+    client.register_schema(_SUBJECT, Schema(json.dumps(schema), 'AVRO'))
+
+    ser = AvroSerializer(client, schema_str=None, conf=ser_conf)
+    ser_ctx = SerializationContext(_TOPIC, MessageField.VALUE)
+    obj_bytes = ser(obj, ser_ctx)
+
+    deser = AvroDeserializer(client)
+    obj2 = deser(obj_bytes, ser_ctx)
+    assert obj == obj2
+
+
 def test_avro_serialize_union_with_references():
     conf = {'url': _BASE_URL}
     client = SchemaRegistryClient.new_client(conf)