Ticket #17788: batch_bulk_insert.diff

django/db/backends/__init__.py

@@ -858 +858 @@
         conn = ' %s ' % connector
         return conn.join(sub_expressions)
 
+    def get_batch_size(self, fields, objs):
+        """
+        Returns the maximum allowed batch size for the backend. The fields
+        will be repeated for each object in the batch. You can use the fields
+        together with the objs to approximate the size of the query.
+
+        Use an arbitrarily picked value of 100000 for default batch size.
+        Above this amount backends will have problems with either query string
+        size, or they will use excessive amount of memory.
+        """
+        return 100000
+
 class BaseDatabaseIntrospection(object):
     """
     This class encapsulates all backend-specific introspection utilities

django/db/backends/sqlite3/base.py

@@ -211 +211 @@
         res.extend(["UNION SELECT %s" % ", ".join(["%s"] * len(fields))] * (num_values - 1))
         return " ".join(res)
 
+    def get_batch_size(self, fields, objs):
+        """
+        SQLite has a limit of 1000 variables per query. Intentionally
+        leave a couple of unused variables.
+        """
+        return max(900 / len(fields), 1)
+
 class DatabaseWrapper(BaseDatabaseWrapper):
     vendor = 'sqlite'
     # SQLite requires LIKE statements to include an ESCAPE clause if the value

django/db/models/query.py

@@ -386 +386 @@
         """
         # So this case is fun. When you bulk insert you don't get the primary
         # keys back (if it's an autoincrement), so you can't insert into the
-        # child tables which references this. There are two workarounds, 1)
+        # child tables which references this. There are three workarounds, 1)
         # this could be implemented if you didn't have an autoincrement pk,
         # and 2) you could do it by doing O(n) normal inserts into the parent
         # tables to get the primary keys back, and then doing a single bulk
-        # insert into the childmost table. We're punting on these for now
-        # because they are relatively rare cases.
+        # insert into the childmost table. 3) you could do it if the database
+        # implements RETURNING id queries (Oracle, PostgreSQL).  We're punting
+        # on these for now because they are relatively rare cases. Still,
+        # fixing this would allow "generic" use of bulk create.
         if self.model._meta.parents:
             raise ValueError("Can't bulk create an inherited model")
         if not objs:
@@ -399 +401 @@
         self._for_write = True
         connection = connections[self.db]
         fields = self.model._meta.local_fields
+        def batched_insert(objs, fields, using):
+            if not objs:
+                return
+            batch_size = connection.ops.get_batch_size(fields, objs)
+            self.model._base_manager._insert(objs[0:batch_size], fields=fields, using=self.db)
+            batched_insert(objs[batch_size:], fields, using)
+
         if not transaction.is_managed(using=self.db):
             transaction.enter_transaction_management(using=self.db)
             forced_managed = True
@@ -407 +416 @@
         try:
             if (connection.features.can_combine_inserts_with_and_without_auto_increment_pk
                 and self.model._meta.has_auto_field):
-                self.model._base_manager._insert(objs, fields=fields, using=self.db)
+                batched_insert(objs, fields, self.db)
             else:
                 objs_with_pk, objs_without_pk = partition(lambda o: o.pk is None, objs)
                 if objs_with_pk:
-                    self.model._base_manager._insert(objs_with_pk, fields=fields, using=self.db)
+                    batched_insert(objs_with_pk, fields, self.db)
                 if objs_without_pk:
-                    self.model._base_manager._insert(objs_without_pk, fields=[f for f in fields if not isinstance(f, AutoField)], using=self.db)
+                    batched_insert(objs_without_pk, [f for f in fields if not isinstance(f, AutoField)],
+                                   self.db)
             if forced_managed:
                 transaction.commit(using=self.db)
             else:

tests/regressiontests/bulk_create/models.py

@@ -18 +18 @@
     pass
 
 class State(models.Model):
