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

AUTHORS

@@ -168,6 +168,7 @@ answer newbie questions, and generally made Django that much better:
     eriks@win.tue.nl
     Tomáš Ehrlich <tomas.ehrlich@gmail.com>
     Dirk Eschler <dirk.eschler@gmx.net>
+    Dan Fairs <dan@fezconsulting.com>
     Marc Fargas <telenieko@telenieko.com>
     Szilveszter Farkas <szilveszter.farkas@gmail.com>
     Grigory Fateyev <greg@dial.com.ru>

django/db/backends/__init__.py

@@ -279,6 +279,8 @@ class BaseDatabaseFeatures(object):
     # integer primary keys.
     related_fields_match_type = False
     allow_sliced_subqueries = True
+    has_select_for_update = False
+    has_select_for_update_nowait = False
 
     # Does the default test database allow multiple connections?
     # Usually an indication that the test database is in-memory
@@ -476,6 +478,15 @@ def force_no_ordering(self):
         """
         return []
 
+    def for_update_sql(self, nowait=False):
+        """
+        Returns the FOR UPDATE SQL clause to lock rows for an update operation.
+        """
+        if nowait:
+            return 'FOR UPDATE NOWAIT'
+        else:
+            return 'FOR UPDATE'
+
     def fulltext_search_sql(self, field_name):
         """
         Returns the SQL WHERE clause to use in order to perform a full-text

django/db/backends/mysql/base.py

@@ -124,6 +124,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     allows_group_by_pk = True
     related_fields_match_type = True
     allow_sliced_subqueries = False
+    has_select_for_update = True
+    has_select_for_update_nowait = False
     supports_forward_references = False
     supports_long_model_names = False
     supports_microsecond_precision = False

django/db/backends/oracle/base.py

@@ -70,6 +70,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     needs_datetime_string_cast = False
     interprets_empty_strings_as_nulls = True
     uses_savepoints = True
+    has_select_for_update = True
+    has_select_for_update_nowait = True
     can_return_id_from_insert = True
     allow_sliced_subqueries = False
     supports_subqueries_in_group_by = False

django/db/backends/postgresql_psycopg2/base.py

@@ -70,6 +70,9 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     requires_rollback_on_dirty_transaction = True
     has_real_datatype = True
     can_defer_constraint_checks = True
+    has_select_for_update = True
+    has_select_for_update_nowait = True
+
 
 class DatabaseWrapper(BaseDatabaseWrapper):
     vendor = 'postgresql'

django/db/models/manager.py

@@ -164,6 +164,9 @@ def latest(self, *args, **kwargs):
     def order_by(self, *args, **kwargs):
         return self.get_query_set().order_by(*args, **kwargs)
 
+    def select_for_update(self, *args, **kwargs):
+        return self.get_query_set().select_for_update(*args, **kwargs)
+
     def select_related(self, *args, **kwargs):
         return self.get_query_set().select_related(*args, **kwargs)
 

django/db/models/query.py

@@ -435,6 +435,7 @@ def delete(self):
         del_query._for_write = True
 
         # Disable non-supported fields.
+        del_query.query.select_for_update = False
         del_query.query.select_related = False
         del_query.query.clear_ordering()
 
@@ -583,6 +584,18 @@ def complex_filter(self, filter_obj):
         else:
             return self._filter_or_exclude(None, **filter_obj)
 
+    def select_for_update(self, **kwargs):
+        """
+        Returns a new QuerySet instance that will select objects with a
+        FOR UPDATE lock.
+        """
+        # Default to false for nowait
+        nowait = kwargs.pop('nowait', False)
+        obj = self._clone()
+        obj.query.select_for_update = True
+        obj.query.select_for_update_nowait = nowait
+        return obj
+
     def select_related(self, *fields, **kwargs):
         """
         Returns a new QuerySet instance that will select related objects.

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 @@ def as_sql(self, with_limits=True, with_col_aliases=False):
                         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 @@ def results_iter(self):
         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:

django/db/models/sql/query.py

@@ -125,6 +125,8 @@ def __init__(self, model, where=WhereNode):
         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 @@ def clone(self, klass=None, memo=None, **kwargs):
         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 @@ def get_aggregation(self, using):
 
         query.clear_ordering(True)
         query.clear_limits()
+        query.select_for_update = False
         query.select_related = False
         query.related_select_cols = []
         query.related_select_fields = []

docs/ref/databases.txt

@@ -359,6 +359,13 @@ store a timezone-aware ``time`` or ``datetime`` to a
 :class:`~django.db.models.TimeField` or :class:`~django.db.models.DateTimeField`
 respectively, a ``ValueError`` is raised rather than truncating data.
 
+Row locking with ``QuerySet.select_for_update()``
+-------------------------------------------------
+
+MySQL does not support the ``NOWAIT`` option to the ``SELECT ... FOR UPDATE``
+statement. If ``select_for_update()`` is used with ``nowait=True`` then a
+``DatabaseError`` will be raised.
+
 .. _sqlite-notes:
 
 SQLite notes
@@ -493,6 +500,12 @@ If you're getting this error, you can solve it by:
       This will simply make SQLite wait a bit longer before throwing "database
       is locked" errors; it won't really do anything to solve them.
 
+``QuerySet.select_for_update()`` not supported
+----------------------------------------------
+
+SQLite does not support the ``SELECT ... FOR UPDATE`` syntax. Calling it will
+have no effect.
+
 .. _oracle-notes:
 
 Oracle notes

docs/ref/models/querysets.txt

@@ -966,6 +966,46 @@ For example::
     # queries the database with the 'backup' alias
     >>> Entry.objects.using('backup')
 
+select_for_update
+~~~~~~~~~~~~~~~~~
+
+.. method:: select_for_update(nowait=False)
+
+.. versionadded:: 1.4
+
+Returns a queryset that will lock rows until the end of the transaction,
+generating a ``SELECT ... FOR UPDATE`` SQL statement on supported databases.
+
+For example::
+
+    entries = Entry.objects.select_for_update().filter(author=request.user)
+
+All matched entries will be locked until the end of the transaction block,
+meaning that other transactions will be prevented from changing or acquiring
+locks on them.
+
+Usually, if another transaction has already acquired a lock on one of the
+selected rows, the query will block until the lock is released. If this is
+not the behaviour you want, call ``select_for_update(nowait=True)``. This will
+make the call non-blocking. If a conflicting lock is already acquired by
+another transaction, ``django.db.utils.DatabaseError`` will be raised when
+the queryset is evaluated.
+
+Note that using ``select_for_update`` will cause the current transaction to be
+set dirty, if under transaction management. This is to ensure that Django issues
+a ``COMMIT`` or ``ROLLBACK``, releasing any locks held by the ``SELECT FOR
+UPDATE``.
+
+Currently, the ``postgresql_psycopg2``, ``oracle``, and ``mysql``
+database backends support ``select_for_update()``. However, MySQL has no
+support for the ``nowait`` argument.
+
+Passing ``nowait=True`` to ``select_for_update`` using database backends that
+do not support ``nowait``, such as MySQL, will cause a ``DatabaseError`` to be
+raised. This is in order to prevent code unexpectedly blocking.
+
+Using ``select_for_update`` on backends which do not support
+``SELECT ... FOR UPDATE`` (such as SQLite) will have no effect.
 
 Methods that do not return QuerySets
 ------------------------------------

tests/modeltests/select_for_update/__init__.py

@@ -0,0 +1 @@
+#

tests/modeltests/select_for_update/models.py

@@ -0,0 +1,4 @@
+from django.db import models
+
+class Person(models.Model):
+    name = models.CharField(max_length=30)