def update_returning_pk(qs: QuerySet, updates: dict) -> Set[Any]:
"""
Updates QuerySet items returning primary key values.
This method should not depend on database engine, though can have optimization performances for some engines.
:param qs: QuerySet to update
:param updates: Update items as passed to QuerySet.update(**updates) method
:return: A set of primary keys
"""
qs._for_write = True
if django_pg_returning_available(qs.db) and hasattr(
qs, 'update_returning'):
pk_name = qs.model._meta.pk.name
qs = qs.only(pk_name).update_returning(**updates)
pks = set(qs.values_list(pk_name, flat=True))
else:
with transaction.atomic(using=qs.db):
pks = set(qs.select_for_update().values_list('pk', flat=True))
QuerySet.update(qs, **updates)
return pks
def get_safe_query_set(q: QuerySet, action: str) -> QuerySet:
if action in ['list', 'retrive', 'check', 'summary']:
return q
return q.select_for_update()
