Dequeue transferred records using an explicit database handle

morango/sync/context.py

@@ -4,6 +4,7 @@
 from morango.models.certificates import Filter
 from morango.models.core import SyncSession
 from morango.models.core import TransferSession
+from morango.sync.db import SyncDatabase
 from morango.utils import CAPABILITIES
 from morango.utils import parse_capabilities_from_server_request
 
@@ -20,6 +21,7 @@ class SessionContext(object):
         "is_push",
         "capabilities",
         "error",
+        "db",
     )
     max_backoff_interval = 5
 
@@ -30,6 +32,7 @@ def __init__(
         sync_filter=None,
         is_push=None,
         capabilities=None,
+        db=None,
     ):
         """
         :param sync_session: The sync session instance
@@ -42,13 +45,16 @@ def __init__(
         :type is_push: bool
         :param capabilities: Capabilities set that is combined (intersected) with our own capabilities
         :type capabilities: set|None
+        :param db: Database handle for local sync operations, defaulting to routed access
+        :type db: SyncDatabase|None
         """
         self.sync_session = sync_session
         self.transfer_session = transfer_session
         self.filter = sync_filter
         self.is_push = is_push
         self.capabilities = set(capabilities or []) & CAPABILITIES
         self.error = None
+        self.db = db if db is not None else SyncDatabase.default()
 
         if self.transfer_session:
             self.sync_session = transfer_session.sync_session or self.sync_session
@@ -179,6 +185,7 @@ def __getstate__(self):
 
     def __setstate__(self, state):
         """Re-apply dict state after serialization"""
+        self.db = SyncDatabase.default()
         sync_session_id = state.get("sync_session_id", None)
         if sync_session_id is not None:
             self.sync_session = SyncSession.objects.get(pk=sync_session_id)
@@ -498,6 +505,7 @@ def __getstate__(self):
 
     def __setstate__(self, state):
         """Re-apply dict state after serialization"""
+        self.db = SyncDatabase.default()
         self.children = state.get("children", [])
         self._counter = state.get("counter", 0)
         self._stage = state.get("stage", None)

morango/sync/operations.py

@@ -25,10 +25,10 @@
 from morango.models.certificates import Filter
 from morango.models.core import Buffer
 from morango.models.core import DatabaseMaxCounter
-from morango.models.core import InstanceIDModel
 from morango.models.core import RecordMaxCounter
 from morango.models.core import RecordMaxCounterBuffer
 from morango.models.core import Store
+from morango.models.core import SyncSession
 from morango.models.core import TransferSession
 from morango.models.core import UUIDField
 from morango.models.fsic_utils import calculate_directional_fsic_diff
@@ -332,13 +332,13 @@ def _queue_into_buffer_v2(transfersession, chunk_size=200):
             )
 
 
-def _update_legacy_self_ref_order_for_model(queryset):
+def _update_legacy_self_ref_order_for_model(queryset, db):
     # 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(
+    parent = db.manager(Store).filter(
         id=Cast(OuterRef("_self_ref_fk"), UUIDField()),
         _self_ref_order__isnull=False,
     )
@@ -349,45 +349,52 @@ def _update_legacy_self_ref_order_for_model(queryset):
         pass
 
 
-def _update_legacy_self_ref_order(transfer_session):
-    profile = transfer_session.sync_session.profile
-    transferred_store_records = Store.objects.filter(
+def _update_legacy_self_ref_order(transfer_session, db):
+    profile = db.manager(SyncSession).get(pk=transfer_session.sync_session_id).profile
+    transferred_store_records = db.manager(Store).filter(
         last_transfer_session_id=transfer_session.id,
         profile=profile,
     )
 
     for Model in syncable_models.get_models(profile):
         queryset = transferred_store_records.filter(model_name=Model.morango_model_name)
         if self_referential_fk(Model):
-            _update_legacy_self_ref_order_for_model(queryset)
+            _update_legacy_self_ref_order_for_model(queryset, db)
         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):
+def _dequeue_into_store(transfer_session, fsic, v2_format, self_ref_order, db):
     """
     Takes data from the buffers and merges into the store and record max counters.
 
+    :param transfer_session: Transfer session whose buffered records will be merged
+    :param fsic: JSON-encoded filter-specific instance counters from the sender
+    :param v2_format: Whether the counters use the v2 format
+    :param self_ref_order: Whether the sender supplied self-referential ordering
+    :param db: Database handle for the dequeue transaction, SQL, and instance counter
+    :type db: morango.sync.db.SyncDatabase
+
     ALGORITHM: Incrementally insert and delete on a case by case basis to ensure subsequent cases
     are not affected by previous cases.
     """
 
