Schema registry code passes local tests

Author: Fireye04

docs/documenteer.toml

@@ -9,7 +9,17 @@ package = "sasquatchbackpack"
 rst_epilog_file = "_rst_epilog.rst"
 disable_primary_sidebars = ["index", "changelog"]
 extensions = ['sphinx_click']
-nitpick_ignore = [["py:class", "faststream.kafka.KafkaBroker"]]
+nitpick_ignore = [
+	[
+		"py:class",
+		"faststream.kafka.KafkaBroker",
+	],
+	[
+		"py:class",
+		"dataclasses_avroschema.pydantic.main.AvroBaseModel",
+	],
+]
 
 [sphinx.intersphinx.projects]
+safir = "https://safir.lsst.io"
 python = "https://docs.python.org/3/"

src/sasquatchbackpack/sasquatch.py

@@ -40,13 +40,18 @@ class DataSource(ABC):
         Specific source name, used as an identifier
     """
 
-    def __init__(
-        self, topic_name: str, schema: str, *, uses_redis: bool
-    ) -> None:
+    def __init__(self, topic_name: str, *, uses_redis: bool) -> None:
         self.topic_name = topic_name
-        self.schema = schema
         self.uses_redis = uses_redis
 
+    @abstractmethod
+    def get_schema(self, namespace: str) -> AvroBaseModel:
+        pass
+
+    @abstractmethod
+    def assemble_schema(self, records: dict, namespace: str) -> AvroBaseModel:
+        pass
+
     @abstractmethod
     def get_records(self) -> list[dict]:
         pass
@@ -66,6 +71,7 @@ class RedisManager:
     """
 
     def __init__(self, address: str) -> None:
+        nest_asyncio.apply()
         self.address = address
         self.model = redis.from_url(self.address)
 
@@ -173,10 +179,10 @@ async def _dispatch(
     except Exception as e:
         return e
 
-    await schema_manager.register_model(schema)
+    await schema_manager.register_model(type(schema))
 
     for record in records:
-        avro: bytes = schema_manager.serialize(
+        avro: bytes = await schema_manager.serialize(
             source.assemble_schema(record, namespace)
         )
 

src/sasquatchbackpack/sources/usgs/scripts.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 
+from dataclasses_avroschema.pydantic import AvroBaseModel
 from libcomcat.search import search
 
 from sasquatchbackpack.sasquatch import DataSource
@@ -99,9 +100,7 @@ def __init__(
         self,
         config: USGSConfig,
     ) -> None:
-        super().__init__(
-            config.topic_name, config.schema, uses_redis=config.uses_redis
-        )
+        super().__init__(config.topic_name, uses_redis=config.uses_redis)
         self.duration = config.duration
         self.config = config
         self.radius = config.radius
@@ -144,9 +143,7 @@ def get_records(self) -> list[dict]:
                 f"A connection error occurred while fetching records: {ce}"
             ) from ce
 
-    def assemble_schema(
-        self, record: dict, namespace: str
-    ) -> EarthquakeSchema.avro_schema():
+    def assemble_schema(self, record: dict, namespace: str) -> AvroBaseModel:
         schema = {
             "timestamp": record["timestamp"],
             "id": record["id"],
@@ -156,13 +153,19 @@ def assemble_schema(
             "magnitude": record["magnitude"],
             "namespace": namespace,
         }
-        return EarthquakeSchema(**schema)
+        return EarthquakeSchema.parse_obj(data=schema)
 
-    def get_schema(self, namespace: str) -> EarthquakeSchema.avro_schema():
+    def get_schema(self, namespace: str) -> AvroBaseModel:
         schema = {
+            "timestamp": 1,
+            "id": "default",
+            "latitude": 1.0,
+            "longitude": 1.0,
+            "depth": 1.0,
+            "magnitude": 1.0,
             "namespace": namespace,
         }
-        return EarthquakeSchema(**schema)
+        return EarthquakeSchema.parse_obj(data=schema)
 
     def get_redis_key(self, datapoint: dict) -> str:
         """Allow USGS API to format its own redis keys.

tests/dispatcher_test.py

@@ -18,19 +18,32 @@ class TestSchema(AvroBaseModel):
     class Meta:
         """Schema metadata."""
 
-        namespace = "$namespace"
+        namespace = "Default"
         schema_name = "testSchema"
 
 
 class TestSource(sasquatch.DataSource):
     def __init__(self, current_records: list[dict[str, str]]) -> None:
         super().__init__(
             "test",
-            TestSchema.avro_schema().replace("double", "float"),
             uses_redis=True,
         )
         self.records = current_records
 
+    def assemble_schema(self, record: dict, namespace: str) -> AvroBaseModel:
+        schema = {
+            "id": record["value"]["id"],
+            "namespace": namespace,
+        }
+        return TestSchema.parse_obj(data=schema)
+
+    def get_schema(self, namespace: str) -> AvroBaseModel:
+        schema = {
+            "id": "default",
+            "namespace": namespace,
+        }
+        return TestSchema.parse_obj(data=schema)
+
     def get_records(self) -> list[dict]:
         return [{"value": {"id": record["id"]}} for record in self.records]