Streaming deserialize: add source stream and task

morango/models/core.py

@@ -23,6 +23,7 @@
 from django.db.models.expressions import CombinedExpression
 from django.db.models.fields.related import ForeignKey
 from django.db.models.functions import Cast
+from django.db.models.functions import NullIf
 from django.utils import timezone
 from django.utils.functional import cached_property
 
@@ -440,10 +441,30 @@ def char_ids_list(self):
             .values_list("fixed_id", flat=True)
         )
 
+    def filter_deserialization_error(self, has_error: bool) -> "StoreQueryset":
+        """
+        Filters the queryset to return Store records that have or have not any
+        deserialization errors
+        """
+        # this nullIf assertion is generally more performant than an OR'd statement on
+        # unindexed columns, in both SQLite and PostgreSQL
+        return self.annotate(
+            _deserialization_error=NullIf(
+                F("deserialization_error"), Value(""), output_field=models.TextField()
+            )
+        ).filter(_deserialization_error__isnull=not has_error)
+
+    def filter_has_deserialization_error(self) -> "StoreQueryset":
+        """Filters the queryset to return Store records that have deserialization errors"""
+        return self.filter_deserialization_error(True)
+
+    def exclude_has_deserialization_error(self) -> "StoreQueryset":
+        """Filters the queryset to return Store records that have no deserialization error"""
+        return self.filter_deserialization_error(False)
+
 
-class StoreManager(models.Manager):
-    def get_queryset(self):
-        return StoreQueryset(self.model, using=self._db)
+class StoreManager(models.Manager.from_queryset(StoreQueryset)):
+    pass
 
 
 class Store(AbstractStore):

morango/registry.py

@@ -110,6 +110,18 @@ def get_model_querysets(self, profile) -> Generator[QuerySet, None, None]:
                 queryset = queryset.order_by(*self._get_nulls_last_ordering(ordering))
             yield queryset
 
+    def get_store_querysets(self, profile) -> Generator[QuerySet, None, None]:
+        """
+        Complementary method to `get_model_querysets` but for Store querysets
+        """
+        from morango.models.core import Store
+
+        for model in self.get_models(profile):
+            store_qs = Store.objects.filter(profile=profile, model_name=model.morango_model_name)
+            if self.get_self_referential_fk(model) is not None:
+                store_qs = store_qs.order_by(*self._get_nulls_last_ordering(("_self_ref_order",)))
+            yield store_qs
+
     @staticmethod
     def _get_nulls_last_ordering(ordering):
         normalized = []

morango/sync/stream/deserialize.py

@@ -0,0 +1,92 @@
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Optional
+from typing import Type
+
+from morango.models.certificates import Filter
+from morango.models.core import Store
+from morango.models.core import SyncableModel
+from morango.registry import syncable_models
+from morango.sync.stream.source import MorangoSource
+from morango.sync.stream.source import SourceTask
+
+
+class DeserializeTask(SourceTask):
+    """Carrier class for providing context through the deserialization pipeline."""
+
+    __slots__ = ("store", "app_model", "fk_cache", "errors")
+
+    def __init__(self, store: Store, fk_cache: Dict):
+        self.store = store
+        self.fk_cache: Dict = fk_cache
+        self.app_model: Optional[SyncableModel] = None
+        self.errors: List[Exception] = []
+
+    @property
+    def id(self) -> str:
+        return self.store.id
+
+    @property
+    def model(self) -> Type[SyncableModel]:
+        return syncable_models.get_model(self.store.profile, self.store.model_name)
+
+    @property
+    def has_errors(self) -> bool:
+        return len(self.errors) > 0
+
+    def set_app_model(self, app_model: Optional[SyncableModel]) -> None:
+        self.app_model = app_model
+
+    def add_error(self, error: Exception) -> None:
+        self.errors.append(error)
+
+
+class StoreModelSource(MorangoSource[DeserializeTask]):
+    """
+    Yields ``DeserializeTask`` objects for dirty store models that match the optional
+    *sync_filter*.
+    """
+
+    def __init__(
+        self,
+        profile: str,
+        sync_filter: Optional[Filter] = None,
+        dirty_only: bool = True,
+        partition_order: str = "asc",
+        fk_cache: Optional[Dict] = None,
+        skip_errored: bool = False,
+    ):
+        """
+        :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"
+        :param fk_cache: Dictionary cache for FK references
+        :param skip_errored: Whether to skip Store records with deserialization errors
+        """
+        super().__init__(profile, sync_filter, dirty_only, partition_order)
+        self.fk_cache = fk_cache if fk_cache is not None else {}
+        self.skip_errored = skip_errored
+
+    def begin(self) -> None:
+        """Reset fk_cache at the beginning of stream"""
+        super().begin()
+        self.fk_cache.clear()
+
+    def stream_for_filter(
+        self, partition_condition: Optional[str]
+    ) -> Generator[DeserializeTask, None, None]:
+        # the registry yields models in foreign key dependency order, so streaming model by model
+        # ensures a record's foreign key targets are deserialized before it is
+        for store_qs in syncable_models.get_store_querysets(self.profile):
+            qs = store_qs
+            if partition_condition is not None:
+                qs = qs.filter(partition__startswith=partition_condition)
+            if self.dirty_only:
+                qs = qs.filter(dirty_bit=True)
+            if self.skip_errored:
+                qs = qs.exclude_has_deserialization_error()
+
+            for store_model in qs.iterator():
+                yield DeserializeTask(store_model, self.fk_cache)

