Enable multi-database support

CHANGELOG.md

@@ -8,15 +8,17 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ## [Unreleased](https://github.com/model-bakers/model_bakery/tree/main)
 
 ### Added
-- [dev] Add instructions and script for running `postgres` and `postgis` tests.
 - Add ability to pass `str` values to `foreign_key` for recipes from other modules [PR #120](https://github.com/model-bakers/model_bakery/pull/120)
+- Add new parameter `_using` to support multi database Django applications [PR #126](https://github.com/model-bakers/model_bakery/pull/126)
+- [dev] Add instructions and script for running `postgres` and `postgis` tests.
 
 ### Changed
 - Fixed _model parameter annotations [PR #115](https://github.com/model-bakers/model_bakery/pull/115)
 - Fixes bug when field has callable `default` [PR #117](https://github.com/model-bakers/model_bakery/pull/117)
 - [dev] Drop Python 3.5 support as it is retired (https://www.python.org/downloads/release/python-3510/)
 
 ### Removed
+- [dev] Remove support for Django<2.2 ([more about Django supported versions](https://www.djangoproject.com/download/#supported-versions))
 
 ## [1.2.0](https://pypi.org/project/model-bakery/1.2.0/)
 

docs/source/basic_usage.rst

@@ -215,3 +215,23 @@ It also works with ``prepare``:
 
     customers = baker.prepare('shop.Customer', _quantity=3)
     assert len(customers) == 3
+
+Multi-database support
+----------------------
+
+Model Bakery supports django application with more than one database.
+If you want to determine which database bakery should use, you have the ``_using`` parameter:
+
+
+.. code-block:: python
+
+    from model_bakery import baker
+
+    custom_db = "your_custom_db"
+    assert custom_db in settings.DATABASES
+    history = baker.make('shop.PurchaseHistory', _using=custom_db)
+    assert history in PurchaseHistory.objects.using(custom_db).all()
+    assert history.customer in Customer.objects.using(custom_db).all()
+    # default database tables with no data
+    assert not PurchaseHistory.objects.exists()
+    assert not Customer.objects.exists()

model_bakery/baker.py

@@ -53,6 +53,7 @@ def make(
     _save_kwargs: Optional[Dict] = None,
     _refresh_after_create: bool = False,
     _create_files: bool = False,
+    _using: str = "",
     **attrs: Any
 ):
     """Create a persisted instance from a given model its associated models.
@@ -61,7 +62,9 @@ def make(
     fields you want to define its values by yourself.
     """
     _save_kwargs = _save_kwargs or {}
-    baker = Baker.create(_model, make_m2m=make_m2m, create_files=_create_files)
+    baker = Baker.create(
+        _model, make_m2m=make_m2m, create_files=_create_files, _using=_using
+    )
     if _valid_quantity(_quantity):
         raise InvalidQuantityException
 
@@ -80,14 +83,18 @@ def make(
 
 
 def prepare(
-    _model: Union[str, Type[ModelBase]], _quantity=None, _save_related=False, **attrs
+    _model: Union[str, Type[ModelBase]],
+    _quantity=None,
+    _save_related=False,
+    _using="",
+    **attrs
 ) -> Model:
     """Create but do not persist an instance from a given model.
 
     It fill the fields with random values or you can specify which
     fields you want to define its values by yourself.
     """
-    baker = Baker.create(_model)
+    baker = Baker.create(_model, _using=_using)
     if _valid_quantity(_quantity):
         raise InvalidQuantityException
 
@@ -105,13 +112,17 @@ def _recipe(name: str) -> Any:
     return import_from_str(".".join((app, "baker_recipes", recipe_name)))
 
 
-def make_recipe(baker_recipe_name, _quantity=None, **new_attrs):
-    return _recipe(baker_recipe_name).make(_quantity=_quantity, **new_attrs)
+def make_recipe(baker_recipe_name, _quantity=None, _using="", **new_attrs):
+    return _recipe(baker_recipe_name).make(
+        _quantity=_quantity, _using=_using, **new_attrs
+    )
 
 
-def prepare_recipe(baker_recipe_name, _quantity=None, _save_related=False, **new_attrs):
+def prepare_recipe(
+    baker_recipe_name, _quantity=None, _save_related=False, _using="", **new_attrs
+):
     return _recipe(baker_recipe_name).prepare(
-        _quantity=_quantity, _save_related=_save_related, **new_attrs
+        _quantity=_quantity, _save_related=_save_related, _using=_using, **new_attrs
     )
 
 
@@ -235,16 +246,18 @@ def create(
         _model: Union[str, Type[ModelBase]],
         make_m2m: bool = False,
         create_files: bool = False,
+        _using: str = "",
     ) -> "Baker":
         """Create the baker class defined by the `BAKER_CUSTOM_CLASS` setting."""
         baker_class = _custom_baker_class() or cls
-        return baker_class(_model, make_m2m, create_files)
+        return baker_class(_model, make_m2m, create_files, _using=_using)
 
     def __init__(
         self,
         _model: Union[str, Type[ModelBase]],
         make_m2m: bool = False,
         create_files: bool = False,
+        _using: str = "",
     ) -> None:
         self.make_m2m = make_m2m
         self.create_files = create_files
@@ -253,6 +266,7 @@ def __init__(
         self.model_attrs = {}  # type: Dict[str, Any]
         self.rel_attrs = {}  # type: Dict[str, Any]
         self.rel_fields = []  # type: List[str]
+        self._using = _using
 
         if isinstance(_model, str):
             self.model = self.finder.get_model(_model)
@@ -309,6 +323,8 @@ def _make(
         **attrs: Any
     ) -> Model:
         _save_kwargs = _save_kwargs or {}
+        if self._using:
+            _save_kwargs["using"] = self._using
 
         self._clean_attrs(attrs)
         for field in self.get_fields():
@@ -341,8 +357,8 @@ def _make(
         instance = self.instance(
             self.model_attrs,
             _commit=commit,
-            _save_kwargs=_save_kwargs,
             _from_manager=_from_manager,
+            _save_kwargs=_save_kwargs,
         )
         if commit:
             for related in self.get_related():
@@ -491,7 +507,7 @@ def _handle_m2m(self, instance: Model):
                         m2m_relation.source_field_name: instance,
                         m2m_relation.target_field_name: value,
                     }
-                    make(through_model, **base_kwargs)
+                    make(through_model, _using=self._using, **base_kwargs)
 
     def _remote_field(
         self, field: Union[ForeignKey, OneToOneField]
@@ -534,6 +550,7 @@ def generate_value(self, field: Field, commit: bool = True) -> Any:
 
         # attributes like max_length, decimal_places are taken into account when
         # generating the value.
+        field._using = self._using
         generator_attrs = get_required_values(generator, field)
 
         if field.name in self.rel_fields:

model_bakery/random_gen.py

@@ -249,7 +249,7 @@ def gen_related(model, **attrs):
     return make(model, **attrs)
 
 
-gen_related.required = [_fk_model]  # type: ignore[attr-defined]
+gen_related.required = [_fk_model, "_using"]  # type: ignore[attr-defined]
 gen_related.prepare = _prepare_related  # type: ignore[attr-defined]
 
 
@@ -259,7 +259,7 @@ def gen_m2m(model, **attrs):
     return make(model, _quantity=MAX_MANY_QUANTITY, **attrs)
 
 
-gen_m2m.required = [_fk_model]  # type: ignore[attr-defined]
+gen_m2m.required = [_fk_model, "_using"]  # type: ignore[attr-defined]
 
 
 # GIS generators

model_bakery/recipe.py

@@ -21,7 +21,7 @@ def __init__(self, _model: Union[str, Type[ModelBase]], **attrs) -> None:
         # _iterator_backups will hold values of the form (backup_iterator, usable_iterator).
         self._iterator_backups = {}  # type: Dict[str, Any]
 
-    def _mapping(self, new_attrs: Dict[str, Any]) -> Dict[str, Any]:
+    def _mapping(self, _using, new_attrs: Dict[str, Any]) -> Dict[str, Any]:
         _save_related = new_attrs.get("_save_related", True)
         rel_fields_attrs = dict((k, v) for k, v in new_attrs.items() if "__" in k)
         new_attrs = dict((k, v) for k, v in new_attrs.items() if "__" not in k)
@@ -47,22 +47,24 @@ def _mapping(self, new_attrs: Dict[str, Any]) -> Dict[str, Any]:
                         a[key] = rel_fields_attrs.pop(key)
                 recipe_attrs = baker.filter_rel_attrs(k, **a)
                 if _save_related:
-                    mapping[k] = v.recipe.make(**recipe_attrs)
+                    mapping[k] = v.recipe.make(_using=_using, **recipe_attrs)
                 else:
-                    mapping[k] = v.recipe.prepare(**recipe_attrs)
+                    mapping[k] = v.recipe.prepare(_using=_using, **recipe_attrs)
             elif isinstance(v, related):
                 mapping[k] = v.make()
         mapping.update(new_attrs)
         mapping.update(rel_fields_attrs)
         return mapping
 
-    def make(self, **attrs: Any) -> Union[Model, List[Model]]:
-        return baker.make(self._model, **self._mapping(attrs))
+    def make(self, _using="", **attrs: Any) -> Union[Model, List[Model]]:
+        return baker.make(self._model, _using=_using, **self._mapping(_using, attrs))
 
-    def prepare(self, **attrs: Any) -> Union[Model, List[Model]]:
+    def prepare(self, _using="", **attrs: Any) -> Union[Model, List[Model]]:
         defaults = {"_save_related": False}
         defaults.update(attrs)
-        return baker.prepare(self._model, **self._mapping(defaults))
+        return baker.prepare(
+            self._model, _using=_using, **self._mapping(_using, defaults)
+        )
 
     def extend(self, **attrs) -> "Recipe":
         attr_mapping = self.attr_mapping.copy()

tests/conftest.py

@@ -18,23 +18,28 @@ def pytest_configure():
     if test_db == "sqlite":
         db_engine = "django.db.backends.sqlite3"
         db_name = ":memory:"
+        extra_db_name = ":memory:"
     elif test_db == "postgresql":
         using_postgres_flag = True
         db_engine = "django.db.backends.postgresql_psycopg2"
         db_name = "postgres"
         installed_apps = ["django.contrib.postgres"] + installed_apps
+        extra_db_name = "extra_db"
     elif test_db == "postgis":
         using_postgres_flag = True
         db_engine = "django.contrib.gis.db.backends.postgis"
         db_name = "postgres"
+        extra_db_name = "extra_db"
         installed_apps = [
             "django.contrib.postgres",
             "django.contrib.gis",
         ] + installed_apps
     else:
         raise NotImplementedError("Tests for % are not supported", test_db)
 
+    EXTRA_DB = "extra"
     settings.configure(
+        EXTRA_DB=EXTRA_DB,
         DATABASES={
             "default": {
                 "ENGINE": db_engine,
@@ -44,7 +49,16 @@ def pytest_configure():
                 "PORT": os.environ.get("PGPORT", ""),
                 "USER": os.environ.get("PGUSER", ""),
                 "PASSWORD": os.environ.get("PGPASSWORD", ""),
-            }
+            },
+            # Extra DB used to test multi database support
+            EXTRA_DB: {
+                "ENGINE": db_engine,
+                "NAME": extra_db_name,
+                "HOST": "localhost",
+                "PORT": os.environ.get("PGPORT", ""),
+                "USER": os.environ.get("PGUSER", ""),
+                "PASSWORD": os.environ.get("PGPASSWORD", ""),
+            },
         },
         INSTALLED_APPS=installed_apps,
         LANGUAGE_CODE="en",

tests/generic/baker_recipes.py

@@ -10,6 +10,7 @@
     DummyDefaultFieldsModel,
     DummyUniqueIntegerFieldModel,
     Person,
+    School,
 )
 
 person = Recipe(
@@ -72,6 +73,8 @@
     breed="Super basset",
 )
 
+paulo_freire_school = Recipe(School, name="Escola Municipal Paulo Freire")
+
 
 class SmallDogRecipe(Recipe):
     pass

tests/generic/models.py

@@ -317,7 +317,7 @@ class Ambiguous(models.Model):
 
 
 class School(models.Model):
-    name = models.CharField(max_length=10)
+    name = models.CharField(max_length=50)
     students = models.ManyToManyField(Person, through="SchoolEnrollment")