Merge pull request #16 from ambitioninc/develop

bulk_upsert operates on models rather than dictionaries
ambitioninc · Apr 7, 2014 · 9197529 · 9197529
2 parents fdb6428 + c9bfe0b
commit 9197529
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 66 deletions.
diff --git a/README.md b/README.md
@@ -103,8 +103,8 @@ Performs an update on an object or an insert if the object does not exist.
     print model_obj.int_field, model_obj.float_field
     2, 4.0
 
-## bulk_upsert(objs, unique_fields, update_fields)<a name="bulk_upsert"></a>
-Performs a bulk update or insert on a list of dictionaries. Matches all objects in the queryset with the objs provided using the field values in unique_fields. If an existing object is matched, it is updated with the values from the provided objects. Objects that don't match anything are bulk inserted.
+## bulk_upsert(model_objs, unique_fields, update_fields)<a name="bulk_upsert"></a>
+Performs a bulk update or insert on a list of model objects. Matches all objects in the queryset with the objs provided using the field values in unique_fields. If an existing object is matched, it is updated with the values from the provided objects. Objects that don't match anything are bulk inserted.
 
 **Args**:
 - objs: A list of dictionaries that have fields corresponding to the model in the manager.
@@ -118,9 +118,9 @@ Performs a bulk update or insert on a list of dictionaries. Matches all objects
     # Start off with no objects in the database. Call a bulk_upsert on the TestModel, which includes
     # a char_field, int_field, and float_field
     TestModel.objects.bulk_upsert([
-        {'float_field': 1.0, 'char_field': '1', 'int_field': 1},
-        {'float_field': 2.0, 'char_field': '2', 'int_field': 2},
-        {'float_field': 3.0, 'char_field': '3', 'int_field': 3},
+        TestModel(float_field=1.0, char_field='1', int_field=1),
+        TestModel(float_field=2.0, char_field='2', int_field=2),
+        TestModel(float_field=3.0, char_field='3', int_field=3),
     ], ['int_field'], ['char_field'])
 
     # All objects should have been created
@@ -130,9 +130,9 @@ Performs a bulk update or insert on a list of dictionaries. Matches all objects
     # Now perform a bulk upsert on all the char_field values. Since the objects existed previously
     # (known by the int_field uniqueness constraint), the char fields should be updated
     TestModel.objects.bulk_upsert([
-        {'float_field': 1.0, 'char_field': '0', 'int_field': 1},
-        {'float_field': 2.0, 'char_field': '0', 'int_field': 2},
-        {'float_field': 3.0, 'char_field': '0', 'int_field': 3},
+        TestModel(float_field=1.0, char_field='0', int_field=1),
+        TestModel(float_field=2.0, char_field='0', int_field=2),
+        TestModel(float_field=3.0, char_field='0', int_field=3),
     ], ['int_field'], ['char_field'])
 
     # No more new objects should have been created, and every char field should be 0
@@ -142,10 +142,10 @@ Performs a bulk update or insert on a list of dictionaries. Matches all objects
     # Do the exact same operation, but this time add an additional object that is not already
     # stored. It will be inserted.
     TestModel.objects.bulk_upsert([
-        {'float_field': 1.0, 'char_field': '1', 'int_field': 1},
-        {'float_field': 2.0, 'char_field': '2', 'int_field': 2},
-        {'float_field': 3.0, 'char_field': '3', 'int_field': 3},
-        {'float_field': 4.0, 'char_field': '4', 'int_field': 4},
+        TestModel(float_field=1.0, char_field='1', int_field=1),
+        TestModel(float_field=2.0, char_field='2', int_field=2),
+        TestModel(float_field=3.0, char_field='3', int_field=3),
+        TestModel(float_field=4.0, char_field='4', int_field=4),
     ], ['int_field'], ['char_field'])
 
     # There should be one more object
@@ -156,10 +156,10 @@ Performs a bulk update or insert on a list of dictionaries. Matches all objects
     # filter for int_field=1. In this case, only one object has the ability to be updated.
     # All of the other objects will be inserted
     TestModel.objects.filter(int_field=1).bulk_upsert([
-        {'float_field': 1.0, 'char_field': '1', 'int_field': 1},
-        {'float_field': 2.0, 'char_field': '2', 'int_field': 2},
-        {'float_field': 3.0, 'char_field': '3', 'int_field': 3},
-        {'float_field': 4.0, 'char_field': '4', 'int_field': 4},
+        TestModel(float_field=1.0, char_field='1', int_field=1),
+        TestModel(float_field=2.0, char_field='2', int_field=2),
+        TestModel(float_field=3.0, char_field='3', int_field=3),
+        TestModel(float_field=4.0, char_field='4', int_field=4),
     ], ['int_field'], ['char_field'])
 
     # There should be three more objects

