Merge pull request #160 from somewes/develop

fix queries

Author: somewes

entity/sync.py

@@ -118,20 +118,51 @@ def _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_al
     return super_entities_by_ctype
 
 
-def _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, sync_all):
+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():
+
+        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
+        else:
+            unfetched_model_ids = model_ids
+
+        # Check if new records
+        if unfetched_model_ids:
+
+            # Fetch the records and add them to the model_objs_map
+            model_qset = entity_registry.entity_registry.get(ctype.model_class()).queryset
+            model_objs_to_sync = model_qset.filter(id__in=unfetched_model_ids)
+            for model_obj in model_objs_to_sync:
+                model_objs_by_ctype[ctype].append(model_obj)
+                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, sync_all):
     """
     Given the model IDs to sync, fetch all model objects to sync
     """
     model_objs_to_sync = {}
-    for ctype, model_ids_to_sync_for_ctype in model_ids_to_sync.items():
-        model_qset = entity_registry.entity_registry.get(ctype.model_class()).queryset
 
-        if not sync_all:
-            model_objs_to_sync[ctype] = model_qset.filter(id__in=model_ids_to_sync_for_ctype)
-        else:
-            model_objs_to_sync[ctype] = [
-                model_objs_map[ctype, model_id] for model_id in model_ids_to_sync_for_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] = [
+            model_objs_map[ctype, model_id]
+            for model_id in model_ids_to_sync_for_ctype
+        ]
 
     return model_objs_to_sync
 
@@ -231,32 +262,10 @@ def sync(self):
         # IDs of sub/super entity relationships
         super_entities_by_ctype = _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_all)
 
-        # 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
-            ids_of_fetched_models = {
-                model.id
-                for model in model_objs_by_ctype[ctype]
-            }
-
-            # Compute the set diff to see if any new records were created
-            created_model_ids = model_ids - ids_of_fetched_models
-
-            # Check if new records
-            if created_model_ids:
-
-                # Fetch the records and add them to the model_objs_map
-                new_records = ctype.model_class().objects.filter(id__in=created_model_ids)
-                for new_record in new_records:
-                    model_objs_by_ctype[ctype].append(new_record)
-                    model_objs_map[(ctype, new_record.id)] = new_record
-
         # 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, sync_all)
+        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

@@ -309,23 +309,34 @@ def test_sync_all_accounts_teams_deleted_account_during_sync(self):
         accounts[3].save()
 
         def wrapped_super_entities(*args, **kwargs):
-            if Account.objects.filter(email='[email protected]').exists():
-                Account.objects.filter(email='[email protected]').delete()
+            account = Account.objects.filter(email='[email protected]', is_active=True).first()
+            if account:
+                account.is_active = False
+                account.save()
 
             return _get_super_entities_by_ctype(*args, **kwargs)
 
-        # Sync all the entities. There should be 6 (4 accounts 2 teams)
+        # Sync the accounts
         with patch('entity.sync._get_super_entities_by_ctype', wraps=wrapped_super_entities):
             sync_entities(*accounts)
-            # Sync again to hit other wrapped function branch and make sure it doesn't error
+            account = Account.objects.get(email='[email protected]')
+            entity = Entity.all_objects.get_for_obj(account)
+            self.assertEqual(entity.is_active, False)
+
+        # Fetch accounts and sync again - hits other block in wrapped function
+        with patch('entity.sync._get_super_entities_by_ctype', wraps=wrapped_super_entities):
+            accounts = Account.objects.all()
             sync_entities(*accounts)
+            account = Account.objects.get(email='[email protected]')
+            entity = Entity.all_objects.get_for_obj(account)
+            self.assertEqual(entity.is_active, False)
 
         self.assertEquals(Entity.objects.filter(entity_type=ContentType.objects.get_for_model(Account)).count(), 4)
         self.assertEquals(Entity.objects.filter(entity_type=ContentType.objects.get_for_model(Team)).count(), 2)
         self.assertEquals(Entity.objects.all().count(), 6)
 
         # There should be six entity relationships
-        self.assertEquals(EntityRelationship.objects.all().count(), 3)
+        self.assertEquals(EntityRelationship.objects.all().count(), 4)
 
     def test_sync_all_accounts_teams_inactive_entity_kind(self):
         """

entity/version.py

@@ -1 +1 @@
-__version__ = '5.0.1'
+__version__ = '5.0.2'

release_notes.md

@@ -1,5 +1,7 @@
 ## Release Notes
 
+- 5.0.2:
+    - Only fetch new entities on a sync_all
 - 5.0.1:
     - Fix bug where an entity created or deleted during a sync would cause errors
 - 5.0.0: