add get/update or create queryset method and tests

This commit is contained in:
collerek
2020-09-24 17:17:10 +02:00
parent 9f034829d6
commit a13c13550d
4 changed files with 66 additions and 1 deletions

View File

@ -203,6 +203,22 @@ class QuerySet:
self.check_single_result_rows_count(rows)
return rows[0]
async def get_or_create(self, **kwargs: Any) -> "Model":
try:
return await self.get(**kwargs)
except NoMatch:
return await self.create(**kwargs)
async def update_or_create(self, **kwargs: Any) -> "Model":
pk_name = self.model_cls.Meta.pkname
if "pk" in kwargs:
kwargs[pk_name] = kwargs.pop("pk")
if pk_name not in kwargs or kwargs.get(pk_name) is None:
return await self.create(**kwargs)
else:
model = await self.get(pk=kwargs[pk_name])
return await model.update(**kwargs)
async def all(self, **kwargs: Any) -> List["Model"]: # noqa: A003
if kwargs:
return await self.filter(**kwargs).all()