diff --git a/manager_utils/VERSION b/manager_utils/VERSION
@@ -1 +1 @@
-0.3.5
+0.3.6
diff --git a/manager_utils/manager_utils.py b/manager_utils/manager_utils.py
@@ -20,7 +20,7 @@ def id_dict(self):
         """
         return {obj.id: obj for obj in self}
 
-    def bulk_upsert(self, objs, unique_fields, update_fields):
+    def bulk_upsert(self, model_objs, unique_fields, update_fields):
         """
         Performs a bulk update or insert on a queryset.
         """
@@ -30,22 +30,23 @@ def bulk_upsert(self, objs, unique_fields, update_fields):
             raise ValueError('Must provide update_fields argument')
 
         # Create a look up table for all of the objects in the queryset keyed on the unique_fields
-        model_obj_dict = {
-            tuple(getattr(model_obj, field) for field in unique_fields): model_obj for model_obj in self
+        extant_model_objs = {
+            tuple(getattr(extant_model_obj, field) for field in unique_fields): extant_model_obj
+            for extant_model_obj in self
         }
 
         # Find all of the objects to update and all of the objects to create
         model_objs_to_update, model_objs_to_create = [], []
-        for obj in objs:
-            model_obj = model_obj_dict.get(tuple(obj[field] for field in unique_fields), None)
-            if model_obj is None:
+        for model_obj in model_objs:
+            extant_model_obj = extant_model_objs.get(tuple(getattr(model_obj, field) for field in unique_fields), None)
+            if extant_model_obj is None:
                 # If the object needs to be created, make a new instance of it
-                model_objs_to_create.append(self.model(**obj))
+                model_objs_to_create.append(model_obj)
             else:
                 # If the object needs to be updated, update its fields
                 for field in update_fields:
-                    setattr(model_obj, field, obj[field])
-                model_objs_to_update.append(model_obj)
+                    setattr(extant_model_obj, field, getattr(model_obj, field))
+                model_objs_to_update.append(extant_model_obj)
 
         # Apply bulk updates and creates
         self.model.objects.bulk_update(model_objs_to_update, update_fields)
@@ -104,9 +105,9 @@ def id_dict(self):
         """
         return self.get_queryset().id_dict()
 
-    def bulk_upsert(self, objs, unique_fields, update_fields):
+    def bulk_upsert(self, model_objs, unique_fields, update_fields):
         """
-        Performs a bulk update or insert on a list of dictionaries. Matches all objects in the queryset
+        Performs a bulk update or insert on a list of model objects. Matches all objects in the queryset
         with the objs provided using the field values in unique_fields.
         If an existing object is matched, it is updated with the values from the provided objects. Objects
         that don't match anything are bulk inserted.
@@ -124,9 +125,9 @@ def bulk_upsert(self, objs, unique_fields, update_fields):
             # Start off with no objects in the database. Call a bulk_upsert on the TestModel, which includes
             # a char_field, int_field, and float_field
             TestModel.objects.bulk_upsert([
-                {'float_field': 1.0, 'char_field': '1', 'int_field': 1},
-                {'float_field': 2.0, 'char_field': '2', 'int_field': 2},
-                {'float_field': 3.0, 'char_field': '3', 'int_field': 3},
+                TestModel(float_field=1.0, char_field='1', int_field=1),
+                TestModel(float_field=2.0, char_field='2', int_field=2),
+                TestModel(float_field=3.0, char_field='3', int_field=3),
             ], ['int_field'], ['char_field'])
 
             # All objects should have been created
@@ -136,9 +137,9 @@ def bulk_upsert(self, objs, unique_fields, update_fields):
             # Now perform a bulk upsert on all the char_field values. Since the objects existed previously
             # (known by the int_field uniqueness constraint), the char fields should be updated
             TestModel.objects.bulk_upsert([
-                {'float_field': 1.0, 'char_field': '0', 'int_field': 1},
-                {'float_field': 2.0, 'char_field': '0', 'int_field': 2},
-                {'float_field': 3.0, 'char_field': '0', 'int_field': 3},
+                TestModel(float_field=1.0, char_field='0', int_field=1),
+                TestModel(float_field=2.0, char_field='0', int_field=2),
+                TestModel(float_field=3.0, char_field='0', int_field=3),
             ], ['int_field'], ['char_field'])
 
             # No more new objects should have been created, and every char field should be 0
@@ -148,10 +149,10 @@ def bulk_upsert(self, objs, unique_fields, update_fields):
             # Do the exact same operation, but this time add an additional object that is not already
             # stored. It will be inserted.
             TestModel.objects.bulk_upsert([
-                {'float_field': 1.0, 'char_field': '1', 'int_field': 1},
-                {'float_field': 2.0, 'char_field': '2', 'int_field': 2},
-                {'float_field': 3.0, 'char_field': '3', 'int_field': 3},
-                {'float_field': 4.0, 'char_field': '4', 'int_field': 4},
+                TestModel(float_field=1.0, char_field='1', int_field=1),
+                TestModel(float_field=2.0, char_field='2', int_field=2),
+                TestModel(float_field=3.0, char_field='3', int_field=3),
+                TestModel(float_field=4.0, char_field='4', int_field=4),
             ], ['int_field'], ['char_field'])
 
             # There should be one more object
@@ -162,24 +163,24 @@ def bulk_upsert(self, objs, unique_fields, update_fields):
             # filter for int_field=1. In this case, only one object has the ability to be updated.
             # All of the other objects will be inserted
             TestModel.objects.filter(int_field=1).bulk_upsert([
-                {'float_field': 1.0, 'char_field': '1', 'int_field': 1},
-                {'float_field': 2.0, 'char_field': '2', 'int_field': 2},
-                {'float_field': 3.0, 'char_field': '3', 'int_field': 3},
-                {'float_field': 4.0, 'char_field': '4', 'int_field': 4},
+                TestModel(float_field=1.0, char_field='1', int_field=1),
+                TestModel(float_field=2.0, char_field='2', int_field=2),
+                TestModel(float_field=3.0, char_field='3', int_field=3),
+                TestModel(float_field=4.0, char_field='4', int_field=4),
             ], ['int_field'], ['char_field'])
 
             # There should be three more objects
             print TestModel.objects.count()
             7
         """
