_bulk_create flag is now creating related objects as well (#206)

* Move query count to a test utilitary file

* Bulk create now also creates related objects

* Add docs on _bulk_create flag and its limitations

* Fix broken url on the docs

* Update changelog

* Link PR in the changelog

* Fix mypy erros by expliciting that both methods return None

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ## [Unreleased](https://github.com/model-bakers/model_bakery/tree/main)
 
 ### Added
+- `_bulk_create` flag is not populating related objects as well [PR #206](https://github.com/model-bakers/model_bakery/pull/206)
 - Add support for iterators on GFK fields when using _quantity param [PR #207](https://github.com/model-bakers/model_bakery/pull/207)
 
 ### Changed

docs/source/basic_usage.rst

@@ -216,6 +216,22 @@ It also works with ``prepare``:
     customers = baker.prepare('shop.Customer', _quantity=3)
     assert len(customers) == 3
 
+The ``make`` method also accepts a parameter ``_bulk_create`` to use Django's `bulk_create <https://docs.djangoproject.com/en/3.0/ref/models/querysets/#bulk-create>`_ method instead of calling ``obj.save()`` for each created instance.
+
+**Disclaimer**: Django's ``bulk_create`` does not updates the created object primary key as explained in their docs. Because of that, there's no way for model-bakery to avoid calling ``save`` method for all the foreign keys.
+
+So, for example, if you're trying to create 20 instances of a model with a foreign key using ``_bulk_create`` this will result in 21 queries (20 for each foreign key object and one to bulk create your 20 instances).
+
+If you want to avoid that, you'll have to perform individual bulk creations per foreign keys as the following example:
+
+.. code-block:: python
+
+    from model_bakery import baker
+
+    baker.prepare(User, _quantity=5, _bulk_create=True)
+    user_iter = User.objects.all().iterator()
+    baker.prepare(Profile, user=user_iter, _quantity=5, _bulk_create=True)
+
 Multi-database support
 ----------------------
 

model_bakery/baker.py

@@ -70,12 +70,7 @@ def make(
         raise InvalidQuantityException
 
     if _quantity and _bulk_create:
-        return baker.model._base_manager.bulk_create(
-            [
-                baker.prepare(_save_kwargs=_save_kwargs, **attrs)
-                for _ in range(_quantity)
-            ]
-        )
+        return bulk_create(baker, _quantity, _save_kwargs=_save_kwargs, **attrs)
     elif _quantity:
         return [
             baker.make(
@@ -619,3 +614,36 @@ def filter_rel_attrs(field_name: str, **rel_attrs) -> Dict[str, Any]:
             clean_dict[k] = v
 
     return clean_dict
+
+
+def bulk_create(baker, quantity, **kwargs) -> None:
+    """
+    Bulk create entries and all related FKs as well.
+
+    Important: there's no way to avoid save calls since Django does
+    not return the created objects after a bulk_insert call.
+    """
+
+    def _save_related_objs(model, objects) -> None:
+        fk_fields = [
+            f
+            for f in model._meta.fields
+            if isinstance(f, OneToOneField) or isinstance(f, ForeignKey)
+        ]
+
+        for fk in fk_fields:
+            fk_objects = []
+            for obj in objects:
+                fk_obj = getattr(obj, fk.name, None)
+                if fk_obj and not fk_obj.pk:
+                    fk_objects.append(fk_obj)
+
+            if fk_objects:
+                _save_related_objs(fk.related_model, fk_objects)
+                for i, fk_obj in enumerate(fk_objects):
+                    fk_obj.save()
+                    setattr(objects[i], fk.name, fk_obj)
+
+    entries = [baker.prepare(**kwargs) for _ in range(quantity)]
+    _save_related_objs(baker.model, entries)
+    return baker.model._base_manager.bulk_create(entries)

tests/generic/models.py

@@ -62,6 +62,7 @@ class User(models.Model):
     profile = models.ForeignKey(
         Profile, blank=True, null=True, on_delete=models.CASCADE
     )
+    username = models.CharField(max_length=32)
 
 
 class PaymentBill(models.Model):

tests/test_baker.py

@@ -5,7 +5,6 @@
 
 import pytest
 from django.conf import settings
-from django.db import connection
 from django.db.models import Manager
 from django.db.models.signals import m2m_changed
 from django.test import TestCase
@@ -20,6 +19,7 @@
 from model_bakery.timezone import tz_aware
 from tests.generic import models
 from tests.generic.forms import DummyGenericIPAddressFieldForm
+from tests.utils import QueryCount
 
 
 def test_import_seq_from_baker():
@@ -30,42 +30,6 @@ def test_import_seq_from_baker():
         pytest.fail("{} raised".format(ImportError.__name__))
 
 
-class QueryCount:
-    """
-    Keep track of db calls.
-
-    Example:
-    ========
-
-        qc = QueryCount()
-
-        with qc.start_count():
-            MyModel.objects.get(pk=1)
-            MyModel.objects.create()
-
-        qc.count  # 2
-
-    """
-
-    def __init__(self):
-        self.count = 0
-
-    def __call__(self, execute, sql, params, many, context):
-        """
-        `django.db.connection.execute_wrapper` callback.
-
-        https://docs.djangoproject.com/en/3.1/topics/db/instrumentation/
-        """
-        self.count += 1
-        execute(sql, params, many, context)
-
-    def start_count(self):
-        """Reset query count to 0 and return context manager for wrapping db queries."""
-        self.count = 0
-
-        return connection.execute_wrapper(self)
-
-
 class TestsModelFinder:
     def test_unicode_regression(self):
         obj = baker.prepare("generic.Person")
@@ -401,6 +365,49 @@ def test_create_one_to_one(self):
         assert isinstance(lonely_person.only_friend, models.Person)
         assert models.Person.objects.all().count() == 1
 
+    def test_create_multiple_one_to_one(self):
+        baker.make(models.LonelyPerson, _quantity=5)
+        assert models.LonelyPerson.objects.all().count() == 5
+        assert models.Person.objects.all().count() == 5
+
+    def test_bulk_create_multiple_one_to_one(self):
+        queries = QueryCount()
+
+        with queries.start_count():
+            baker.make(models.LonelyPerson, _quantity=5, _bulk_create=True)
+            assert queries.count == 6
+
+        assert models.LonelyPerson.objects.all().count() == 5
+        assert models.Person.objects.all().count() == 5
+
+    def test_chaining_bulk_create_reduces_query_count(self):
+        queries = QueryCount()
+
+        qtd = 5
+        with queries.start_count():
+            baker.make(models.Person, _quantity=qtd, _bulk_create=True)
+            person_iter = models.Person.objects.all().iterator()
+            baker.make(
+                models.LonelyPerson,
+                only_friend=person_iter,
+                _quantity=5,
+                _bulk_create=True,
+            )
+            assert queries.count == 3  # 2 bulk create and 1 select on Person
+
+        assert models.LonelyPerson.objects.all().count() == 5
+        assert models.Person.objects.all().count() == 5
+
+    def test_bulk_create_multiple_fk(self):
+        queries = QueryCount()
+
+        with queries.start_count():
+            baker.make(models.PaymentBill, _quantity=5, _bulk_create=True)
+            assert queries.count == 6
+
+        assert models.PaymentBill.objects.all().count() == 5
+        assert models.User.objects.all().count() == 5
+
     def test_create_many_to_many_if_flagged(self):
         store = baker.make(models.Store, make_m2m=True)
         assert store.employees.count() == 5

tests/utils.py

@@ -0,0 +1,37 @@
+from django.db import connection
+
+
+class QueryCount:
+    """
+    Keep track of db calls.
+
+    Example:
+    ========
+
+        qc = QueryCount()
+
+        with qc.start_count():
+            MyModel.objects.get(pk=1)
+            MyModel.objects.create()
+
+        qc.count  # 2
+
+    """
+
+    def __init__(self):
+        self.count = 0
+
+    def __call__(self, execute, sql, params, many, context):
+        """
+        `django.db.connection.execute_wrapper` callback.
+
+        https://docs.djangoproject.com/en/3.1/topics/db/instrumentation/
+        """
+        self.count += 1
+        execute(sql, params, many, context)
+
+    def start_count(self):
+        """Reset query count to 0 and return context manager for wrapping db queries."""
+        self.count = 0
+
+        return connection.execute_wrapper(self)