morango/sync/stream/serialize.py

@@ -6,7 +6,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
@@ -19,14 +18,15 @@
 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
 
 logger = logging.getLogger(__name__)
 
 
-class SerializeTask(object):
+class SerializeTask(SourceTask):
     """Carrier class for providing context through the pipeline"""
 
     __slots__ = (
@@ -35,7 +35,6 @@ class SerializeTask(object):
         "store",
         "counter",
         "_self_ref_fk_value",
-        "_self_ref_fk_value",
         "_self_ref_order",
     )
 
@@ -47,6 +46,10 @@ def __init__(self, model: Type[SyncableModel], obj: SyncableModel):
         self._self_ref_fk_value: Optional[str] = None
         self._self_ref_order: Optional[int] = 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
@@ -80,59 +83,22 @@ def set_self_ref_order(self, value: Optional[int]):
         self._self_ref_order = value
 
 
-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]]):

morango/sync/stream/source.py

@@ -0,0 +1,104 @@
+import abc
+from typing import Generator
+from typing import Iterator
+from typing import Optional
+from typing import TypeVar
+
+from morango.models.certificates import Filter
+from morango.sync.stream.core import Source
+
+
+class SourceTask(abc.ABC):
+    """Typing for source object passed through streaming pipeline"""
+
+    __slots__ = ()
+
+    @property
+    @abc.abstractmethod
+    def id(self) -> str:
+        pass
+
+
+T = TypeVar("T", bound=SourceTask)
+
+
+class MorangoSource(Source[T], abc.ABC):
+    """
+    Common source functionality for Morango sources, such as SyncableModels and Store records.
+    """
+
+    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: Optional[set] = None
+
+    def begin(self) -> None:
+        """Initialize seen set at the beginning of the stream"""
+        self._seen = set()
+
+    def prefix_conditions(self) -> Generator[Optional[str], None, None]:
+        """
+        Generates partition prefixes for queries based on the sync filter and partition order.
+
+        This method outputs prefixes in sorted order according to the specified partition
+        order. If no sync filter is provided, it yields `None` to indicate a query
+        without filtering by partition.
+
+        :return: A generator yielding partition prefixes or `None` if no filtering is applied.
+        """
+        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 prefix
+
+    def stream(self) -> Generator[T, None, None]:
+        """
+        Streams unique objects based on prefix conditions. This generator method iterates over
+        partition conditions defined in the sync_filter and passes through to `stream_for_filter`
+        to stream back objects, ensuring that only objects with unique `id` values are yielded.
+
+        :return: A generator yielding unique objects.
+        """
+        for partition_condition in self.prefix_conditions():
+            for obj in self.stream_for_filter(partition_condition):
+                # partition filtering could result in overlaps, and since we're walking
+                # through the partitions one by one, we should avoid duplicates. Morango
+                # syncable models and store records have unique IDs across the entire profile
+                if obj.id not in self._seen:
+                    # without sync filters, we do not need to worry about repeating objects
+                    if self.sync_filter is not None:
+                        self._seen.add(obj.id)
+                    yield obj
+
+    @abc.abstractmethod
+    def stream_for_filter(self, partition_condition: Optional[str]) -> Iterator[T]:
+        """
+        This method is intended to generate an iterator that yields data based on the given
+        filtering condition.
+
+        :param partition_condition: A string representing a partition filter prefix condition
+        :return: An iterator yielding items
+        """
+        pass