from datetime import datetime
from typing import List, Optional
from pydantic import validator, BaseModel
from pymongo.results import UpdateResult, DeleteResult
from fastapi_contrib.common.utils import async_timing, get_now
from fastapi_contrib.db.utils import get_db_client, get_next_id
[docs]class NotSet(object):
...
notset = NotSet()
[docs]class MongoDBModel(BaseModel):
"""
Base Model to use for any information saving in MongoDB.
Provides `id` field as a base, populated by id-generator.
Use it as follows:
.. code-block:: python
class MyModel(MongoDBModel):
additional_field1: str
optional_field2: int = 42
class Meta:
collection = "mymodel_collection"
mymodel = MyModel(additional_field1="value")
mymodel.save()
assert mymodel.additional_field1 == "value"
assert mymodel.optional_field2 == 42
assert isinstance(mymodel.id, int)
"""
id: int = None
[docs] @validator("id", pre=True, always=True)
def set_id(cls, v, values, **kwargs) -> int:
"""
If id is supplied (ex. from DB) then use it, otherwise generate new.
"""
if v:
return v
return get_next_id()
[docs] @classmethod
def get_db_collection(cls) -> str:
return cls.Meta.collection
[docs] @classmethod
@async_timing
async def get(cls, **kwargs) -> Optional["MongoDBModel"]:
db = get_db_client()
result = await db.get(cls, **kwargs)
if not result:
return None
result["id"] = result.pop("_id")
return cls(**result)
[docs] @classmethod
@async_timing
async def delete(cls, **kwargs) -> DeleteResult:
db = get_db_client()
result = await db.delete(cls, **kwargs)
return result
[docs] @classmethod
@async_timing
async def count(cls, **kwargs) -> int:
db = get_db_client()
result = await db.count(cls, **kwargs)
return result
[docs] @classmethod
@async_timing
async def list(cls, raw=True, _limit=0, _offset=0, _sort=None, **kwargs):
db = get_db_client()
cursor = db.list(
cls, _limit=_limit, _offset=_offset, _sort=_sort, **kwargs
)
result = []
async for document in cursor:
document["id"] = document.pop("_id")
result.append(document)
if not raw:
return (cls(**record) for record in result)
return result
[docs] @async_timing
async def save(
self,
include: set = None,
exclude: set = None,
rewrite_fields: dict = None,
) -> int:
db = get_db_client()
if not rewrite_fields:
rewrite_fields = {}
for field, value in rewrite_fields.items():
setattr(self, field, value)
insert_result = await db.insert(self, include=include, exclude=exclude)
self.id = insert_result.inserted_id
return self.id
[docs] @classmethod
@async_timing
async def update_one(cls, filter_kwargs: dict, **kwargs) -> UpdateResult:
db = get_db_client()
result = await db.update_one(
cls, filter_kwargs=filter_kwargs, **kwargs
)
return result
[docs] @classmethod
@async_timing
async def update_many(cls, filter_kwargs: dict, **kwargs) -> UpdateResult:
db = get_db_client()
result = await db.update_many(
cls, filter_kwargs=filter_kwargs, **kwargs
)
return result
[docs] @classmethod
@async_timing
async def create_indexes(cls) -> Optional[List[str]]:
if hasattr(cls.Meta, "indexes"):
db = get_db_client()
collection = db.get_collection(cls.Meta.collection)
return await collection.create_indexes(cls.Meta.indexes)
[docs] class Config:
anystr_strip_whitespace = True
[docs]class MongoDBTimeStampedModel(MongoDBModel):
"""
TimeStampedModel to use when you need to have `created` field,
populated at your model creation time.
Use it as follows:
.. code-block:: python
class MyTimeStampedModel(MongoDBTimeStampedModel):
class Meta:
collection = "timestamped_collection"
mymodel = MyTimeStampedModel()
mymodel.save()
assert isinstance(mymodel.id, int)
assert isinstance(mymodel.created, datetime)
"""
created: datetime = None
[docs] @validator("created", pre=True, always=True)
def set_created_now(cls, v: datetime) -> datetime:
"""
If created is supplied (ex. from DB) -> use it, otherwise generate new.
"""
if v:
return v
return get_now()