Commit fcd1bfb1 authored by Christos Stavrakakis's avatar Christos Stavrakakis
Browse files

Re-implement select_for_update

Previous implementation of 'select_for_update' method was based on an
attribute of the 'ForUpdateManager'. However, the model manager is shared
between all threads, which lead to race conditions. This commit
reimplements 'select_for_update', by making it return a custom QuerySet
(ForUpdateQuerySet), which overrides the 'get' and 'filter' methods.
parent cad4254d
......@@ -49,22 +49,24 @@ class ForUpdateManager(Manager):
"""
def __init__(self, *args, **kwargs):
super(ForUpdateManager, self).__init__(*args, **kwargs)
self._select_for_update = False
def select_for_update(self, *args, **kwargs):
return ForUpdateQuerySet(self.model, using=self.db)
class ForUpdateQuerySet(QuerySet):
""" QuerySet implmenting SELECT .. FOR UPDATE statement
This QuerySet overrides filter and get methods in order to implement
select_for_update() statement, by appending 'FOR UPDATE' to the end
for the SQL query.
"""
def filter(self, *args, **kwargs):
query = self.get_query_set().filter(*args, **kwargs)
if self._select_for_update:
self._select_for_update = False
return for_update(query)
else:
return query
query = super(ForUpdateQuerySet, self).filter(*args, **kwargs)
return for_update(query)
def get(self, *args, **kwargs):
if not self._select_for_update:
return self.get_query_set().get(*args, **kwargs)
query = self.filter(*args, **kwargs)
query = list(query)
num = len(query)
......@@ -80,10 +82,6 @@ class ForUpdateManager(Manager):
"Lookup parameters were %s" %
(self.model._meta.object_name, num, kwargs))
def select_for_update(self, *args, **kwargs):
self._select_for_update = True
return self
def for_update(query):
""" Rewrite query using SELECT .. FOR UPDATE.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment