Fix bad schema caching logic in serializer

Schema versions were cached by schema, which is incorrect when using
the topic_record_name naming strategy.

Author: cdebost

src/aws_schema_registry/avro.py

@@ -5,10 +5,10 @@
 
 import fastavro
 
-from aws_schema_registry.schema import DataFormat
+from aws_schema_registry.schema import DataFormat, Schema
 
 
-class AvroSchema:
+class AvroSchema(Schema):
     """Implementation of the `Schema` protocol for Avro schemas.
 
     Arguments:
@@ -19,29 +19,47 @@ class AvroSchema:
             itself
     """
 
-    data_format: DataFormat = 'AVRO'
-    parsed: dict
-
     def __init__(self, string: str, return_record_name: bool = False):
-        self.string = string
-        self.parsed = fastavro.parse_schema(json.loads(string))
-        # https://github.com/fastavro/fastavro/issues/415
-        self.name = self.parsed.get('name', self.parsed['type'])
+        self._dict = json.loads(string)
+        self._parsed = fastavro.parse_schema(self._dict)
         self.return_record_name = return_record_name
 
+    def __hash__(self):
+        return hash(str(self))
+
+    def __eq__(self, other):
+        return isinstance(other, AvroSchema) and \
+            self._parsed == other._parsed and \
+            self.return_record_name == other.return_record_name
+
+    def __str__(self):
+        return json.dumps(self._dict)
+
+    def __repr__(self):
+        return '<AvroSchema %s>' % self._dict
+
+    @property
+    def data_format(self) -> DataFormat:
+        return 'AVRO'
+
+    @property
+    def fqn(self) -> str:
+        # https://github.com/fastavro/fastavro/issues/415
+        return self._parsed.get('name', self._parsed['type'])
+
     def read(self, bytes_: bytes):
         b = BytesIO(bytes_)
         value = fastavro.schemaless_reader(
             b,
-            self.parsed,
+            self._parsed,
             return_record_name=self.return_record_name
         )
         b.close()
         return value
 
     def write(self, data) -> bytes:
         b = BytesIO()
-        fastavro.schemaless_writer(b, self.parsed, data)
+        fastavro.schemaless_writer(b, self._parsed, data)
         value = b.getvalue()
         b.close()
         return value

src/aws_schema_registry/naming.py

@@ -61,7 +61,7 @@ def record_name_strategy(topic: str, is_key: bool, schema: Schema) -> str:
     However, this requires that the fully-qualified record names uniquely
     and consistently identify a schema across the entire registry.
     """
-    return schema.name
+    return schema.fqn
 
 
 def topic_record_name_strategy(topic: str, is_key: bool,
@@ -72,4 +72,4 @@ def topic_record_name_strategy(topic: str, is_key: bool,
     Additionally allows different topics to use the same record name for
     incompatible schemas.
     """
-    return f'{topic}-{schema.name}'
+    return f'{topic}-{schema.fqn}'

src/aws_schema_registry/schema.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 import sys
-from typing import Any, Optional
+from typing import Any, Optional, Hashable
 from uuid import UUID
 
 if sys.version_info[1] < 8:  # for py37
-    from typing_extensions import Literal, Protocol
+    from typing_extensions import Literal
 else:
-    from typing import Literal, Protocol
+    from typing import Literal
 
 DataFormat = Literal['AVRO', 'JSON']
 
@@ -33,14 +34,26 @@
 SchemaVersionStatus = Literal['AVAILABLE', 'PENDING', 'FAILURE', 'DELETING']
 
 
-class Schema(Protocol):
-    data_format: DataFormat
-    name: str
-    string: str
+class Schema(ABC, Hashable):
+    """Abstract base class for a schema implementation."""
+
+    @property
+    @abstractmethod
+    def data_format(self) -> DataFormat:
+        """The data format of this schema."""
+
+    @property
+    @abstractmethod
+    def fqn(self) -> str:
+        """The fully-qualified name of this schema."""
 
