Removed n+1 query from bulk_create_with_history utility (#975)

Author: twolfson

AUTHORS.rst

@@ -113,6 +113,7 @@ Authors
 - Steven Buss (`sbuss <https://github.com/sbuss>`_)
 - Steven Klass
 - Tim Schilling (`tim-schilling <https://github.com/tim-schilling>`_)
+- Todd Wolfson (`twolfson <https://github.com/twolfson>`_)
 - Tommy Beadle (`tbeadle <https://github.com/tbeadle>`_)
 - Trey Hunner (`treyhunner <https://github.com/treyhunner>`_)
 - Ulysses Vilela

CHANGES.rst

@@ -4,6 +4,7 @@ Changes
 Unreleased
 ----------
 
+- Removed n+1 query from ``bulk_create_with_history`` utility (gh-975)
 - Started using ``exists`` query instead of ``count`` in ``populate_history`` command (gh-982)
 
 3.1.1 (2022-04-23)

simple_history/tests/tests/test_utils.py

@@ -186,6 +186,38 @@ def test_bulk_create_history_with_duplicates_ignore_conflicts(self):
         self.assertEqual(PollWithUniqueQuestion.objects.count(), 2)
         self.assertEqual(PollWithUniqueQuestion.history.count(), 2)
 
+    def test_bulk_create_history_with_no_ids_return(self):
+        pub_date = timezone.now()
+        objects = [
+            Poll(question="Question 1", pub_date=pub_date),
+            Poll(question="Question 2", pub_date=pub_date),
+            Poll(question="Question 3", pub_date=pub_date),
+            Poll(question="Question 4", pub_date=pub_date),
+            Poll(question="Question 5", pub_date=pub_date),
+        ]
+
+        _bulk_create = Poll._default_manager.bulk_create
+
+        def mock_bulk_create(*args, **kwargs):
+            _bulk_create(*args, **kwargs)
+            return [
+                Poll(question="Question 1", pub_date=pub_date),
+                Poll(question="Question 2", pub_date=pub_date),
+                Poll(question="Question 3", pub_date=pub_date),
+                Poll(question="Question 4", pub_date=pub_date),
+                Poll(question="Question 5", pub_date=pub_date),
+            ]
+
+        with patch.object(
+            Poll._default_manager, "bulk_create", side_effect=mock_bulk_create
+        ):
+            with self.assertNumQueries(3):
+                result = bulk_create_with_history(objects, Poll)
+            self.assertEqual(
+                [poll.question for poll in result], [poll.question for poll in objects]
+            )
+            self.assertNotEqual(result[0].id, None)
+
 
 class BulkCreateWithHistoryTransactionTestCase(TransactionTestCase):
     def setUp(self):
@@ -242,7 +274,7 @@ def test_bulk_create_no_ids_return(self, hist_manager_mock):
         model = Mock(
             _default_manager=Mock(
                 bulk_create=Mock(return_value=[Place(name="Place 1")]),
-                filter=Mock(return_value=objects),
+                filter=Mock(return_value=Mock(order_by=Mock(return_value=objects))),
             ),
             _meta=Mock(get_fields=Mock(return_value=[])),
         )

simple_history/utils.py

@@ -2,7 +2,7 @@
 
 import django
 from django.db import transaction
-from django.db.models import ForeignKey, ManyToManyField
+from django.db.models import Case, ForeignKey, ManyToManyField, Q, When
 from django.forms.models import model_to_dict
 
 from simple_history.exceptions import AlternativeManagerError, NotHistoricalModelError
@@ -111,16 +111,35 @@ def bulk_create_with_history(
                 default_date=default_date,
             )
     if second_transaction_required:
-        obj_list = []
         with transaction.atomic(savepoint=False):
-            for obj in objs_with_id:
+            # Generate a common query to avoid n+1 selections
+            #   https://github.com/jazzband/django-simple-history/issues/974
+            cumulative_filter = None
+            obj_when_list = []
+            for i, obj in enumerate(objs_with_id):
                 attributes = dict(
                     filter(
                         lambda x: x[1] is not None,
                         model_to_dict(obj, exclude=exclude_fields).items(),
                     )
                 )
-                obj_list += model_manager.filter(**attributes)
+                q = Q(**attributes)
+                cumulative_filter = (cumulative_filter | q) if cumulative_filter else q
+                # https://stackoverflow.com/a/49625179/1960509
+                # DEV: If an attribute has `then` as a key
+                #   then they'll also run into issues with `bulk_update`
+                #   due to shared implementation
+                #   https://github.com/django/django/blob/4.0.4/django/db/models/query.py#L624-L638
+                obj_when_list.append(When(**attributes, then=i))
+            obj_list = (
+                list(
+                    model_manager.filter(cumulative_filter).order_by(
+                        Case(*obj_when_list)
+                    )
+                )
+                if objs_with_id
+                else []
+            )
             history_manager.bulk_history_create(
                 obj_list,
                 batch_size=batch_size,