Deserialize to DataAndSchema instead of just data

Author: cdebost

README.md

@@ -66,7 +66,7 @@ from kafka import KafkaConsumer
 # Create the schema registry client, which is a façade around the boto3 glue client
 client = SchemaRegistryClient(glue_client,
                               registry_name='my-registry')
-                              
+
 # Create the serializer
 serializer = SchemaRegistrySerializer(client)
 
@@ -77,7 +77,7 @@ producer = KafkaProducer(value_serializer=serializer)
 # In this example we're using Avro, so we'll load an .avsc file.
 with open('user.avsc', 'r') as schema_file:
     schema = AvroSchema(schema_file.read())
-    
+
 # Send message data along with schema
 data = {
     'name': 'John Doe',
@@ -101,7 +101,7 @@ from kafka import KafkaConsumer
 # Create the schema registry client, which is a façade around the boto3 glue client
 client = SchemaRegistryClient(glue_client,
                               registry_name='my-registry')
-                              
+
 # Create the deserializer
 deserializer = SchemaRegistryDeserializer(client)
 
@@ -110,7 +110,10 @@ consumer = KafkaConsumer('my-topic', value_deserializer=deserializer)
 
 # Now use the consumer normally
 for message in consumer:
-    # `message.value` will be the deserialized value
+    # The deserializer produces DataAndSchema instances
+    value: DataAndSchema = message.value
+    value.data
+    value.schema
 ```
 
 ## Contributing

src/aws_schema_registry/schema.py

@@ -64,3 +64,6 @@ class SchemaVersion:
     data_format: DataFormat
     status: SchemaVersionStatus
     version_number: Optional[int] = None
+
+    def __hash__(self):
+        return hash(self.definition)

src/aws_schema_registry/serde.py

@@ -3,7 +3,7 @@
 import functools
 import logging
 import sys
-from typing import Any, Dict, NamedTuple
+from typing import Any, NamedTuple
 from uuid import UUID
 
 if sys.version_info[1] < 8:  # for py37
@@ -109,8 +109,6 @@ class SchemaRegistryDeserializer:
             registries or handle schema-less data.
     """
 
-    _writer_schemas: Dict[UUID, Schema]
-
     def __init__(
         self,
         client: SchemaRegistryClient,
@@ -120,7 +118,6 @@ def __init__(
         self.client = client
         self.return_record_name = return_record_name
         self.secondary_deserializer = secondary_deserializer
-        self._writer_schemas = {}
 
     def deserialize(self, topic: str, bytes_: bytes):
         if bytes_ is None:
@@ -135,20 +132,18 @@ def deserialize(self, topic: str, bytes_: bytes):
                     'no secondary deserializer provided to handle'
                     ' unrecognized data encoding'
                 ) from e
-        writer_schema = self._writer_schemas.get(schema_version_id)
-        if not writer_schema:
-            LOG.info('Schema version %s not found locally, fetching from'
-                     ' registry...', schema_version_id)
-            schema_version = self.client.get_schema_version(
-                version_id=schema_version_id
-            )
-            writer_schema = self._create_writer_schema(schema_version)
-            self._writer_schemas[schema_version_id] = writer_schema
-            LOG.info('Schema version %s fetched', schema_version_id)
-        return writer_schema.read(data_bytes)
-
-    def _create_writer_schema(self, schema_version: SchemaVersion) -> Schema:
-        if schema_version.data_format == 'AVRO':
-            return AvroSchema(schema_version.definition)
-        elif schema_version.data_format == 'JSON':
+        writer_schema_version = self._get_schema_version(schema_version_id)
+        writer_schema = self._schema_for_version(writer_schema_version)
+        return DataAndSchema(writer_schema.read(data_bytes), writer_schema)
+
+    @functools.lru_cache(maxsize=None)
+    def _get_schema_version(self, version_id: UUID):
+        LOG.info('Fetching schema version %s...', version_id)
+        return self.client.get_schema_version(version_id)
+
+    @functools.lru_cache(maxsize=None)
+    def _schema_for_version(self, version: SchemaVersion) -> Schema:
+        if version.data_format == 'AVRO':
+            return AvroSchema(version.definition)
+        elif version.data_format == 'JSON':
             raise NotImplementedError('JSON schema not supported')

tests/integration/java/test_java_integration.py

@@ -54,7 +54,10 @@ def test_interop_with_java_library(glue_client, registry, boto_session):
         }
     )
     proc.check_returncode()
-    assert deserializer.deserialize('test', proc.stdout) == data
+    deserialized = deserializer.deserialize('test', proc.stdout)
+    assert deserialized
+    assert deserialized.data == data
+    assert deserialized.schema == SCHEMA
 
 
 def compile_java():

tests/integration/kafka/test_kafka_integration.py

@@ -94,5 +94,7 @@ def test_produce_consume_with_ser_de_schema_registry(
     assert len(batch) == 1
     messages = batch[list(batch.keys())[0]]
     assert len(messages) == 2
-    assert messages[0].value == data1
-    assert messages[1].value == data2
+    assert messages[0].value.data == data1
+    assert messages[0].value.schema == SCHEMA_V1
+    assert messages[1].value.data == data2
+    assert messages[1].value.schema == SCHEMA_V2