add post-signal in bulk-update & bulk-create

This commit is contained in:
huangsong
2022-01-14 14:21:57 +08:00
parent 1ffb28d7b0
commit fd38e6e334
3 changed files with 216 additions and 152 deletions

View File

@ -56,6 +56,21 @@ class SavePrepareMixin(RelationMixin, AliasMixin):
new_kwargs = cls.translate_columns_to_aliases(new_kwargs)
return new_kwargs
@classmethod
def prepare_model_to_update(cls, new_kwargs: dict) -> dict:
"""
Combines all preparation methods before updating.
:param new_kwargs: dictionary of model that is about to be saved
:type new_kwargs: Dict[str, str]
:return: dictionary of model that is about to be updated
:rtype: Dict[str, str]
"""
new_kwargs = cls.parse_non_db_fields(new_kwargs)
new_kwargs = cls.substitute_models_with_pks(new_kwargs)
new_kwargs = cls.reconvert_str_to_bytes(new_kwargs)
new_kwargs = cls.translate_columns_to_aliases(new_kwargs)
return new_kwargs
@classmethod
def _remove_not_ormar_fields(cls, new_kwargs: dict) -> dict:
"""

View File

@ -1,3 +1,4 @@
import asyncio
from typing import (
Any,
Dict,
@ -1049,17 +1050,23 @@ class QuerySet(Generic[T]):
:param objects: list of ormar models already initialized and ready to save.
:type objects: List[Model]
"""
ready_objects = []
for objt in objects:
new_kwargs = objt.dict()
new_kwargs = objt.prepare_model_to_save(new_kwargs)
ready_objects.append(new_kwargs)
ready_objects = [
obj.prepare_model_to_save(obj.dict())
for obj in objects
]
expr = self.table.insert().values(ready_objects)
# shouldn't use the execute_many, it's `queries.foreach(execute)`
await self.database.execute(expr)
expr = self.table.insert()
await self.database.execute_many(expr, ready_objects)
# FIXME: add pre_save signals
async def after_create(entity):
entity.set_save_status(True)
await entity.signals.post_save.send(
sender=entity.__class__, instance=entity)
for objt in objects:
objt.set_save_status(True)
await asyncio.gather(
*[after_create(entity) for entity in objects]
)
async def bulk_update( # noqa: CCR001
self, objects: List["T"], columns: List[str] = None
@ -1095,19 +1102,17 @@ class QuerySet(Generic[T]):
columns = [self.model.get_column_alias(k) for k in columns]
for objt in objects:
new_kwargs = objt.dict()
if pk_name not in new_kwargs or new_kwargs.get(pk_name) is None:
for obj in objects:
new_kwargs = obj.dict()
if new_kwargs.get(pk_name) is None:
raise ModelPersistenceError(
"You cannot update unsaved objects. "
f"{self.model.__name__} has to have {pk_name} filled."
)
new_kwargs = self.model.parse_non_db_fields(new_kwargs)
new_kwargs = self.model.substitute_models_with_pks(new_kwargs)
new_kwargs = self.model.reconvert_str_to_bytes(new_kwargs)
new_kwargs = self.model.translate_columns_to_aliases(new_kwargs)
new_kwargs = {"new_" + k: v for k, v in new_kwargs.items() if k in columns}
ready_objects.append(new_kwargs)
new_kwargs = obj.prepare_model_to_update(new_kwargs)
ready_objects.append({
"new_" + k: v for k, v in new_kwargs.items() if k in columns
})
pk_column = self.model_meta.table.c.get(self.model.get_column_alias(pk_name))
pk_column_name = self.model.get_column_alias(pk_name)
@ -1127,5 +1132,13 @@ class QuerySet(Generic[T]):
expr = str(expr)
await self.database.execute_many(expr, ready_objects)
for objt in objects:
objt.set_save_status(True)
# FIXME: add pre-update-signals
async def after_update(entity):
entity.set_save_status(True)
await entity.signals.post_update.send(
sender=entity.__class__, instance=entity
)
await asyncio.gather(
*[after_update(entity) for entity in objects]
)