-        return self.get_queryset().bulk_upsert(objs, unique_fields, update_fields)
+        return self.get_queryset().bulk_upsert(model_objs, unique_fields, update_fields)
 
-    def bulk_create(self, objs, batch_size=None):
+    def bulk_create(self, model_objs, batch_size=None):
         """
         Overrides Django's bulk_create function to emit a post_bulk_operation signal when bulk_create
         is finished.
         """
-        ret_val = super(ManagerUtilsMixin, self).bulk_create(objs, batch_size=batch_size)
+        ret_val = super(ManagerUtilsMixin, self).bulk_create(model_objs, batch_size=batch_size)
         post_bulk_operation.send(sender=self, model=self.model)
         return ret_val
 

diff --git a/test_project/tests/manager_utils_tests.py b/test_project/tests/manager_utils_tests.py
@@ -37,9 +37,9 @@ def test_no_updates(self):
         Tests the case when no updates were previously stored (i.e objects are only created)
         """
         TestModel.objects.bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field'], ['char_field', 'float_field'])
 
         for i, model_obj in enumerate(TestModel.objects.order_by('int_field')):
@@ -58,9 +58,9 @@ def test_all_updates_unique_int_field(self):
 
         # Update using the int field as a uniqueness constraint
         TestModel.objects.bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field'], ['char_field', 'float_field'])
 
         # Verify that the fields were updated
@@ -81,9 +81,9 @@ def test_all_updates_unique_int_field_update_float_field(self):
 
         # Update using the int field as a uniqueness constraint
         TestModel.objects.bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field'], update_fields=['float_field'])
 
         # Verify that the float field was updated
@@ -104,9 +104,9 @@ def test_some_updates_unique_int_field_update_float_field(self):
 
         # Update using the int field as a uniqueness constraint. The first two are updated while the third is created
         TestModel.objects.bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field'], ['float_field'])
 
         # Verify that the float field was updated for the first two models and the char field was not updated for
@@ -128,9 +128,9 @@ def test_some_updates_unique_int_char_field_update_float_field(self):
 
         # Update using the int field as a uniqueness constraint. The first two are updated while the third is created
         TestModel.objects.bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field', 'char_field'], ['float_field'])
 
         # Verify that the float field was updated for the first two models and the char field was not updated for
@@ -152,9 +152,9 @@ def test_no_updates_unique_int_char_field(self):
 
         # Update using the int and char field as a uniqueness constraint. All three objects are created
         TestModel.objects.bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field', 'char_field'], ['float_field'])
 
         # Verify that no updates occured
@@ -180,9 +180,9 @@ def test_some_updates_unique_int_char_field_queryset(self):
 
         # Update using the int field as a uniqueness constraint on a queryset. Only one object should be updated.
         TestModel.objects.filter(int_field=0).bulk_upsert([
-            {'int_field': 0, 'char_field': '0', 'float_field': 0},
-            {'int_field': 1, 'char_field': '1', 'float_field': 1},
-            {'int_field': 2, 'char_field': '2', 'float_field': 2},
+            TestModel(int_field=0, char_field='0', float_field=0),
+            TestModel(int_field=1, char_field='1', float_field=1),
+            TestModel(int_field=2, char_field='2', float_field=2),
         ], ['int_field'], ['float_field'])
 
         # Verify that two new objecs were inserted