Add the multi tenancy support on the document level. (#57)

* Added the multi tanency support on the document level.
* updated the create many method and corrected the annotation
* Resolved the lint issues
* Resolve the creation issue
* Added the from_collection attribute in the metaclass
* Added the test cases
* Added the fixture
* use the using method to delete the record.
* Corrected the test case for the save method

mongoz/core/db/documents/document.py

@@ -49,7 +49,11 @@ async def update(
         """
         Updates a record on an instance level.
         """
-        collection = collection or self.meta.collection._collection  # type: ignore
+        if collection is None:
+            if isinstance(self.meta.from_collection, AsyncIOMotorCollection):
+                collection = self.meta.from_collection
+            elif isinstance(self.meta.collection, Collection):
+                collection = self.meta.collection._collection
         field_definitions = {
             name: (annotations, ...)
             for name, annotations in self.__annotations__.items()
@@ -69,7 +73,7 @@ async def update(
             data.update(values)
 
             await self.signals.pre_update.send(sender=self.__class__, instance=self)
-            await collection.update_one({"_id": self.id}, {"$set": data})
+            await collection.update_one({"_id": self.id}, {"$set": data})  # type: ignore
             await self.signals.post_update.send(sender=self.__class__, instance=self)
 
             for k, v in data.items():
@@ -87,7 +91,10 @@ async def create_many(cls: Type["Document"], models: List["Document"]) -> List["
             raise TypeError(f"All models must be of type {cls.__name__}")
 
         data = (model.model_dump(exclude={"id"}) for model in models)
-        results = await cls.meta.collection._collection.insert_many(data)  # type: ignore
+        if isinstance(cls.meta.from_collection, AsyncIOMotorCollection):
+            results = await cls.meta.from_collection.insert_many(data)
+        else:
+            results = await cls.meta.collection._collection.insert_many(data)  # type: ignore
         for model, inserted_id in zip(models, results.inserted_ids, strict=True):
             model.id = inserted_id
         return models
@@ -288,10 +295,14 @@ async def delete(self, collection: Union[AsyncIOMotorCollection, None] = None) -
         """Delete the document."""
         is_operation_allowed(self)
 
-        collection = collection or self.meta.collection._collection  # type: ignore
+        if collection is None:
+            if isinstance(self.meta.from_collection, AsyncIOMotorCollection):
+                collection = self.meta.from_collection
+            elif isinstance(self.meta.collection, Collection):
+                collection = self.meta.collection._collection
         await self.signals.pre_delete.send(sender=self.__class__, instance=self)
 
-        result = await collection.delete_one({"_id": self.id})
+        result = await collection.delete_one({"_id": self.id})  # type: ignore
         await self.signals.post_delete.send(sender=self.__class__, instance=self)
         return cast(int, result.deleted_count)
 
@@ -349,14 +360,18 @@ async def save(
             await movie.save()
         """
         is_operation_allowed(self)
-        collection = collection or self.meta.collection._collection  # type: ignore
+        if collection is None:
+            if isinstance(self.meta.from_collection, AsyncIOMotorCollection):
+                collection = self.meta.from_collection
+            elif isinstance(self.meta.collection, Collection):
+                collection = self.meta.collection._collection
 
         if not self.id:
             return await self.create()
 
         await self.signals.pre_save.send(sender=self.__class__, instance=self)
 
-        await collection.update_one(
+        await collection.update_one(  # type: ignore
             {"_id": self.id}, {"$set": self.model_dump(exclude={"id", "_id"})}
         )
         for k, v in self.model_dump(exclude={"id"}).items():

mongoz/core/db/documents/document_row.py

@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Sequence, Type, Union, cast
 
+from motor.motor_asyncio import AsyncIOMotorCollection
+
 from mongoz.core.db.documents.base import MongozBaseModel
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -19,6 +21,7 @@ def from_row(
         is_defer_fields: bool = False,
         only_fields: Union[Sequence[str], None] = None,
         defer_fields: Union[Sequence[str], None] = None,
+        from_collection:  Union[AsyncIOMotorCollection, None] = None
     ) -> Union[Type["Document"], None]:
         """
         Class method to convert a dictionary row result into a Document row type.
@@ -55,6 +58,7 @@ def from_row(
                     item[column] = value
 
         model = cast("Type[Document]", cls(**item))  # type: ignore
+        model.Meta.from_collection = from_collection
         return model
 
     @classmethod

mongoz/core/db/documents/metaclasses.py

@@ -15,6 +15,7 @@
     no_type_check,
 )
 
+from motor.motor_asyncio import AsyncIOMotorCollection
 from pydantic._internal._model_construction import ModelMetaclass
 
 from mongoz.core.connection.collections import Collection
@@ -47,6 +48,7 @@ class MetaInfo:
         "database",
         "manager",
         "autogenerate_index",
+        "from_collection"
     )
 
     def __init__(self, meta: Any = None, **kwargs: Any) -> None:
@@ -65,6 +67,7 @@ def __init__(self, meta: Any = None, **kwargs: Any) -> None:
         self.signals: Optional[Broadcaster] = {}  # type: ignore
         self.manager: "Manager" = getattr(meta, "manager", Manager())
         self.autogenerate_index: bool = getattr(meta, "autogenerate_index", False)
+        self.from_collection: Union[AsyncIOMotorCollection, None] = getattr(meta, "from_collection", None)
 
     def model_dump(self) -> Dict[Any, Any]:
         return {k: getattr(self, k, None) for k in self.__slots__}

mongoz/core/db/querysets/core/manager.py

@@ -364,6 +364,7 @@ async def _all(self) -> List[T]:
                 only_fields=manager._only_fields,
                 is_defer_fields=is_defer_fields,
                 defer_fields=manager._defer_fields,
+                from_collection=manager._collection
             )
             async for document in cursor
         ]

tests/models/manager/test_using.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, AsyncGenerator
 
 import pydantic
 import pytest
@@ -23,7 +23,14 @@ class Meta:
         database = "test_db"
 
 
-async def test_model_using() -> None:
+@pytest.fixture(scope="function", autouse=True)
+async def prepare_database() -> AsyncGenerator:
+    await Movie.objects.using("test_my_db").delete()
+    yield
+    await Movie.objects.using("test_my_db").delete()
+
+
+async def test_model_using_create() -> None:
     await Movie.objects.create(name="Harshali", year=2024)
     await Movie.objects.using("test_my_db").create(name="Harshali Zode", year=2024)
 
@@ -42,3 +49,54 @@ async def test_model_using() -> None:
     with pytest.raises(DocumentNotFound):
         await Movie.objects.filter(name="Harshali Zode").get()
         await Movie.objects.using("test_my_db").filter(name="Harshali").get()
+
+
+async def test_model_using_update() -> None:
+    await Movie.objects.using("test_my_db").create(name="Harshali", year=2024)
+
+    movie = await Movie.objects.using("test_my_db").get()
+    assert movie.name == "Harshali"
+
+    await movie.update(name="Harshali Zode")
+
+    movie = await Movie.objects.using("test_my_db").get()
+    assert movie.name == "Harshali Zode"
+
+    movie = await Movie.objects.using("test_my_db").filter(_id=movie.id).get()
+    assert movie.name == "Harshali Zode"
+
+    with pytest.raises(DocumentNotFound):
+        await Movie.objects.filter(name="Harshali Zode").get()
+        await Movie.objects.using("test_my_db").filter(name="Harshali").get()
+
+
+async def test_model_delete() -> None:
+    await Movie.objects.using("test_my_db").create(name="Harshali Zode", year=2024)
+
+    movie = await Movie.objects.using("test_my_db").get()
+    assert movie.name == "Harshali Zode"
+
+    await movie.delete()
+
+    with pytest.raises(DocumentNotFound):
+        movie = await Movie.objects.using("test_my_db").get()
+
+
+async def test_model_save() -> None:
+    await Movie.objects.using("test_my_db").create(name="Harshali", year=2024)
+
+    movie = await Movie.objects.using("test_my_db").get()
+    assert movie.name == "Harshali"
+
+    movie.name = "Harshali Zode"
+    await movie.save()
+
+    movie = await Movie.objects.using("test_my_db").get()
+    assert movie.name == "Harshali Zode"
+
+    movie = await Movie.objects.using("test_my_db").filter(_id=movie.id).get()
+    assert movie.name == "Harshali Zode"
+
+    with pytest.raises(DocumentNotFound):
+        await Movie.objects.filter(name="Harshali Zode").get()
+        await Movie.objects.using("test_my_db").filter(name="Harshali").get()