Fixed #18556 -- Allowed RelatedManager.add() to execute 1 query where possible.

Thanks Loic Bistuer for review.

8 files changed, 151 insertions, 28 deletions

diff --git a/django/contrib/contenttypes/fields.py b/django/contrib/contenttypes/fields.py
index 6548c39b4b..0980ed21d1 100644
--- a/django/contrib/contenttypes/fields.py
+++ b/django/contrib/contenttypes/fields.py
@@ -501,15 +501,38 @@ def create_generic_related_manager(superclass, rel):
                     False,
                     self.prefetch_cache_name)
 
-        def add(self, *objs):
+        def add(self, *objs, **kwargs):
+            bulk = kwargs.pop('bulk', True)
             db = router.db_for_write(self.model, instance=self.instance)
-            with transaction.atomic(using=db, savepoint=False):
+
+            def check_and_update_obj(obj):
+                if not isinstance(obj, self.model):
+                    raise TypeError("'%s' instance expected, got %r" % (
+                        self.model._meta.object_name, obj
+                    ))
+                setattr(obj, self.content_type_field_name, self.content_type)
+                setattr(obj, self.object_id_field_name, self.pk_val)
+
+            if bulk:
+                pks = []
                 for obj in objs:
-                    if not isinstance(obj, self.model):
-                        raise TypeError("'%s' instance expected" % self.model._meta.object_name)
-                    setattr(obj, self.content_type_field_name, self.content_type)
-                    setattr(obj, self.object_id_field_name, self.pk_val)
-                    obj.save()
+                    if obj._state.adding or obj._state.db != db:
+                        raise ValueError(
+                            "%r instance isn't saved. Use bulk=False or save "
+                            "the object first. but must be." % obj
+                        )
+                    check_and_update_obj(obj)
+                    pks.append(obj.pk)
+
+                self.model._base_manager.using(db).filter(pk__in=pks).update(**{
+                    self.content_type_field_name: self.content_type,
+                    self.object_id_field_name: self.pk_val,
+                })
+            else:
+                with transaction.atomic(using=db, savepoint=False):
+                    for obj in objs:
+                        check_and_update_obj(obj)
+                        obj.save()
         add.alters_data = True
 
         def remove(self, *objs, **kwargs):
@@ -542,13 +565,14 @@ def create_generic_related_manager(superclass, rel):
             # could be affected by `manager.clear()`. Refs #19816.
             objs = tuple(objs)
 
+            bulk = kwargs.pop('bulk', True)
             clear = kwargs.pop('clear', False)
 
             db = router.db_for_write(self.model, instance=self.instance)
             with transaction.atomic(using=db, savepoint=False):
                 if clear:
                     self.clear()
-                    self.add(*objs)
+                    self.add(*objs, bulk=bulk)
                 else:
                     old_objs = set(self.using(db).all())
                     new_objs = []
@@ -559,7 +583,7 @@ def create_generic_related_manager(superclass, rel):
                             new_objs.append(obj)
 
                     self.remove(*old_objs)
-                    self.add(*new_objs)
+                    self.add(*new_objs, bulk=bulk)
         set.alters_data = True
 
         def create(self, **kwargs):
diff --git a/django/db/models/fields/related.py b/django/db/models/fields/related.py
index c10b2658b5..77b7bcb8d3 100644
--- a/django/db/models/fields/related.py
+++ b/django/db/models/fields/related.py
@@ -765,16 +765,36 @@ def create_foreign_related_manager(superclass, rel):
             cache_name = self.field.related_query_name()
             return queryset, rel_obj_attr, instance_attr, False, cache_name
 
-        def add(self, *objs):
+        def add(self, *objs, **kwargs):
+            bulk = kwargs.pop('bulk', True)
             objs = list(objs)
             db = router.db_for_write(self.model, instance=self.instance)
-            with transaction.atomic(using=db, savepoint=False):
+
+            def check_and_update_obj(obj):
+                if not isinstance(obj, self.model):
+                    raise TypeError("'%s' instance expected, got %r" % (
+                        self.model._meta.object_name, obj,
+                    ))
+                setattr(obj, self.field.name, self.instance)
+
+            if bulk:
+                pks = []
                 for obj in objs:
