From 4d442bd53c733e6f851af812853e82bb550aefff Mon Sep 17 00:00:00 2001 From: Ashia Date: Sat, 23 Oct 2021 23:03:31 -0600 Subject: [PATCH 1/3] Fix bulk create not working with multi-database setup --- model_bakery/baker.py | 21 +++++++++++++---- tests/test_baker.py | 52 ++++++++++++++++++++++++++++++++++++++----- 2 files changed, 64 insertions(+), 9 deletions(-) diff --git a/model_bakery/baker.py b/model_bakery/baker.py index 9cc1b9a4..f9aaac58 100644 --- a/model_bakery/baker.py +++ b/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) ] @@ -648,6 +648,9 @@ def bulk_create(baker, quantity, **kwargs) -> List[Model]: """ def _save_related_objs(model, objects) -> None: + _save_kwargs = {} + if baker._using: + _save_kwargs = {"using": baker._using} fk_fields = [ f for f in model._meta.fields @@ -664,9 +667,19 @@ 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: + manager = baker.model._base_manager.using(baker._using) + else: + manager = baker.model._base_manager + return manager.bulk_create(entries) diff --git a/tests/test_baker.py b/tests/test_baker.py index baefab7e..8cbace86 100644 --- a/tests/test_baker.py +++ b/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_and_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() From 469d924b5c9048a0b25871d6e81e7cab88658caf Mon Sep 17 00:00:00 2001 From: Ashia Date: Wed, 27 Oct 2021 23:14:51 -0600 Subject: [PATCH 2/3] 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 --- model_bakery/baker.py | 9 ++++++--- tests/test_baker.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/model_bakery/baker.py b/model_bakery/baker.py index f9aaac58..242baa64 100644 --- a/model_bakery/baker.py +++ b/model_bakery/baker.py @@ -646,11 +646,11 @@ 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: - _save_kwargs = {} - if baker._using: - _save_kwargs = {"using": baker._using} fk_fields = [ f for f in model._meta.fields @@ -679,7 +679,10 @@ def _save_related_objs(model, objects) -> None: _save_related_objs(baker.model, 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) diff --git a/tests/test_baker.py b/tests/test_baker.py index 8cbace86..b4ac5176 100644 --- a/tests/test_baker.py +++ b/tests/test_baker.py @@ -817,7 +817,7 @@ def test_related_fk_database_specified_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_and_bulk_create( + 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 From 1583d1c9e8844aaccc684386d1260d8a5f06863f Mon Sep 17 00:00:00 2001 From: Ashia Date: Wed, 27 Oct 2021 23:19:28 -0600 Subject: [PATCH 3/3] Add changelog message --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1dc90cf1..fff24bc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/). ### Changed - 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) ### Removed