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

Add "using" method in model queryset manager, which will be used to change database at runtime while querying. #28

Merged
merged 7 commits into from
May 2, 2024
10 changes: 10 additions & 0 deletions docs/queries.md
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,16 @@ users = await User.objects.filter(id__desc=False) # same as asc True
users = await User.objects.filter(id__neq=1) # same as asc True
```

### Using

Change the database while querying, only need to supply the database name in order to change the database.

=== "Manager"

```python
users = await User.objects.using("my_mongo_db").all()
```

### Query

The `query` is what is used by the `queryset` instead of the `manager`. In other words, the `query`
Expand Down
12 changes: 10 additions & 2 deletions mongoz/core/db/documents/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@
import bson
import pydantic
from bson.errors import InvalidId
from motor.motor_asyncio import AsyncIOMotorCollection
from pydantic import BaseModel

from mongoz.core.connection.collections import Collection
from mongoz.core.db.documents.document_row import DocumentRow
from mongoz.core.db.documents.metaclasses import EmbeddedModelMetaClass
from mongoz.core.db.fields.base import MongozField
Expand All @@ -19,7 +21,9 @@ class Document(DocumentRow):
Representation of an Mongoz Document.
"""

async def create(self: "Document") -> "Document":
async def create(
self: "Document", collection: Union[AsyncIOMotorCollection, None] = None
) -> "Document":
"""
Inserts a document.
"""
Expand All @@ -28,7 +32,11 @@ async def create(self: "Document") -> "Document":
await self.signals.pre_save.send(sender=self.__class__, instance=self)

data = self.model_dump(exclude={"id"})
result = await self.meta.collection._collection.insert_one(data) # type: ignore
if collection is not None:
result = await collection.insert_one(data)
else:
if isinstance(self.meta.collection, Collection):
result = await self.meta.collection._collection.insert_one(data) # noqa
self.id = result.inserted_id

await self.signals.post_save.send(sender=self.__class__, instance=self)
Expand Down
32 changes: 30 additions & 2 deletions mongoz/core/db/querysets/core/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,32 @@ def __init__(
def __get__(self, instance: Any, owner: Any) -> "Manager":
return self.__class__(model_class=owner)

def using(self, database_name: str) -> "Manager":
"""
**Type** Public

**Arguments:**
- database_name (str): string contains the database name.

**Returns:**
- Object: self instance.

**Raises:**
- None

This method is use to select the database:
- get the data base using the get_database method form the meta \
class registry using the database_name that provided in \
argument.
- store the database object as database.
- get the collection from the data base based on \
self._collection.name
- return the self instance.
"""
database = self.model_class.meta.registry.get_database(database_name) # type: ignore
self._collection = database.get_collection(self._collection.name)._collection
return self

def clone(self) -> Any:
manager = self.__class__.__new__(self.__class__)
manager.model_class = self.model_class
Expand Down Expand Up @@ -196,7 +222,7 @@ def filter_query(self, exclude: bool = False, **kwargs: Any) -> "Manager":
else:
filter_clauses += clauses

return cast(
manager = cast(
"Manager",
self.__class__(
model_class=self.model_class,
Expand All @@ -206,6 +232,8 @@ def filter_query(self, exclude: bool = False, **kwargs: Any) -> "Manager":
defer_fields=self._defer_fields,
),
)
manager._collection = self._collection
return manager

def filter(self, **kwargs: Any) -> "Manager":
"""
Expand Down Expand Up @@ -355,7 +383,7 @@ async def create(self, **kwargs: Any) -> "Document":
Creates a mongo db document.
"""
manager: "Manager" = self.clone()
instance = await manager.model_class(**kwargs).create()
instance = await manager.model_class(**kwargs).create(manager._collection)
return cast("Document", instance)

async def delete(self) -> int:
Expand Down
44 changes: 44 additions & 0 deletions tests/models/manager/test_using.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
from typing import List, Optional

import pydantic
import pytest

import mongoz
from mongoz import Document, ObjectId
from mongoz.exceptions import DocumentNotFound
from tests.conftest import client

pytestmark = pytest.mark.anyio
pydantic_version = pydantic.__version__[:3]


class Movie(Document):
name: str = mongoz.String()
year: int = mongoz.Integer()
tags: Optional[List[str]] = mongoz.Array(str, null=True)
uuid: Optional[ObjectId] = mongoz.ObjectId(null=True)

class Meta:
registry = client
database = "test_db"


async def test_model_using() -> None:
await Movie.objects.create(name="Harshali", year=2024)
await Movie.objects.using("test_my_db").create(name="Harshali Zode", year=2024)

movie = await Movie.objects.get()
assert movie.name == "Harshali"

movie = await Movie.objects.using("test_my_db").get()
assert movie.name == "Harshali Zode"

movie = await Movie.objects.using("test_my_db").filter(name="Harshali Zode").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()
Loading