Fixed #2705: added a `select_for_update()` clause to querysets.

A number of people worked on this patch over the years -- Hawkeye, Colin Grady,
KBS, sakyamuni, anih, jdemoor, and Issak Kelly. Thanks to them all, and
apologies if I missed anyone.

Special thanks to Dan Fairs for picking it up again at the end and seeing this
through to commit.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@16058 bcc190cf-cafb-0310-a4f2-bffc1f526a37

2 files changed, 20 insertions, 0 deletions

diff --git a/django/db/models/sql/compiler.py b/django/db/models/sql/compiler.py
index d425c8b32b..ace9cf4c88 100644
--- a/django/db/models/sql/compiler.py
+++ b/django/db/models/sql/compiler.py
@@ -1,11 +1,13 @@
 from django.core.exceptions import FieldError
 from django.db import connections
+from django.db import transaction
 from django.db.backends.util import truncate_name
 from django.db.models.sql.constants import *
 from django.db.models.sql.datastructures import EmptyResultSet
 from django.db.models.sql.expressions import SQLEvaluator
 from django.db.models.sql.query import get_proxied_model, get_order_dir, \
      select_related_descend, Query
+from django.db.utils import DatabaseError
 
 class SQLCompiler(object):
     def __init__(self, query, connection, using):
@@ -117,6 +119,14 @@ class SQLCompiler(object):
                         result.append('LIMIT %d' % val)
                 result.append('OFFSET %d' % self.query.low_mark)
 
+        if self.query.select_for_update and self.connection.features.has_select_for_update:
+            # If we've been asked for a NOWAIT query but the backend does not support it,
+            # raise a DatabaseError otherwise we could get an unexpected deadlock.
+            nowait = self.query.select_for_update_nowait
+            if nowait and not self.connection.features.has_select_for_update_nowait:
+                raise DatabaseError('NOWAIT is not supported on this database backend.')
+            result.append(self.connection.ops.for_update_sql(nowait=nowait))
+
         return ' '.join(result), tuple(params)
 
     def as_nested_sql(self):
@@ -677,6 +687,11 @@ class SQLCompiler(object):
         resolve_columns = hasattr(self, 'resolve_columns')
         fields = None
         has_aggregate_select = bool(self.query.aggregate_select)
+        # Set transaction dirty if we're using SELECT FOR UPDATE to ensure
+        # a subsequent commit/rollback is executed, so any database locks
+        # are released.
+        if self.query.select_for_update and transaction.is_managed(self.using):
+            transaction.set_dirty(self.using)
         for rows in self.execute_sql(MULTI):
             for row in rows:
                 if resolve_columns:
diff --git a/django/db/models/sql/query.py b/django/db/models/sql/query.py
index 59f7bcbf14..f41272cf0f 100644
--- a/django/db/models/sql/query.py
+++ b/django/db/models/sql/query.py
@@ -125,6 +125,8 @@ class Query(object):
         self.order_by = []
         self.low_mark, self.high_mark = 0, None  # Used for offset/limit
         self.distinct = False
+        self.select_for_update = False
+        self.select_for_update_nowait = False
         self.select_related = False
         self.related_select_cols = []
 
@@ -254,6 +256,8 @@ class Query(object):
         obj.order_by = self.order_by[:]
         obj.low_mark, obj.high_mark = self.low_mark, self.high_mark
         obj.distinct = self.distinct
+        obj.select_for_update = self.select_for_update
+        obj.select_for_update_nowait = self.select_for_update_nowait
         obj.select_related = self.select_related
         obj.related_select_cols = []
         obj.aggregates = copy.deepcopy(self.aggregates, memo=memo)
@@ -360,6 +364,7 @@ class Query(object):
 
         query.clear_ordering(True)
         query.clear_limits()
+        query.select_for_update = False
         query.select_related = False
         query.related_select_cols = []
         query.related_select_fields = []