-    def read(self, bytes_: bytes) -> Any: ...
+    @abstractmethod
+    def read(self, bytes_: bytes) -> Any:
+        """Read bytes into a record."""
 
-    def write(self, data) -> bytes: ...
+    @abstractmethod
+    def write(self, data) -> bytes:
+        """Write a record into bytes."""
 
 
 @dataclass

src/aws_schema_registry/serde.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import logging
 import sys
 from typing import Any, Dict, NamedTuple
@@ -58,8 +59,6 @@ class SchemaRegistrySerializer:
             alternate strategies.
     """
 
-    _cache: Dict[Schema, SchemaVersion]
-
     def __init__(
         self,
         client: SchemaRegistryClient,
@@ -71,32 +70,29 @@ def __init__(
         self.is_key = is_key
         self.compatibility_mode: CompatibilityMode = compatibility_mode
         self.schema_naming_strategy = schema_naming_strategy
-        self._cache = {}
 
-    def serialize(self, topic, data_and_schema: DataAndSchema):
+    def serialize(self, topic: str, data_and_schema: DataAndSchema):
         if data_and_schema is None:
             return None
         if not isinstance(data_and_schema, DataAndSchema):
             raise TypeError('AvroSerializer can only serialize',
                             f' {DataAndSchema}, got {type(data_and_schema)}')
         data, schema = data_and_schema
-        schema_version = self._cache.get(schema)
-        if not schema_version:
-            schema_name = self.schema_naming_strategy(
-                topic, self.is_key, schema
-            )
-            LOG.info('Schema %s not cached locally, registering...',
-                     schema_name)
-            schema_version = self.client.get_or_register_schema_version(
-                definition=schema.string,
-                schema_name=schema_name,
-                data_format=schema.data_format,
-                compatibility_mode=self.compatibility_mode
-            )
-            self._cache[schema] = schema_version
+        schema_version = self._get_schema_version(topic, schema)
         serialized = schema.write(data)
         return encode(serialized, schema_version.version_id)
 
+    @functools.lru_cache(maxsize=None)
+    def _get_schema_version(self, topic: str, schema: Schema) -> SchemaVersion:
+        schema_name = self.schema_naming_strategy(topic, self.is_key, schema)
+        LOG.info('Fetching schema %s...', schema_name)
+        return self.client.get_or_register_schema_version(
+            definition=str(schema),
+            schema_name=schema_name,
+            data_format=schema.data_format,
+            compatibility_mode=self.compatibility_mode
+        )
+
 
 class SchemaRegistryDeserializer:
     """Kafka serializer that uses the AWS Schema Registry.
@@ -146,10 +142,13 @@ def deserialize(self, topic: str, bytes_: bytes):
             schema_version = self.client.get_schema_version(
                 version_id=schema_version_id
             )
-            if schema_version.data_format == 'AVRO':
-                writer_schema = AvroSchema(schema_version.definition)
-            elif schema_version.data_format == 'JSON':
-                raise NotImplementedError('JSON schema not supported')
+            writer_schema = self._create_writer_schema(schema_version)
             self._writer_schemas[schema_version_id] = writer_schema
-            LOG.info('Schema %s fetched', schema_version_id)
+            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':
+            raise NotImplementedError('JSON schema not supported')

tests/test_avro.py

@@ -3,13 +3,13 @@
 
 def test_fully_qualified_name():
     s = AvroSchema('{"type": "record", "namespace": "foo", "name": "Bar"}')
-    assert s.name == "foo.Bar"
+    assert s.fqn == "foo.Bar"
 
 
 def test_primitive_name():
     # fastavro does not fulfill this part of the Avro spec
     s = AvroSchema('{"type": "string"}')
-    assert s.name == 'string'
+    assert s.fqn == 'string'
 
 
 def test_readwrite():