fetch using qset and change query count back

Author: somewes

entity/sync.py

@@ -118,22 +118,26 @@ def _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_al
     return super_entities_by_ctype
 
 
-def _fetch_entity_models(model_ids_to_sync, model_objs_map, model_objs_by_ctype):
+def _fetch_entity_models(model_ids_to_sync, model_objs_map, model_objs_by_ctype, sync_all):
     """
     Fetch the entity models per content type. This will also handle the
     case where accounts are created before _get_super_entities_by_ctype and
     the model_ids_to_sync do not match the model_objs_map
     """
     for ctype, model_ids in model_ids_to_sync.items():
 
-        # Build a set of ids of already fetched models
-        fetched_model_ids = {
-            model.id
-            for model in model_objs_by_ctype[ctype]
-        }
+        if sync_all:
+
+            # Build a set of ids of already fetched models
+            fetched_model_ids = {
+                model.id
+                for model in model_objs_by_ctype[ctype]
+            }
 
-        # Compute the set diff to see if any records are missing
-        unfetched_model_ids = model_ids - fetched_model_ids
+            # Compute the set diff to see if any records are missing
+            unfetched_model_ids = model_ids - fetched_model_ids
+        else:
+            unfetched_model_ids = model_ids
 
         # Check if new records
         if unfetched_model_ids:
@@ -146,13 +150,13 @@ def _fetch_entity_models(model_ids_to_sync, model_objs_map, model_objs_by_ctype)
                 model_objs_map[(ctype, model_obj.id)] = model_obj
 
 
-def _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, model_objs_by_ctype):
+def _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, model_objs_by_ctype, sync_all):
     """
     Given the model IDs to sync, fetch all model objects to sync
     """
     model_objs_to_sync = {}
 
-    _fetch_entity_models(model_ids_to_sync, model_objs_map, model_objs_by_ctype)
+    _fetch_entity_models(model_ids_to_sync, model_objs_map, model_objs_by_ctype, sync_all)
 
     for ctype, model_ids_to_sync_for_ctype in model_ids_to_sync.items():
         model_objs_to_sync[ctype] = [
@@ -261,7 +265,7 @@ def sync(self):
         # Now that we have all models we need to sync, fetch them so that we can extract
         # metadata and entity kinds. If we are syncing all entities, we've already fetched
         # everything and can fill in this data struct without doing another DB hit
-        model_objs_to_sync = _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, model_objs_by_ctype)
+        model_objs_to_sync = _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, model_objs_by_ctype, sync_all)
 
         # Obtain all entity kind tuples associated with the models
         entity_kind_tuples_to_sync = set()

entity/tests/sync_tests.py

@@ -954,7 +954,7 @@ def test_optimal_queries_registered_entity_with_no_qset(self):
         team_group = G(TeamGroup)
 
         ContentType.objects.clear_cache()
-        with self.assertNumQueries(14):
+        with self.assertNumQueries(15):
             team_group.save()
 
     def test_optimal_queries_registered_entity_w_qset(self):
@@ -964,7 +964,7 @@ def test_optimal_queries_registered_entity_w_qset(self):
         account = G(Account)
 
         ContentType.objects.clear_cache()
-        with self.assertNumQueries(17):
+        with self.assertNumQueries(18):
             account.save()
 
     def test_sync_all_optimal_queries(self):