-                    if not isinstance(obj, self.model):
-                        raise TypeError("'%s' instance expected, got %r" %
-                                        (self.model._meta.object_name, obj))
-                    setattr(obj, self.field.name, self.instance)
-                    obj.save()
+                    check_and_update_obj(obj)
+                    if obj._state.adding or obj._state.db != db:
+                        raise ValueError(
+                            "%r instance isn't saved. Use bulk=False or save "
+                            "the object first." % obj
+                        )
+                    pks.append(obj.pk)
+                self.model._base_manager.using(db).filter(pk__in=pks).update(**{
+                    self.field.name: self.instance,
+                })
+            else:
+                with transaction.atomic(using=db, savepoint=False):
+                    for obj in objs:
+                        check_and_update_obj(obj)
+                        obj.save()
         add.alters_data = True
 
         def create(self, **kwargs):
@@ -835,6 +855,7 @@ def create_foreign_related_manager(superclass, rel):
             # could be affected by `manager.clear()`. Refs #19816.
             objs = tuple(objs)
 
+            bulk = kwargs.pop('bulk', True)
             clear = kwargs.pop('clear', False)
 
             if self.field.null:
@@ -842,7 +863,7 @@ def create_foreign_related_manager(superclass, rel):
                 with transaction.atomic(using=db, savepoint=False):
                     if clear:
                         self.clear()
-                        self.add(*objs)
+                        self.add(*objs, bulk=bulk)
                     else:
                         old_objs = set(self.using(db).all())
                         new_objs = []
@@ -852,10 +873,10 @@ def create_foreign_related_manager(superclass, rel):
                             else:
                                 new_objs.append(obj)
 
-                        self.remove(*old_objs)
-                        self.add(*new_objs)
+                        self.remove(*old_objs, bulk=bulk)
+                        self.add(*new_objs, bulk=bulk)
             else:
-                self.add(*objs)
+                self.add(*objs, bulk=bulk)
         set.alters_data = True
 
     return RelatedManager
diff --git a/docs/ref/models/relations.txt b/docs/ref/models/relations.txt
index 9e27e27151..d480ec18c2 100644
--- a/docs/ref/models/relations.txt
+++ b/docs/ref/models/relations.txt
@@ -36,7 +36,7 @@ Related objects reference
       In this example, the methods below will be available both on
       ``topping.pizza_set`` and on ``pizza.toppings``.
 
-    .. method:: add(*objs)
+    .. method:: add(*objs, bulk=True)
 
         Adds the specified model objects to the related object set.
 
@@ -48,7 +48,13 @@ Related objects reference
 
         In the example above, in the case of a
         :class:`~django.db.models.ForeignKey` relationship,
