Fix bulk create not working with multi-database setup (#252)

* Fix bulk create not working with multi-database setup

* Code review updates

Added a comment, extra line, and removed extra and in test name
Moved _save_kwargs out of recursive function to fix tests

* Add changelog message

Co-authored-by: berin <bernardoxhc@gmail.com>

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ### Changed
 - Validate `increment_by` parameter of `seq` helper when `value` is an instance of `datetime` [PR #247](https://github.com/model-bakers/model_bakery/pull/247)
 - Fix a simple typo in `bulk_create` disclaimer in docs
+- Fix bulk_create not working with multi-database setup [PR #252](https://github.com/model-bakers/model_bakery/pull/252)
 - Conditionally support NullBooleanField, it's under deprecation and will be removed in Django 4.0 [PR #25](https://github.com/model-bakers/model_bakery/pull/250)
 - Fix Django max version pin in requirements file [PR #251](https://github.com/model-bakers/model_bakery/pull/251)
 

model_bakery/baker.py

@@ -90,7 +90,7 @@ def make(
             baker.make(
                 _save_kwargs=_save_kwargs,
                 _refresh_after_create=_refresh_after_create,
-                **attrs
+                **attrs,
             )
             for _ in range(_quantity)
         ]
@@ -646,6 +646,9 @@ def bulk_create(baker, quantity, **kwargs) -> List[Model]:
     Important: there's no way to avoid save calls since Django does
     not return the created objects after a bulk_insert call.
     """
+    _save_kwargs = {}
+    if baker._using:
+        _save_kwargs = {"using": baker._using}
 
     def _save_related_objs(model, objects) -> None:
         fk_fields = [
@@ -664,9 +667,22 @@ def _save_related_objs(model, objects) -> None:
             if fk_objects:
                 _save_related_objs(fk.related_model, fk_objects)
                 for i, fk_obj in enumerate(fk_objects):
-                    fk_obj.save()
+                    fk_obj.save(**_save_kwargs)
                     setattr(objects[i], fk.name, fk_obj)
 
-    entries = [baker.prepare(**kwargs) for _ in range(quantity)]
+    entries = [
+        baker.prepare(
+            **kwargs,
+        )
+        for _ in range(quantity)
+    ]
     _save_related_objs(baker.model, entries)
-    return baker.model._base_manager.bulk_create(entries)
+
+    if baker._using:
+        # Try to use the desired DB and let Django fail if spanning
+        # relationships without the proper router setup
+        manager = baker.model._base_manager.using(baker._using)
+    else:
+        manager = baker.model._base_manager
+
+    return manager.bulk_create(entries)

tests/test_baker.py

@@ -7,7 +7,7 @@
 from django.conf import settings
 from django.db.models import Manager
 from django.db.models.signals import m2m_changed
-from django.test import TestCase
+from django.test import TestCase, override_settings
 
 from model_bakery import baker, random_gen
 from model_bakery.baker import MAX_MANY_QUANTITY
@@ -786,7 +786,7 @@ def test_allow_user_to_specify_database_via_using(self):
         assert qs.count() == 1
         assert profile in qs
 
-    def test_related_fk_database_speficied_via_using_kwarg(self):
+    def test_related_fk_database_specified_via_using_kwarg(self):
         dog = baker.make(models.Dog, _using=settings.EXTRA_DB)
         dog_qs = models.Dog.objects.using(settings.EXTRA_DB).all()
         assert dog_qs.count() == 1
@@ -796,7 +796,17 @@ def test_related_fk_database_speficied_via_using_kwarg(self):
         assert person_qs.count() == 1
         assert dog.owner in person_qs
 
-    def test_related_fk_database_speficied_via_using_kwarg_combined_with_quantity(self):
+    def test_allow_user_to_specify_database_via_using_combined_with_bulk_create(
+        self,
+    ):
+        baker.make(
+            models.Profile, _using=settings.EXTRA_DB, _quantity=10, _bulk_create=True
+        )
+        qs = models.Profile.objects.using(settings.EXTRA_DB).all()
+
+        assert qs.count() == 10
+
+    def test_related_fk_database_specified_via_using_kwarg_combined_with_quantity(self):
         dogs = baker.make(models.Dog, _using=settings.EXTRA_DB, _quantity=5)
         dog_qs = models.Dog.objects.using(settings.EXTRA_DB).all()
         person_qs = models.Person.objects.using(settings.EXTRA_DB).all()
@@ -807,6 +817,38 @@ def test_related_fk_database_speficied_via_using_kwarg_combined_with_quantity(se
             assert dog in dog_qs
             assert dog.owner in person_qs
 
+    def test_related_fk_database_specified_via_using_kwarg_combined_with_bulk_create(
+        self,
+    ):
+        # A custom router must be used when using bulk create and saving
+        # related objects in a multi-database setting.
+        class AllowRelationRouter:
+            """Custom router that allows saving of relations."""
+
+            def allow_relation(self, obj1, obj2, **hints):
+                """Allow all relations."""
+                return True
+
+        with override_settings(DATABASE_ROUTERS=[AllowRelationRouter()]):
+            baker.make(
+                models.PaymentBill,
+                _quantity=5,
+                _bulk_create=True,
+                _using=settings.EXTRA_DB,
+            )
+
+        assert models.PaymentBill.objects.all().using(settings.EXTRA_DB).count() == 5
+        assert models.User.objects.all().using(settings.EXTRA_DB).count() == 5
+
+        # Default router will not be able to save the related objects
+        with pytest.raises(ValueError):
+            baker.make(
+                models.PaymentBill,
+                _quantity=5,
+                _bulk_create=True,
+                _using=settings.EXTRA_DB,
+            )
+
     def test_allow_recipe_to_specify_database_via_using(self):
         dog = baker.make_recipe("generic.homeless_dog", _using=settings.EXTRA_DB)
         qs = models.Dog.objects.using(settings.EXTRA_DB).all()
@@ -836,13 +878,13 @@ def test_recipe_related_fk_database_specified_via_using_kwarg_and_quantity(self)
             assert dog in dog_qs
             assert dog.owner in person_qs
 
-    def test_related_m2m_database_speficied_via_using_kwarg(self):
+    def test_related_m2m_database_specified_via_using_kwarg(self):
         baker.make(models.Dog, _using=settings.EXTRA_DB, make_m2m=True)
         dog_qs = models.Dog.objects.using(settings.EXTRA_DB).all()
         assert dog_qs.count() == MAX_MANY_QUANTITY + 1
         assert not models.Dog.objects.exists()
 
-    def test_related_m2m_database_speficied_via_using_kwarg_and_quantity(self):
+    def test_related_m2m_database_specified_via_using_kwarg_and_quantity(self):
         qtd = 3
         baker.make(models.Dog, _using=settings.EXTRA_DB, make_m2m=True, _quantity=qtd)
         dog_qs = models.Dog.objects.using(settings.EXTRA_DB).all()