-    with begin_transaction(Filter(transfer_session.filter)):
-        with connection.cursor() as cursor:
-            DBBackend._dequeuing_delete_rmcb_records(cursor, transfer_session.id)
-            DBBackend._dequeuing_delete_buffered_records(cursor, transfer_session.id)
-            current_id = InstanceIDModel.get_current_instance_and_increment_counter()
-            DBBackend._dequeuing_merge_conflict_buffer(cursor, current_id, transfer_session.id)
-            DBBackend._dequeuing_merge_conflict_rmcb(cursor, transfer_session.id)
-            DBBackend._dequeuing_update_rmcs_last_saved_by(cursor, current_id, transfer_session.id)
-            DBBackend._dequeuing_delete_mc_rmcb(cursor, transfer_session.id)
-            DBBackend._dequeuing_delete_mc_buffer(cursor, transfer_session.id)
-            DBBackend._dequeuing_insert_remaining_buffer(cursor, transfer_session.id)
-            DBBackend._dequeuing_insert_remaining_rmcb(cursor, transfer_session.id)
+    with db.begin_transaction(Filter(transfer_session.filter)):
+        with db.cursor() as cursor:
+            db.backend._dequeuing_delete_rmcb_records(cursor, transfer_session.id)
+            db.backend._dequeuing_delete_buffered_records(cursor, transfer_session.id)
+            current_id = db.current_instance_and_increment_counter()
+            db.backend._dequeuing_merge_conflict_buffer(cursor, current_id, transfer_session.id)
+            db.backend._dequeuing_merge_conflict_rmcb(cursor, transfer_session.id)
+            db.backend._dequeuing_update_rmcs_last_saved_by(cursor, current_id, transfer_session.id)
+            db.backend._dequeuing_delete_mc_rmcb(cursor, transfer_session.id)
+            db.backend._dequeuing_delete_mc_buffer(cursor, transfer_session.id)
+            db.backend._dequeuing_insert_remaining_buffer(cursor, transfer_session.id)
+            db.backend._dequeuing_insert_remaining_rmcb(cursor, transfer_session.id)
             if not self_ref_order:
-                _update_legacy_self_ref_order(transfer_session)
-            DBBackend._dequeuing_delete_remaining_rmcb(cursor, transfer_session.id)
-            DBBackend._dequeuing_delete_remaining_buffer(cursor, transfer_session.id)
+                _update_legacy_self_ref_order(transfer_session, db)
+            db.backend._dequeuing_delete_remaining_rmcb(cursor, transfer_session.id)
+            db.backend._dequeuing_delete_remaining_buffer(cursor, transfer_session.id)
 
         DatabaseMaxCounter.update_fsics(
             json.loads(fsic),
@@ -712,6 +719,7 @@ def handle(self, context):
                 fsic,
                 v2_format=FSIC_V2_FORMAT in context.capabilities,
                 self_ref_order=SELF_REF_ORDER in context.capabilities,
+                db=context.db,
             )
 
         return transfer_statuses.COMPLETED

tests/testapp/tests/sync/test_context.py

@@ -15,12 +15,18 @@
 from morango.sync.context import NetworkSessionContext
 from morango.sync.context import SessionContext
 from morango.sync.controller import SessionController
+from morango.sync.db import SyncDatabase
 
 from ..helpers import create_dummy_store_data
 from ..helpers import TestSessionContext
 
 
 class SessionContextTestCase(SimpleTestCase):
+    def test_database_handle(self):
+        self.assertIsNone(TestSessionContext().db.alias)
+        db = SyncDatabase("default2")
+        self.assertIs(TestSessionContext(db=db).db, db)
+
     def test_init__nothing(self):
         context = TestSessionContext()
         context.update_state(
@@ -259,6 +265,15 @@ def test_init(self):
 
 
 class ContextPicklingTestCase(TestCase):
+    def test_local_restores_default_database_handle(self):
+        context = LocalSessionContext(db=SyncDatabase("default2"))
+        self.assertNotIn("db", context.__getstate__())
+
+        restored = pickle.loads(pickle.dumps(context))
+
+        self.assertIsInstance(restored.db, SyncDatabase)
+        self.assertIsNone(restored.db.alias)
+
     def test_basic(self):
         data = create_dummy_store_data()
         transfer_session = data["tx"]
@@ -319,7 +334,7 @@ def test_composite(self):
         )
         network = NetworkSessionContext(conn)
 
-        composite = CompositeSessionContext([local, network])
+        composite = CompositeSessionContext([local, network], db=SyncDatabase("default2"))
         for _ in range(2):
             composite.update_state(
                 stage=transfer_stages.INITIALIZING, stage_status=transfer_statuses.COMPLETED
@@ -333,8 +348,10 @@ def test_composite(self):
         self.assertEqual(composite._counter, unpickled_context._counter)
         self.assertEqual(composite.stage, unpickled_context.stage)
         self.assertEqual(composite.stage_status, unpickled_context.stage_status)
+        self.assertIsNone(unpickled_context.db.alias)
         for i, context_type in enumerate([LocalSessionContext, NetworkSessionContext]):
             self.assertIsInstance(unpickled_context.children[i], context_type)
+            self.assertIsNone(unpickled_context.children[i].db.alias)
 
 
 class CompositeSessionContextTestCase(SimpleTestCase):