added an entity model manager that sends out signals when bulk operations occur. Entities are resynced when bulk operations occur and it is tested

Author: Wes Kendall

entity/__init__.py

@@ -1,3 +1,3 @@
 from .sync import sync_entities
 
-from .models import EntityModelMixin, Entity
+from .models import EntityModelMixin, Entity, BaseEntityModel

entity/models.py

@@ -1,11 +1,16 @@
 from django.contrib.contenttypes import generic
 from django.contrib.contenttypes.models import ContentType
 from django.db import models
+from django.db.models.query import QuerySet
 from django.db.models.signals import post_save, post_delete
-from django.dispatch import receiver
+from django.dispatch import receiver, Signal
 from jsonfield import JSONField
 
 
+# Create a signal that is emitted after bulk operations occur
+post_bulk_operation = Signal()
+
+
 class Entity(models.Model):
     """
     Describes an entity and its relevant metadata. Also defines
@@ -177,6 +182,41 @@ def is_super_entity_relationship_active(self, model_obj):
         return True
 
 
+class EntityQuerySet(QuerySet):
+    """
+    Overrides bulk operations on a queryset to emit a signal when they occur.
+    """
+    def update(self, **kwargs):
+        ret_val = super(EntityQuerySet, self).update(**kwargs)
+        post_bulk_operation.send(sender=self)
+        return ret_val
+
+
+class EntityModelManager(models.Manager):
+    """
+    Defines a model manager for Entity models. This model manager provides additional
+    functionality on top of the regular managers, such as emitting signals when bulk
+    updates and creates are issued.
+    """
+    def get_queryset(self):
+        return EntityQuerySet(self.model)
+
+    def bulk_create(self, objs, batch_size=None):
+        ret_val = super(EntityModelManager, self).bulk_create(objs, batch_size=batch_size)
+        post_bulk_operation.send(sender=self)
+        return ret_val
+
+
+class BaseEntityModel(models.Model, EntityModelMixin):
+    """
+    Defines base properties for an Entity model defined by a third-party application.
+    """
+    class Meta:
+        abstract = True
+
+    objects = EntityModelManager()
+
+
 def sync_entity_signal_handler(sender, model_obj, is_deleted):
     """
     Filters post save/delete signals for entities by checking if they
@@ -189,6 +229,16 @@ def sync_entity_signal_handler(sender, model_obj, is_deleted):
         sync_entity(model_obj, is_deleted)
 
 
+def sync_entities_signal_handler(sender):
+    """
+    When a bulk operation occurs on a model manager, sync all the entities
+    if the model of the manager is an entity class.
+    """
+    if issubclass(sender.model, EntityModelMixin):
+        from .sync import sync_entities
+        sync_entities()
+
+
 @receiver(post_delete)
 def delete_entity_signal_handler(sender, *args, **kwargs):
     """
@@ -205,3 +255,11 @@ def save_entity_signal_handler(sender, *args, **kwargs):
     the entity mirror table.
     """
     sync_entity_signal_handler(sender, kwargs['instance'], False)
+
+
+@receiver(post_bulk_operation)
+def bulk_operation_signal_handler(sender, *args, **kwargs):
+    """
+    When a bulk operation has happened on a model, sync all the entities again.
+    """
+    sync_entities_signal_handler(sender)

test_project/models.py

@@ -1,16 +1,17 @@
 from django.db import models
 
-from entity import EntityModelMixin, Entity
+from entity import BaseEntityModel, Entity
+from entity.models import EntityModelManager
 
 
-class TeamGroup(models.Model, EntityModelMixin):
+class TeamGroup(BaseEntityModel):
     """
     A grouping of teams.
     """
     name = models.CharField(max_length=256)
 
 
-class Team(models.Model, EntityModelMixin):
+class Team(BaseEntityModel):
     """
     A team entity model. Encapsulates accounts.
     """
@@ -27,7 +28,7 @@ def get_super_entities(self):
         return [self.team_group] if self.team_group is not None else []
 
 
-class Account(models.Model, EntityModelMixin):
+class Account(BaseEntityModel):
     """
     An account entity model
     """
@@ -82,8 +83,10 @@ class DummyModel(models.Model):
     """
     dummy_data = models.CharField(max_length=64)
 
+    objects = EntityModelManager()
 
-class BaseEntityClass(models.Model, EntityModelMixin):
+
+class BaseEntityClass(BaseEntityModel):
     """
     A base class that inherits EntityModelMixin. Helps ensure that mutliple-inherited
     entities are still synced properly.

test_project/tests/sync_tests.py

@@ -107,6 +107,47 @@ class TestEntitySignalSync(EntityTestCase):
     Tests that entities (from the test models) are properly synced upon post_save and
     post_delete calls.
     """
+    def test_post_bulk_create(self):
+        """
+        Tests that entities can have bulk creates applied to them and still be synced.
+        """
+        # Bulk create five accounts
+        accounts = [Account() for i in range(5)]
+        Account.objects.bulk_create(accounts)
+        # Verify that there are 5 entities
+        self.assertEquals(Entity.objects.all().count(), 5)
+
+    def test_post_bulk_update(self):
+        """
+        Calls a bulk update on a list of entities. Verifies that the models are appropriately
+        synced.
+        """
+        # Create five accounts
+        for i in range(5):
+            Account.objects.create(email='[email protected]')
+        # Verify that there are five entities all with the '[email protected]' email
+        for entity in Entity.objects.all():
+            self.assertEquals(entity.entity_meta['email'], '[email protected]')
+        self.assertEquals(Entity.objects.all().count(), 5)
+
+        # Bulk update the account emails to a different one
+        Account.objects.all().update(email='[email protected]')
+
+        # Verify that the email was updated properly in all entities
+        for entity in Entity.objects.all():
+            self.assertEquals(entity.entity_meta['email'], '[email protected]')
+        self.assertEquals(Entity.objects.all().count(), 5)
+
+    def test_post_bulk_update_dummy(self):
+        """
+        Tests that even if the dummy model is using the special model manager for bulk
+        updates, it still does not get synced since it doesn't inherit EntityModelMixin.
+        """
+        # Create five dummy models with a bulk update
+        DummyModel.objects.bulk_create([DummyModel() for i in range(5)])
+        # There should be no synced entities
+        self.assertEquals(Entity.objects.all().count(), 0)
+
     def test_post_save_dummy_data(self):
         """
         Tests that dummy data that does not inherit from EntityModelMixin is not synced