refactor self-ref order fallback to use model-scoped db updates

Author: ozer550

morango/sync/operations.py

@@ -8,7 +8,10 @@
 from django.core import exceptions
 from django.db import connection
 from django.db.models import CharField
+from django.db.models import Exists
+from django.db.models import OuterRef
 from django.db.models import Q
+from django.db.models import Subquery
 from django.db.models import signals
 from django.db.utils import OperationalError
 from django.utils import timezone
@@ -680,49 +683,38 @@ def _queue_into_buffer_v2(transfersession, chunk_size=200):
             )
 
 
+def _update_legacy_self_ref_order_for_model(queryset):
+    # root nodes set the _self_ref_order to 0
+    queryset.filter(_self_ref_fk="").exclude(_self_ref_order=0).update(_self_ref_order=0)
+    # reset the _self_ref_order to None for all records that have a parent
+    queryset.exclude(_self_ref_fk="").exclude(_self_ref_order=None).update(
+        _self_ref_order=None
+    )
+
+    parent = Store.objects.filter(
+        id=OuterRef("_self_ref_fk"),
+        _self_ref_order__isnull=False,
+    )
+    parent_order = parent.values("_self_ref_order")[:1]
+    pending = queryset.exclude(_self_ref_fk="").filter(_self_ref_order=None)
+
+    while pending.filter(Exists(parent)).update(_self_ref_order=Subquery(parent_order)):
+        pass
+
+
 def _update_legacy_self_ref_order(transfer_session):
+    profile = transfer_session.sync_session.profile
     transferred_store_records = Store.objects.filter(
-        last_transfer_session_id=transfer_session.id
+        last_transfer_session_id=transfer_session.id,
+        profile=profile,
     )
-    records_by_id = {record.id: record for record in transferred_store_records}
-    records_by_model = defaultdict(list)
-
-    for record in records_by_id.values():
-        try:
-            Model = syncable_models.get_model(record.profile, record.model_name)
-        except KeyError:
-            continue
 
+    for Model in syncable_models.get_models(profile):
+        queryset = transferred_store_records.filter(model_name=Model.morango_model_name)
         if self_referential_fk(Model):
-            records_by_model[(record.profile, record.model_name)].append(record)
-        elif record._self_ref_order is not None:
-            record._self_ref_order = None
-            record.save(update_fields=["_self_ref_order"])
-
-    for records in records_by_model.values():
-        pending_record_ids = set(record.id for record in records)
-        while pending_record_ids:
-            updated = False
-            for record_id in list(pending_record_ids):
-                record = records_by_id[record_id]
-                if not record._self_ref_fk:
-                    order = 0
-                else:
-                    parent = records_by_id.get(record._self_ref_fk)
-                    if parent is not None and parent.id in pending_record_ids:
-                        continue
-                    order = Store.objects.filter(id=record._self_ref_fk).values_list(
-                        "_self_ref_order", flat=True
-                    ).first()
-
-                if record._self_ref_order != order:
-                    record._self_ref_order = order
-                    record.save(update_fields=["_self_ref_order"])
-                pending_record_ids.remove(record_id)
-                updated = True
-
-            if not updated:
-                break
+            _update_legacy_self_ref_order_for_model(queryset)
+        else:
+            queryset.exclude(_self_ref_order=None).update(_self_ref_order=None)
 
 
 def _dequeue_into_store(transfer_session, fsic, v2_format=False, self_ref_order=True):

tests/testapp/tests/sync/test_operations.py

@@ -37,6 +37,7 @@
     _deserialize_from_store,
     _queue_into_buffer_v1,
     _queue_into_buffer_v2,
+    _update_legacy_self_ref_order,
 )
 from morango.sync.syncsession import TransferClient
 
@@ -654,7 +655,7 @@ def setUp(self):
         conn.server_info = dict(capabilities=[])
         self.data["mc"] = MorangoProfileController("facilitydata")
         session = SyncSession.objects.create(
-            id=uuid.uuid4().hex, profile="", last_activity_timestamp=timezone.now()
+            id=uuid.uuid4().hex, profile="facilitydata", last_activity_timestamp=timezone.now()
         )
         self.transfer_session = TransferSession.objects.create(
             id=uuid.uuid4().hex,
@@ -684,6 +685,21 @@ def assert_store_records_not_tagged_with_last_session(self, store_ids):
             except Store.DoesNotExist:
                 pass
 
+    def _make_transferred_store(self, **kwargs):
+        defaults = {
+            "id": uuid.uuid4().hex,
+            "serialized": "{}",
+            "last_saved_instance": self.current_id.id,
+            "last_saved_counter": 1,
+            "model_name": "facility",
+            "profile": "facilitydata",
+            "partition": uuid.uuid4().hex,
+            "source_id": uuid.uuid4().hex,
+            "last_transfer_session_id": self.transfer_session.id,
+        }
+        defaults.update(kwargs)
+        return Store.objects.create(**defaults)
+
     def test_dequeuing_sets_last_session(self):
         store_ids = [self.data[key] for key in ["model2", "model3", "model4", "model5", "model7"]]
         self.assert_store_records_not_tagged_with_last_session(store_ids)
@@ -1009,6 +1025,37 @@ def test_dequeue_into_store__self_ref_order_fallback_for_missing_capability(self
         self.assertEqual(Store.objects.get(id=self.data["model3"])._self_ref_order, 0)
         self.assertEqual(Store.objects.get(id=self.data["model4"])._self_ref_order, 0)
 
+    def test_update_legacy_self_ref_order_nulls_non_self_ref_models(self):
+        store = self._make_transferred_store(
+            model_name=SummaryLog.morango_model_name,
+            _self_ref_order=3,
+        )
+
+        _update_legacy_self_ref_order(self.transfer_session)
+
+        store.refresh_from_db()
+        self.assertIsNone(store._self_ref_order)
+
+    def test_update_legacy_self_ref_order_handles_deeper_self_ref_chains(self):
+        root = self._make_transferred_store(_self_ref_fk="", _self_ref_order=None)
+        child = self._make_transferred_store(
+            _self_ref_fk=root.id,
+            _self_ref_order=None,
+        )
+        grandchild = self._make_transferred_store(
+            _self_ref_fk=child.id,
+            _self_ref_order=None,
+        )
+
+        _update_legacy_self_ref_order(self.transfer_session)
+
+        root.refresh_from_db()
+        child.refresh_from_db()
+        grandchild.refresh_from_db()
+        self.assertEqual(root._self_ref_order, 0)
+        self.assertEqual(child._self_ref_order, 0)
+        self.assertEqual(grandchild._self_ref_order, 0)
+
     def test_local_dequeue_operation(self):
         self.transfer_session.records_transferred = 1
         self.context.filter = [self.transfer_session.filter]