Avro schema references support (#1088)

slominskir · anchitj · web-flow · commit 26d40bed980b · 2023-04-05T17:41:41.000+05:30
* Initial attempt at schema references support: #974 * flake8 line length * Update tests to use Schema * flake8 whitespace after comma * Imports should not be from src * Updating test cases to use Schema object * Convert Registry Schema to fastavro Schema (dict) * Updated docs to indicate Schema object, not Schema String arg * fastavro 1.4.0 renamed _named_schemas * primitive types must be JSON formatted * Fixes #974 * Fixed table formatting * flake8 fixes * Removed named_schemas from schema_registry_client * Added support for nested references * PR Feedback * PR Feedback * Remove unneeded changes * Update unit tests * PR Feedback * PR Feedback * PR Feedback * Fix unit test * PR Feedback * PR Feedback --------- Co-authored-by: Anchit Jain <anjain@confluent.io>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,7 +11,7 @@ v2.1.0 is a feature release with the following features, fixes and enhancements:
 - Added support for password protected private key in CachedSchemaRegistryClient.
 - Add reference support in Schema Registry client. (@RickTalken, #1304)
 - Migrated travis jobs to Semaphore CI (#1503)
-- Add support for passing schema references in JSONSerializer and JSONDeserializer. (#1514)
+- Added support for schema references. (#1514 and @slominskir #1088)
 
 confluent-kafka-python is based on librdkafka v2.1.0, see the
 [librdkafka release notes](https://github.com/edenhill/librdkafka/releases/tag/v2.1.0)
diff --git a/src/confluent_kafka/schema_registry/avro.py b/src/confluent_kafka/schema_registry/avro.py
@@ -67,6 +67,24 @@ def _schema_loads(schema_str):
     return Schema(schema_str, schema_type='AVRO')
 
 
+def _resolve_named_schema(schema, schema_registry_client, named_schemas=None):
+    """
+    Resolves named schemas referenced by the provided schema recursively.
+    :param schema: Schema to resolve named schemas for.
+    :param schema_registry_client: SchemaRegistryClient to use for retrieval.
+    :param named_schemas: Dict of named schemas resolved recursively.
+    :return: named_schemas dict.
+    """
+    if named_schemas is None:
+        named_schemas = {}
+    if schema.references is not None:
+        for ref in schema.references:
+            referenced_schema = schema_registry_client.get_version(ref.subject, ref.version)
+            _resolve_named_schema(referenced_schema.schema, schema_registry_client, named_schemas)
+            parse_schema(loads(referenced_schema.schema.schema_str), named_schemas=named_schemas)
+    return named_schemas
+
+
 class AvroSerializer(Serializer):
     """
     Serializer that outputs Avro binary encoded data with Confluent Schema Registry framing.
@@ -146,7 +164,7 @@ class AvroSerializer(Serializer):
     Args:
         schema_registry_client (SchemaRegistryClient): Schema Registry client instance.
 
-        schema_str (str): Avro `Schema Declaration. <https://avro.apache.org/docs/current/spec.html#schemas>`_
+        schema_str (str or Schema): Avro `Schema Declaration. <https://avro.apache.org/docs/current/spec.html#schemas>`_ Accepts either a string or a `Schema`(Schema) instance.  Note that string definitions cannot reference other schemas. For referencing other schemas, use a Schema instance.
 
         to_dict (callable, optional): Callable(object, SerializationContext) -> dict. Converts object to a dict.
 
@@ -155,15 +173,21 @@ class AvroSerializer(Serializer):
     __slots__ = ['_hash', '_auto_register', '_normalize_schemas', '_use_latest_version',
                  '_known_subjects', '_parsed_schema',
                  '_registry', '_schema', '_schema_id', '_schema_name',
-                 '_subject_name_func', '_to_dict']
+                 '_subject_name_func', '_to_dict', '_named_schemas']
 
     _default_conf = {'auto.register.schemas': True,
                      'normalize.schemas': False,
                      'use.latest.version': False,
                      'subject.name.strategy': topic_subject_name_strategy}
 
-    def __init__(self, schema_registry_client, schema_str,
-                 to_dict=None, conf=None):
+    def __init__(self, schema_registry_client, schema_str, to_dict=None, conf=None):
+        if isinstance(schema_str, str):
+            schema = _schema_loads(schema_str)
+        elif isinstance(schema_str, Schema):
+            schema = schema_str
+        else:
+            raise TypeError('You must pass either schema string or schema object')
+
         self._registry = schema_registry_client
         self._schema_id = None
         self._known_subjects = set()
@@ -200,9 +224,9 @@ def __init__(self, schema_registry_client, schema_str,
             raise ValueError("Unrecognized properties: {}"
                              .format(", ".join(conf_copy.keys())))
 
-        schema = _schema_loads(schema_str)
         schema_dict = loads(schema.schema_str)
-        parsed_schema = parse_schema(schema_dict)
+        self._named_schemas = _resolve_named_schema(schema, schema_registry_client)
+        parsed_schema = parse_schema(schema_dict, named_schemas=self._named_schemas)
 
         if isinstance(parsed_schema, list):
             # if parsed_schema is a list, we have an Avro union and there
@@ -299,8 +323,9 @@ class AvroDeserializer(Deserializer):
         schema_registry_client (SchemaRegistryClient): Confluent Schema Registry
             client instance.
 
-        schema_str (str, optional): The reader schema.
-            If not provided, the writer schema will be used as the reader schema.
+        schema_str (str, Schema, optional): Avro reader schema declaration Accepts either a string or a `Schema`(
+        Schema) instance. If not provided, the writer schema will be used as the reader schema. Note that string
+        definitions cannot reference other schemas. For referencing other schemas, use a Schema instance.
 
         from_dict (callable, optional): Callable(dict, SerializationContext) -> object.
             Converts a dict to an instance of some object.
@@ -315,13 +340,31 @@ class AvroDeserializer(Deserializer):
         `Apache Avro Schema Resolution <https://avro.apache.org/docs/1.8.2/spec.html#Schema+Resolution>`_
     """
 
-    __slots__ = ['_reader_schema', '_registry', '_from_dict', '_writer_schemas', '_return_record_name']
+    __slots__ = ['_reader_schema', '_registry', '_from_dict', '_writer_schemas', '_return_record_name', '_schema',
+                 '_named_schemas']
 
     def __init__(self, schema_registry_client, schema_str=None, from_dict=None, return_record_name=False):
+        schema = None
+        if schema_str is not None:
+            if isinstance(schema_str, str):
+                schema = _schema_loads(schema_str)
+            elif isinstance(schema_str, Schema):
+                schema = schema_str
+            else:
+                raise TypeError('You must pass either schema string or schema object')
+
+        self._schema = schema
         self._registry = schema_registry_client
         self._writer_schemas = {}
 
-        self._reader_schema = parse_schema(loads(schema_str)) if schema_str else None
+        if schema:
+            schema_dict = loads(self._schema.schema_str)
+            self._named_schemas = _resolve_named_schema(self._schema, schema_registry_client)
+            self._reader_schema = parse_schema(schema_dict,
+                                               named_schemas=self._named_schemas)
+        else:
+            self._named_schemas = None
+            self._reader_schema = None
 
         if from_dict is not None and not callable(from_dict):
             raise ValueError("from_dict must be callable with the signature "
@@ -370,10 +413,11 @@ def __call__(self, data, ctx):
             writer_schema = self._writer_schemas.get(schema_id, None)
 
             if writer_schema is None:
-                schema = self._registry.get_schema(schema_id)
-                prepared_schema = _schema_loads(schema.schema_str)
+                registered_schema = self._registry.get_schema(schema_id)
+                self._named_schemas = _resolve_named_schema(registered_schema, self._registry)
+                prepared_schema = _schema_loads(registered_schema.schema_str)
                 writer_schema = parse_schema(loads(
-                    prepared_schema.schema_str))
+                    prepared_schema.schema_str), named_schemas=self._named_schemas)
                 self._writer_schemas[schema_id] = writer_schema
 
             obj_dict = schemaless_reader(payload,
diff --git a/src/confluent_kafka/schema_registry/schema_registry_client.py b/src/confluent_kafka/schema_registry/schema_registry_client.py
@@ -686,12 +686,11 @@ class Schema(object):
     Args:
         schema_str (str): String representation of the schema.
 
-        references ([SchemaReference]): SchemaReferences used in this schema.
-
         schema_type (str): The schema type: AVRO, PROTOBUF or JSON.
-    """
 
-    __slots__ = ['schema_str', 'references', 'schema_type', '_hash']
+        references ([SchemaReference]): SchemaReferences used in this schema.
+    """
+    __slots__ = ['schema_str', 'schema_type', 'references', '_hash']
 
     def __init__(self, schema_str, schema_type, references=[]):
         super(Schema, self).__init__()
diff --git a/tests/integration/schema_registry/test_avro_serializers.py b/tests/integration/schema_registry/test_avro_serializers.py
@@ -15,14 +15,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-
 import pytest
 
 from confluent_kafka import TopicPartition
 from confluent_kafka.serialization import (MessageField,
                                            SerializationContext)
 from confluent_kafka.schema_registry.avro import (AvroSerializer,
                                                   AvroDeserializer)
+from confluent_kafka.schema_registry import Schema, SchemaReference
 
 
 class User(object):
@@ -51,6 +51,145 @@ def __eq__(self, other):
             self.favorite_color == other.favorite_color])
 
 
+class AwardProperties(object):
+    schema_str = """
+        {
+            "namespace": "confluent.io.examples.serialization.avro",
+            "name": "AwardProperties",
+            "type": "record",
+            "fields": [
+                {"name": "year", "type": "int"},
+                {"name": "points", "type": "int"}
+            ]
+        }
+    """
+
+    def __init__(self, points, year):
+        self.points = points
+        self.year = year
+
+    def __eq__(self, other):
+        return all([
+            self.points == other.points,
+            self.year == other.year
+        ])
+
+
+class Award(object):
+    schema_str = """
+        {
+            "namespace": "confluent.io.examples.serialization.avro",
+            "name": "Award",
+            "type": "record",
+            "fields": [
+                {"name": "name", "type": "string"},
+                {"name": "properties", "type": "AwardProperties"}
+            ]
+        }
+    """
+
+    def __init__(self, name, properties):
+        self.name = name
+        self.properties = properties
+
+    def __eq__(self, other):
+        return all([
+            self.name == other.name,
+            self.properties == other.properties
+        ])
+
+
+class AwardedUser(object):
+    schema_str = """
+        {
+            "namespace": "confluent.io.examples.serialization.avro",
+            "name": "AwardedUser",
+            "type": "record",
+            "fields": [
+                {"name": "award", "type": "Award"},
+                {"name": "user", "type": "User"}
+            ]
+        }
+    """
+
+    def __init__(self, award, user):
+        self.award = award
+        self.user = user
+
+    def __eq__(self, other):
+        return all([
+            self.award == other.award,
+            self.user == other.user
+        ])
+
+
+def _register_avro_schemas_and_build_awarded_user_schema(kafka_cluster):
+    sr = kafka_cluster.schema_registry()
+
+    user = User('Bowie', 47, 'purple')
+    award_properties = AwardProperties(10, 2023)
+    award = Award("Best In Show", award_properties)
+    awarded_user = AwardedUser(award, user)
+
+    user_schema_ref = SchemaReference("confluent.io.examples.serialization.avro.User", "user", 1)
+    award_properties_schema_ref = SchemaReference("confluent.io.examples.serialization.avro.AwardProperties",
+                                                  "award_properties", 1)
+    award_schema_ref = SchemaReference("confluent.io.examples.serialization.avro.Award", "award", 1)
+
+    sr.register_schema("user", Schema(User.schema_str, 'AVRO'))
+    sr.register_schema("award_properties", Schema(AwardProperties.schema_str, 'AVRO'))
+    sr.register_schema("award", Schema(Award.schema_str, 'AVRO', [award_properties_schema_ref]))
+
+    references = [user_schema_ref, award_schema_ref]
+    schema = Schema(AwardedUser.schema_str, 'AVRO', references)
+    return awarded_user, schema
+
+
+def _references_test_common(kafka_cluster, awarded_user, serializer_schema, deserializer_schema):
+    """
+    Common (both reader and writer) avro schema reference test.
+    Args:
+        kafka_cluster (KafkaClusterFixture): cluster fixture
+    """
+    topic = kafka_cluster.create_topic("reference-avro")
+    sr = kafka_cluster.schema_registry()
+
+    value_serializer = AvroSerializer(sr, serializer_schema,
+                                      lambda user, ctx:
+                                      dict(award=dict(name=user.award.name,
+                                                      properties=dict(year=user.award.properties.year,
+                                                                      points=user.award.properties.points)),
+                                           user=dict(name=user.user.name,
+                                                     favorite_number=user.user.favorite_number,
+                                                     favorite_color=user.user.favorite_color)))
+
+    value_deserializer = \
+        AvroDeserializer(sr, deserializer_schema,
+                         lambda user, ctx:
+                         AwardedUser(award=Award(name=user.get('award').get('name'),
+                                                 properties=AwardProperties(
+                                                     year=user.get('award').get('properties').get(
+                                                         'year'),
+                                                     points=user.get('award').get('properties').get(
+                                                         'points'))),
+                                     user=User(name=user.get('user').get('name'),
+                                               favorite_number=user.get('user').get('favorite_number'),
+                                               favorite_color=user.get('user').get('favorite_color'))))
+
+    producer = kafka_cluster.producer(value_serializer=value_serializer)
+
+    producer.produce(topic, value=awarded_user, partition=0)
+    producer.flush()
+
+    consumer = kafka_cluster.consumer(value_deserializer=value_deserializer)
+    consumer.assign([TopicPartition(topic, 0)])
+
+    msg = consumer.poll()
+    awarded_user2 = msg.value()
+
+    assert awarded_user2 == awarded_user
+
+
 @pytest.mark.parametrize("avsc, data, record_type",
                          [('basic_schema.avsc', {'name': 'abc'}, "record"),
                           ('primitive_string.avsc', u'Jämtland', "string"),
@@ -185,3 +324,25 @@ def test_avro_record_serialization_custom(kafka_cluster):
     user2 = msg.value()
 
     assert user2 == user
+
+
+def test_avro_reference(kafka_cluster):
+    """
+    Tests Avro schema reference with both serializer and deserializer schemas provided.
+    Args:
+        kafka_cluster (KafkaClusterFixture): cluster fixture
+    """
+    awarded_user, schema = _register_avro_schemas_and_build_awarded_user_schema(kafka_cluster)
+
+    _references_test_common(kafka_cluster, awarded_user, schema, schema)
+
+
+def test_avro_reference_deserializer_none(kafka_cluster):
+    """
+    Tests Avro schema reference with serializer schema provided and deserializer schema set to None.
+    Args:
+        kafka_cluster (KafkaClusterFixture): cluster fixture
+    """
+    awarded_user, schema = _register_avro_schemas_and_build_awarded_user_schema(kafka_cluster)
+
+    _references_test_common(kafka_cluster, awarded_user, schema, None)
diff --git a/tests/schema_registry/test_avro_serializer.py b/tests/schema_registry/test_avro_serializer.py
