Refactor serialize source to use new base class

Author: bjester

morango/sync/stream/serialize.py

@@ -8,7 +8,6 @@
 from typing import Type
 
 from django.core.serializers.json import DjangoJSONEncoder
-from django.db.models import Q
 
 from morango.models.certificates import Filter
 from morango.models.core import DatabaseMaxCounter
@@ -21,15 +20,16 @@
 from morango.registry import syncable_models
 from morango.sync.stream.core import Buffer
 from morango.sync.stream.core import Sink
-from morango.sync.stream.core import Source
 from morango.sync.stream.core import Transform
 from morango.sync.stream.core import Unbuffer
+from morango.sync.stream.source import MorangoSource
+from morango.sync.stream.source import SourceTask
 from morango.utils import self_referential_fk
 
 logger = logging.getLogger(__name__)
 
 
-class SerializeTask(object):
+class SerializeTask(SourceTask):
     """Carrier class for providing context through the pipeline"""
 
     __slots__ = ("model", "obj", "store", "counter")
@@ -40,6 +40,10 @@ def __init__(self, model: Type[SyncableModel], obj: SyncableModel):
         self.store: Optional[Store] = None
         self.counter: Optional[RecordMaxCounter] = None
 
+    @property
+    def id(self) -> str:
+        return self.obj.id
+
     @property
     def is_store_update(self):
         return self.store is not None and not self.store._state.adding
@@ -59,59 +63,22 @@ def self_referential_fk(self) -> Optional[str]:
         return self_referential_fk(self.model)
 
 
-class AppModelSource(Source[SerializeTask]):
+class AppModelSource(MorangoSource[SerializeTask]):
     """
     Yields ``SerializeTask`` objects for every syncable-model record that matches the
     optional *sync_filter*.
     """
 
-    def __init__(
-        self,
-        profile: str,
-        sync_filter: Optional[Filter] = None,
-        dirty_only: bool = True,
-        partition_order: str = "asc",
-    ):
-        """
-        :param profile: The Morango model profile
-        :param sync_filter: The Filter object for this sync
-        :param dirty_only: Whether to filter on dirty records only
-        :param partition_order: Controls how the filter specificity is applied, "asc" or "desc"
-        """
-        self.profile = profile
-        self.sync_filter = sync_filter
-        self.dirty_only = dirty_only
-        self.partition_order = partition_order
-        self._seen = set()
-
-    def prefix_conditions(self) -> Generator[Optional[Q], None, None]:
-        if self.sync_filter is None:
-            # yield None once, so we do one query without a partition filter (everything)
-            yield None
-        else:
-            partitions_prefixes = [str(prefix) for prefix in self.sync_filter]
-            partition_iterator = sorted(
-                partitions_prefixes,
-                reverse=self.partition_order == "desc",
-            )
-
-            for prefix in partition_iterator:
-                yield Q(_morango_partition__startswith=prefix)
-
-    def stream(self) -> Generator[SerializeTask, None, None]:
-        for partition_condition in self.prefix_conditions():
-            for qs in syncable_models.get_model_querysets(self.profile):
-                if partition_condition is not None:
-                    qs = qs.filter(partition_condition)
-                if self.dirty_only:
-                    qs = qs.filter(_morango_dirty_bit=True)
-                for obj in qs.iterator():
-                    # partition filtering could result in overlaps, and since we're walking
-                    # through the partitions one by one, we should avoid duplicates. Morango
-                    # syncable models have unique IDs across the entire profile
-                    if obj.id not in self._seen:
-                        self._seen.add(obj.id)
-                        yield SerializeTask(qs.model, obj)
+    def stream_for_filter(
+        self, partition_condition: Optional[str]
+    ) -> Generator[SerializeTask, None, None]:
+        for qs in syncable_models.get_model_querysets(self.profile):
+            if partition_condition is not None:
+                qs = qs.filter(_morango_partition__startswith=partition_condition)
+            if self.dirty_only:
+                qs = qs.filter(_morango_dirty_bit=True)
+            for obj in qs.iterator():
+                yield SerializeTask(qs.model, obj)
 
 
 class StoreLookup(Transform[List[SerializeTask]]):

tests/testapp/tests/sync/stream/test_serialize.py

@@ -1,7 +1,6 @@
 import json
 
 import mock
-from django.db.models import Q
 from django.test import SimpleTestCase
 
 from morango.models.certificates import Filter
@@ -59,7 +58,7 @@ def test_prefix_conditions__with_filter(self):
         conditions = list(source.prefix_conditions())
         self.assertEqual(len(conditions), 2)
         self.assertEqual(
-            str(conditions[0]), "(AND: ('_morango_partition__startswith', 'a'))"
+            str(conditions[0]), "a"
         )
 
     @mock.patch("morango.sync.stream.serialize.syncable_models.get_model_querysets")
@@ -116,7 +115,7 @@ def test_stream__partition(self, mock_get_model_querysets):
         self.assertEqual(tasks[0].model, model)
         self.assertEqual(tasks[0].obj, obj)
         mock_get_model_querysets.assert_called_once_with("test")
-        qs.filter.assert_called_once_with(Q(_morango_partition__startswith="a"))
+        qs.filter.assert_called_once_with(_morango_partition__startswith="a")
 
 
 class StoreLookupTestCase(SimpleTestCase):