-        ``e.save()`` is called by the related manager to perform the update.
+        :meth:`QuerySet.update() <django.db.models.query.QuerySet.update>`
+        is used to perform the update. This requires the objects to already be
+        saved.
+
+        You can use the ``bulk=False`` argument to instead have the related
+        manager perform the update by calling ``e.save()``.
+
         Using ``add()`` with a many-to-many relationship, however, will not
         call any ``save()`` methods, but rather create the relationships
         using :meth:`QuerySet.bulk_create()
@@ -56,6 +62,12 @@ Related objects reference
         some custom logic when a relationship is created, listen to the
         :data:`~django.db.models.signals.m2m_changed` signal.
 
+        .. versionchanged:: 1.9
+
+            The ``bulk`` parameter was added. In order versions, foreign key
+            updates were always done using ``save()``. Use ``bulk=False`` if
+            you require the old behavior.
+
     .. method:: create(**kwargs)
 
         Creates a new object, saves it and puts it in the related object set.
@@ -135,7 +147,7 @@ Related objects reference
         :class:`~django.db.models.ForeignKey`\s where ``null=True`` and it also
         accepts the ``bulk`` keyword argument.
 
-    .. method:: set(objs, clear=False)
+    .. method:: set(objs, bulk=True, clear=False)
 
         .. versionadded:: 1.9
 
@@ -150,6 +162,8 @@ Related objects reference
         If ``clear=True``, the ``clear()`` method is called instead and the
         whole set is added at once.
 
+        The ``bulk`` argument is passed on to :meth:`add`.
+
         Note that since ``set()`` is a compound operation, it is subject to
         race conditions. For instance, new objects may be added to the database
         in between the call to ``clear()`` and the call to ``add()``.
diff --git a/docs/releases/1.9.txt b/docs/releases/1.9.txt
index 17d0edfa3f..b7863bab03 100644
--- a/docs/releases/1.9.txt
+++ b/docs/releases/1.9.txt
@@ -400,6 +400,11 @@ Models
   managers created by ``ForeignKey``, ``GenericForeignKey``, and
   ``ManyToManyField``.
 
+* The :meth:`~django.db.models.fields.related.RelatedManager.add` method on
+  a reverse foreign key now has a ``bulk`` parameter to allow executing one
+  query regardless of the number of objects being added rather than one query
+  per object.
+
 * Added the ``keep_parents`` parameter to :meth:`Model.delete()
   <django.db.models.Model.delete>` to allow deleting only a child's data in a
   model that uses multi-table inheritance.
@@ -669,6 +674,15 @@ Dropped support for PostgreSQL 9.0
 Upstream support for PostgreSQL 9.0 ended in September 2015. As a consequence,
 Django 1.9 sets 9.1 as the minimum PostgreSQL version it officially supports.
 
+Bulk behavior of ``add()`` method of related managers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To improve performance, the ``add()`` methods of the related managers created
+by ``ForeignKey`` and ``GenericForeignKey`` changed from a series of
+``Model.save()`` calls to a single ``QuerySet.update()`` call. The change means
+that ``pre_save`` and ``post_save`` signals aren't sent anymore. You can use
+the ``bulk=False`` keyword argument to revert to the previous behavior.
+
 Template ``LoaderOrigin`` and ``StringOrigin`` are removed
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/tests/generic_relations/tests.py b/tests/generic_relations/tests.py
index 47545bdef9..0162565969 100644
--- a/tests/generic_relations/tests.py
+++ b/tests/generic_relations/tests.py
@@ -247,6 +247,32 @@ class GenericRelationsTests(TestCase):
             self.comp_func
         )
 
+    def test_add_bulk(self):
+        bacon = Vegetable.objects.create(name="Bacon", is_yucky=False)
+        t1 = TaggedItem.objects.create(content_object=self.quartz, tag="shiny")
+        t2 = TaggedItem.objects.create(content_object=self.quartz, tag="clearish")
+        # One update() query.
+        with self.assertNumQueries(1):
+            bacon.tags.add(t1, t2)
+        self.assertEqual(t1.content_object, bacon)
+        self.assertEqual(t2.content_object, bacon)
+
+    def test_add_bulk_false(self):
+        bacon = Vegetable.objects.create(name="Bacon", is_yucky=False)
+        t1 = TaggedItem.objects.create(content_object=self.quartz, tag="shiny")
+        t2 = TaggedItem.objects.create(content_object=self.quartz, tag="clearish")
+        # One save() for each object.
+        with self.assertNumQueries(2):
+            bacon.tags.add(t1, t2, bulk=False)
+        self.assertEqual(t1.content_object, bacon)
+        self.assertEqual(t2.content_object, bacon)
+
+    def test_add_rejects_unsaved_objects(self):
+        t1 = TaggedItem(content_object=self.quartz, tag="shiny")
+        msg = "<TaggedItem: shiny> instance isn't saved. Use bulk=False or save the object first."
+        with self.assertRaisesMessage(ValueError, msg):
+            self.bacon.tags.add(t1)
+
     def test_set(self):
         bacon = Vegetable.objects.create(name="Bacon", is_yucky=False)
         fatty = bacon.tags.create(tag="fatty")
@@ -266,13 +292,13 @@ class GenericRelationsTests(TestCase):
         bacon.tags.set([])
         self.assertQuerysetEqual(bacon.tags.all(), [])
 
-        bacon.tags.set([fatty, salty], clear=True)
+        bacon.tags.set([fatty, salty], bulk=False, clear=True)
         self.assertQuerysetEqual(bacon.tags.all(), [
             "<TaggedItem: fatty>",
             "<TaggedItem: salty>",
         ])
 
-        bacon.tags.set([fatty], clear=True)
+        bacon.tags.set([fatty], bulk=False, clear=True)
         self.assertQuerysetEqual(bacon.tags.all(), [
             "<TaggedItem: fatty>",
         ])
diff --git a/tests/many_to_one/tests.py b/tests/many_to_one/tests.py
index 7b2a640b2b..ba176c938c 100644
--- a/tests/many_to_one/tests.py
+++ b/tests/many_to_one/tests.py
@@ -57,7 +57,11 @@ class ManyToOneTests(TestCase):
 
         # Create a new article, and add it to the article set.
         new_article2 = Article(headline="Paul's story", pub_date=datetime.date(2006, 1, 17))
-        self.r.article_set.add(new_article2)
+        msg = "<Article: Paul's story> instance isn't saved. Use bulk=False or save the object first."
+        with self.assertRaisesMessage(ValueError, msg):
+            self.r.article_set.add(new_article2)
+
+        self.r.article_set.add(new_article2, bulk=False)
         self.assertEqual(new_article2.reporter.id, self.r.id)
         self.assertQuerysetEqual(self.r.article_set.all(),
             [
diff --git a/tests/many_to_one_null/tests.py b/tests/many_to_one_null/tests.py
index e39f527abe..06f756e35b 100644
--- a/tests/many_to_one_null/tests.py
+++ b/tests/many_to_one_null/tests.py
@@ -115,6 +115,15 @@ class ManyToOneNullTests(TestCase):
         self.assertEqual(1, self.r2.article_set.count())
         self.assertEqual(1, qs.count())
 
+    def test_add_efficiency(self):
+        r = Reporter.objects.create()
+        articles = []
+        for _ in range(3):
+            articles.append(Article.objects.create())
+        with self.assertNumQueries(1):
+            r.article_set.add(*articles)
+        self.assertEqual(r.article_set.count(), 3)
+
     def test_clear_efficiency(self):
         r = Reporter.objects.create()
         for _ in range(3):
diff --git a/tests/multiple_database/tests.py b/tests/multiple_database/tests.py
index 08493ffaf2..4354217d57 100644
--- a/tests/multiple_database/tests.py
+++ b/tests/multiple_database/tests.py
@@ -1043,6 +1043,17 @@ class RouterTestCase(TestCase):
         self.assertEqual(Book.objects.using('default').count(), 1)
         self.assertEqual(Book.objects.using('other').count(), 1)
 
+    def test_invalid_set_foreign_key_assignment(self):
+        marty = Person.objects.using('default').create(name="Marty Alchin")
+        dive = Book.objects.using('other').create(
+            title="Dive into Python",
+            published=datetime.date(2009, 5, 4),
+        )
+        # Set a foreign key set with an object from a different database
+        msg = "<Book: Dive into Python> instance isn't saved. Use bulk=False or save the object first."
+        with self.assertRaisesMessage(ValueError, msg):
+            marty.edited.set([dive])
+
     def test_foreign_key_cross_database_protection(self):
         "Foreign keys can cross databases if they two databases have a common source"
         # Create a book and author on the default database
@@ -1085,7 +1096,7 @@ class RouterTestCase(TestCase):
 
         # Set a foreign key set with an object from a different database
         try:
-            marty.edited = [pro, dive]
+            marty.edited.set([pro, dive], bulk=False)
         except ValueError:
             self.fail("Assignment across primary/replica databases with a common source should be ok")
 
@@ -1107,7 +1118,7 @@ class RouterTestCase(TestCase):
 
         # Add to a foreign key set with an object from a different database
         try:
-            marty.edited.add(dive)
+            marty.edited.add(dive, bulk=False)
         except ValueError:
             self.fail("Assignment across primary/replica databases with a common source should be ok")