Refs #36822 -- Hoisted bulk_batch_size() implementations to base backend.

Author: JaeHyuck Sa

django/db/backends/base/operations.py

@@ -2,12 +2,14 @@
 import decimal
 import json
 from importlib import import_module
+from itertools import chain
 
 import sqlparse
 
 from django.conf import settings
 from django.db import NotSupportedError, transaction
 from django.db.models.expressions import Col
+from django.db.models.fields.composite import CompositePrimaryKey
 from django.utils import timezone
 from django.utils.duration import duration_microseconds
 from django.utils.encoding import force_str
@@ -78,7 +80,17 @@ def bulk_batch_size(self, fields, objs):
         are the fields going to be inserted in the batch, the objs contains
         all the objects to be inserted.
         """
-        return len(objs)
+        if self.connection.features.max_query_params is None or not fields:
+            return len(objs)
+
+        return self.connection.features.max_query_params // len(
+            list(
+                chain.from_iterable(
+                    field.fields if isinstance(field, CompositePrimaryKey) else [field]
+                    for field in fields
+                )
+            )
+        )
 
     def format_for_duration_arithmetic(self, sql):
         raise NotImplementedError(

django/db/backends/oracle/operations.py

@@ -1,15 +1,13 @@
 import datetime
 import uuid
 from functools import lru_cache
-from itertools import chain
 
 from django.conf import settings
 from django.db import NotSupportedError
 from django.db.backends.base.operations import BaseDatabaseOperations
 from django.db.backends.utils import split_tzname_delta, strip_quotes, truncate_name
 from django.db.models import (
     AutoField,
-    CompositePrimaryKey,
     Exists,
     ExpressionWrapper,
     Lookup,
@@ -707,18 +705,6 @@ def subtract_temporals(self, internal_type, lhs, rhs):
             )
         return super().subtract_temporals(internal_type, lhs, rhs)
 
-    def bulk_batch_size(self, fields, objs):
-        """Oracle restricts the number of parameters in a query."""
-        fields = list(
-            chain.from_iterable(
-                field.fields if isinstance(field, CompositePrimaryKey) else [field]
-                for field in fields
-            )
-        )
-        if fields:
-            return self.connection.features.max_query_params // len(fields)
-        return len(objs)
-
     def conditional_expression_supported_in_where_clause(self, expression):
         """
         Oracle supports only EXISTS(...) or filters in the WHERE clause, others

django/db/backends/sqlite3/operations.py

@@ -29,26 +29,6 @@ class DatabaseOperations(BaseDatabaseOperations):
     # SQLite. Use JSON_TYPE() instead.
     jsonfield_datatype_values = frozenset(["null", "false", "true"])
 
-    def bulk_batch_size(self, fields, objs):
-        """
-        SQLite has a variable limit defined by SQLITE_LIMIT_VARIABLE_NUMBER
-        (reflected in max_query_params).
-        """
-        fields = list(
-            chain.from_iterable(
-                (
-                    field.fields
-                    if isinstance(field, models.CompositePrimaryKey)
-                    else [field]
-                )
-                for field in fields
-            )
-        )
-        if fields:
-            return self.connection.features.max_query_params // len(fields)
-        else:
-            return len(objs)
-
     def check_expression_support(self, expression):
         bad_fields = (models.DateField, models.DateTimeField, models.TimeField)
         bad_aggregates = (models.Sum, models.Avg, models.Variance, models.StdDev)

tests/backends/base/test_operations.py

@@ -1,7 +1,7 @@
 import decimal
 
 from django.core.management.color import no_style
-from django.db import NotSupportedError, connection, transaction
+from django.db import NotSupportedError, connection, models, transaction
 from django.db.backends.base.operations import BaseDatabaseOperations
 from django.db.models import DurationField
 from django.db.models.expressions import Col
@@ -11,10 +11,11 @@
     TransactionTestCase,
     override_settings,
     skipIfDBFeature,
+    skipUnlessDBFeature,
 )
 from django.utils import timezone
 
-from ..models import Author, Book
+from ..models import Author, Book, Person
 
 
 class SimpleDatabaseOperationTests(SimpleTestCase):
@@ -201,6 +202,55 @@ def test_subtract_temporals(self):
         with self.assertRaisesMessage(NotSupportedError, msg):
             self.ops.subtract_temporals(duration_field_internal_type, None, None)
 
