Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Expand Up @@ -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)

Expand Down
24 changes: 20 additions & 4 deletions model_bakery/baker.py
Expand Up @@ -90,7 +90,7 @@ def make(
baker.make(
_save_kwargs=_save_kwargs,
_refresh_after_create=_refresh_after_create,
**attrs
**attrs,
)
for _ in range(_quantity)
]
Expand Down Expand Up @@ -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 = [
Expand All @@ -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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a comment here would be useful. Something along these lines:

Suggested change
manager = baker.model._base_manager.using(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)

But the test case you added helps to alleviate any confusion here.

else:
manager = baker.model._base_manager
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add newline between conditional blocks and return statement (for clarity).


return manager.bulk_create(entries)
52 changes: 47 additions & 5 deletions tests/test_baker.py
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand 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()]):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great test setup. I like the simplicity of using override_settings as a context manager here.

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()
Expand Down Expand Up @@ -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()
Expand Down