-    two_letter_code = models.CharField(max_length=2, primary_key=True)
- No newline at end of file
+    two_letter_code = models.CharField(max_length=2, primary_key=True)
+
+class TwoFields(models.Model):
+    f1 = models.CharField(max_length=100)
+    f2 = models.CharField(max_length=100)

tests/regressiontests/bulk_create/tests.py

@@ -2 +2 @@
 
 from operator import attrgetter
 
-from django.test import TestCase, skipUnlessDBFeature
+from django.db import connection
+from django.test import TestCase, skipUnlessDBFeature, skipIfDBFeature
+from django.test.utils import override_settings
 
-from .models import Country, Restaurant, Pizzeria, State
+from .models import Country, Restaurant, Pizzeria, State, TwoFields
 
 
 class BulkCreateTests(TestCase):
@@ -56 +58 @@
             ])
         self.assertQuerysetEqual(State.objects.order_by("two_letter_code"), [
             "CA", "IL", "ME", "NY",
-        ], attrgetter("two_letter_code"))
- No newline at end of file
+        ], attrgetter("two_letter_code"))
+
+    def test_large_batch(self):
+        with override_settings(DEBUG=True):
+            Restaurant.objects.bulk_create([
+                   Restaurant(name=str(i)) for i in range(0, 1001)
+                ])
+            self.assertTrue(len(connection.queries) < 10)
+        self.assertEqual(Restaurant.objects.count(), 1001)
+        self.assertQuerysetEqual(
+            Restaurant.objects.filter(name__in=['999', '1000', '1001']).order_by('name'),
+            ['1000', '999'], attrgetter('name'))
+        # We happen to know SQLite is going to do the cut at 900 objs. Test that boundary.
+        self.assertQuerysetEqual(
+            Restaurant.objects.filter(name__in=['899', '900', '901']).order_by('name'),
+            ['899', '900', '901'], attrgetter('name'))
+
+    def test_large_batch_multifield(self):
+        with override_settings(DEBUG=True):
+            TwoFields.objects.bulk_create([
+                   TwoFields(f1=str(i), f2=str(i)) for i in range(0, 1001)
+                ])
+            self.assertTrue(len(connection.queries) < 10)
+        self.assertEqual(TwoFields.objects.count(), 1001)
+        # Test boundaries
+        self.assertQuerysetEqual(
+            TwoFields.objects.filter(f1__in=['0', '1']).order_by('f1'),
+            ['0', '1'], attrgetter('f1'))
+        self.assertQuerysetEqual(
+            TwoFields.objects.filter(f1__in=['500', '501']).order_by('f1'),
+            ['500', '501'], attrgetter('f1'))
+        self.assertQuerysetEqual(
+            TwoFields.objects.filter(f1__in=['999', '1000', '1001']).order_by('f1'),
+            ['1000', '999'], attrgetter('f1'))
+
+    @skipIfDBFeature('can_combine_inserts_with_and_without_auto_increment_pk')
+    def test_lage_batch_mixed(self):
+        # mixed pk + non-pk.
+        with override_settings(DEBUG=True):
+            TwoFields.objects.bulk_create([
+                   TwoFields(id=i if i % 2 == 0 else None, f1=str(i), f2=str(i)) for i in range(1000, 2001)
+                ])
+            self.assertTrue(len(connection.queries) < 10)
+        self.assertEqual(TwoFields.objects.count(), 1001)
+        self.assertQuerysetEqual(TwoFields.objects.filter(id__in=[1998, 1999, 2000]).order_by('id'), [
+            1998, 2000
+        ], attrgetter('id'))

tests/regressiontests/queries/tests.py

@@ -1807 +1807 @@
         # Test that the "in" lookup works with lists of 1000 items or more.
         Number.objects.all().delete()
         numbers = range(2500)
-        for num in numbers:
-            _ = Number.objects.create(num=num)
+        Number.objects.bulk_create(Number(num=num) for num in numbers)
         self.assertEqual(
             Number.objects.filter(num__in=numbers[:1000]).count(),
             1000