+    @skipUnlessDBFeature("max_query_params")
+    def test_bulk_batch_size_limited(self):
+        max_query_params = connection.features.max_query_params
+        objects = range(max_query_params + 1)
+        first_name_field = Person._meta.get_field("first_name")
+        last_name_field = Person._meta.get_field("last_name")
+        composite_pk = models.CompositePrimaryKey("first_name", "last_name")
+        composite_pk.fields = [first_name_field, last_name_field]
+
+        self.assertEqual(connection.ops.bulk_batch_size([], objects), len(objects))
+        self.assertEqual(
+            connection.ops.bulk_batch_size([first_name_field], objects),
+            max_query_params,
+        )
+        self.assertEqual(
+            connection.ops.bulk_batch_size(
+                [first_name_field, last_name_field], objects
+            ),
+            max_query_params // 2,
+        )
+        self.assertEqual(
+            connection.ops.bulk_batch_size([composite_pk, first_name_field], objects),
+            max_query_params // 3,
+        )
+
+    @skipIfDBFeature("max_query_params")
+    def test_bulk_batch_size_unlimited(self):
+        objects = range(2**16 + 1)
+        first_name_field = Person._meta.get_field("first_name")
+        last_name_field = Person._meta.get_field("last_name")
+        composite_pk = models.CompositePrimaryKey("first_name", "last_name")
+        composite_pk.fields = [first_name_field, last_name_field]
+
+        self.assertEqual(connection.ops.bulk_batch_size([], objects), len(objects))
+        self.assertEqual(
+            connection.ops.bulk_batch_size([first_name_field], objects),
+            len(objects),
+        )
+        self.assertEqual(
+            connection.ops.bulk_batch_size(
+                [first_name_field, last_name_field], objects
+            ),
+            len(objects),
+        )
+        self.assertEqual(
+            connection.ops.bulk_batch_size([composite_pk, first_name_field], objects),
+            len(objects),
+        )
+
 
 class SqlFlushTests(TransactionTestCase):
     available_apps = ["backends"]

tests/backends/oracle/test_operations.py

@@ -1,7 +1,7 @@
 import unittest
 
 from django.core.management.color import no_style
-from django.db import connection, models
+from django.db import connection
 from django.test import TransactionTestCase
 
 from ..models import Person, Tag
@@ -17,31 +17,6 @@ def test_sequence_name_truncation(self):
         )
         self.assertEqual(seq_name, "SCHEMA_AUTHORWITHEVENLOB0B8_SQ")
 
-    def test_bulk_batch_size(self):
-        # Oracle restricts the number of parameters in a query.
-        objects = range(2**16)
-        self.assertEqual(connection.ops.bulk_batch_size([], objects), len(objects))
-        # Each field is a parameter for each object.
-        first_name_field = Person._meta.get_field("first_name")
-        last_name_field = Person._meta.get_field("last_name")
-        self.assertEqual(
-            connection.ops.bulk_batch_size([first_name_field], objects),
-            connection.features.max_query_params,
-        )
-        self.assertEqual(
-            connection.ops.bulk_batch_size(
-                [first_name_field, last_name_field],
-                objects,
-            ),
-            connection.features.max_query_params // 2,
-        )
-        composite_pk = models.CompositePrimaryKey("first_name", "last_name")
-        composite_pk.fields = [first_name_field, last_name_field]
-        self.assertEqual(
-            connection.ops.bulk_batch_size([composite_pk, first_name_field], objects),
-            connection.features.max_query_params // 3,
-        )
-
     def test_sql_flush(self):
         statements = connection.ops.sql_flush(
             no_style(),

tests/backends/sqlite/test_operations.py

@@ -2,7 +2,7 @@
 import unittest
 
 from django.core.management.color import no_style
-from django.db import connection, models
+from django.db import connection
 from django.test import TestCase
 
 from ..models import Person, Tag
@@ -88,29 +88,6 @@ def test_sql_flush_sequences_allow_cascade(self):
             statements[-1],
         )
 
-    def test_bulk_batch_size(self):
-        self.assertEqual(connection.ops.bulk_batch_size([], [Person()]), 1)
-        first_name_field = Person._meta.get_field("first_name")
-        last_name_field = Person._meta.get_field("last_name")
-        self.assertEqual(
-            connection.ops.bulk_batch_size([first_name_field], [Person()]),
-            connection.features.max_query_params,
-        )
-        self.assertEqual(
-            connection.ops.bulk_batch_size(
-                [first_name_field, last_name_field], [Person()]
-            ),
-            connection.features.max_query_params // 2,
-        )
-        composite_pk = models.CompositePrimaryKey("first_name", "last_name")
-        composite_pk.fields = [first_name_field, last_name_field]
-        self.assertEqual(
-            connection.ops.bulk_batch_size(
-                [composite_pk, first_name_field], [Person()]
-            ),
-            connection.features.max_query_params // 3,
-        )
-
     def test_bulk_batch_size_respects_variable_limit(self):
         first_name_field = Person._meta.get_field("first_name")
         last_name_field = Person._meta.get_field("last_name")

tests/composite_pk/tests.py

@@ -149,21 +149,18 @@ def test_in_bulk(self):
 
     def test_in_bulk_batching(self):
         Comment.objects.all().delete()
-        batching_required = connection.features.max_query_params is not None
-        expected_queries = 2 if batching_required else 1
+        num_objects = 10
+        connection.features.__dict__.pop("max_query_params", None)
         with unittest.mock.patch.object(
-            type(connection.features), "max_query_params", 10
+            type(connection.features), "max_query_params", num_objects
         ):
-            num_requiring_batching = (
-                connection.ops.bulk_batch_size([Comment._meta.pk], []) + 1
-            )
             comments = [
                 Comment(id=i, tenant=self.tenant, user=self.user)
-                for i in range(1, num_requiring_batching + 1)
+                for i in range(1, num_objects + 1)
             ]
             Comment.objects.bulk_create(comments)
             id_list = list(Comment.objects.values_list("pk", flat=True))
-            with self.assertNumQueries(expected_queries):
+            with self.assertNumQueries(2):
                 comment_dict = Comment.objects.in_bulk(id_list=id_list)
         self.assertQuerySetEqual(comment_dict, id_list)