Ticket #2705: for_update_tests_14778.diff

django/db/models/sql/compiler.py

@@ -117 +117 @@
                         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:
+            nowait = self.query.select_for_update_nowait and self.connection.features.has_select_for_update
+            result.append("%s" % self.connection.ops.for_update_sql(nowait=nowait))
+
         return ' '.join(result), tuple(params)
 
     def as_nested_sql(self):

django/db/models/sql/query.py

@@ -11 +11 @@
 from django.utils.tree import Node
 from django.utils.datastructures import SortedDict
 from django.utils.encoding import force_unicode
-from django.db import connections, DEFAULT_DB_ALIAS
+from django.db import connections, DEFAULT_DB_ALIAS, DatabaseError
 from django.db.models import signals
 from django.db.models.fields import FieldDoesNotExist
 from django.db.models.query_utils import select_related_descend, InvalidQuery
@@ -23 +23 @@
     ExtraWhere, AND, OR)
 from django.core.exceptions import FieldError
 
-__all__ = ['Query', 'RawQuery']
+__all__ = ['Query', 'RawQuery', 'LockNotAvailable']
 
+
+class LockNotAvailable(DatabaseError): 
+    '''
+    Raised when a query fails because a lock was not available. 
+    '''
+    pass
+
+
 class RawQuery(object):
     """
     A single raw SQL query
@@ -83 +91 @@
         return "<RawQuery: %r>" % (self.sql % self.params)
 
     def _execute_query(self):
-        self.cursor = connections[self.using].cursor()
-        self.cursor.execute(self.sql, self.params)
+        connection = connections[self.using]
+        self.cursor = connection.cursor()
+        try:
+            self.cursor.execute(self.sql, self.params)
+        except DatabaseError, e:
+            if connection.features.has_select_for_update_nowait and connection.ops.signals_lock_not_available(e):
+                raise LockNotAvailable(*e.args)
+            raise
 
-
 class Query(object):
     """
     A single SQL query.
@@ -131 +144 @@
         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 = []
 
@@ -260 +275 @@
         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 = deepcopy(self.aggregates, memo=memo)
@@ -366 +383 @@
 
         query.clear_ordering(True)
         query.clear_limits()
+        query.select_for_update = False
         query.select_related = False
         query.related_select_cols = []
         query.related_select_fields = []

django/db/models/manager.py

@@ -164 +164 @@
     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

@@ -432 +432 @@
         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()
 
@@ -580 +581 @@
         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/backends/mysql/base.py

@@ -23 +23 @@
     raise ImproperlyConfigured("MySQLdb-1.2.1p2 or newer is required; you have %s" % Database.__version__)
 
 from MySQLdb.converters import conversions
-from MySQLdb.constants import FIELD_TYPE, FLAG, CLIENT
+from MySQLdb.constants import FIELD_TYPE, FLAG, CLIENT, ER
 
 from django.db import utils
 from django.db.backends import *
@@ -124 +124 @@
     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
@@ -135 +137 @@
 
 class DatabaseOperations(BaseDatabaseOperations):
     compiler_module = "django.db.backends.mysql.compiler"
+    signals_deadlock = lambda self, e: e.args[0] == ER.LOCK_DEADLOCK
 
     def date_extract_sql(self, lookup_type, field_name):
         # http://dev.mysql.com/doc/mysql/en/date-and-time-functions.html

django/db/backends/oracle/base.py

@@ -48 +48 @@
     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
@@ -285 +287 @@
                                            'column': column_name})
         return output
 
+    def signals_deadlock(self, exception): 
+        return exception.args[0].code == 60 
+
+    def signals_lock_not_available(self, exception): 
+        return exception.args[0].code == 54 
+
     def start_transaction_sql(self):
         return ''
 

django/db/backends/__init__.py

@@ -103 +103 @@
     # 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
@@ -282 +284 @@
         """
         return []
 
+    def for_update_sql(self, nowait=False):
+        """
+        Return FOR UPDATE SQL clause to lock row for update
+        """
+        if nowait:
+            nowaitstr = ' NOWAIT'
+        else:
+            nowaitstr = ''
+        return 'FOR UPDATE' + nowaitstr
+
     def fulltext_search_sql(self, field_name):
         """
         Returns the SQL WHERE clause to use in order to perform a full-text

django/db/backends/postgresql_psycopg2/base.py

@@ -19 +19 @@
 try:
     import psycopg2 as Database
     import psycopg2.extensions
+    from psycopg2 import errorcodes
 except ImportError, e:
     from django.core.exceptions import ImproperlyConfigured
     raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
@@ -70 +71 @@
     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 DatabaseOperations(PostgresqlDatabaseOperations):
+
+    def _pg_error(self, e, code):
+        return getattr(e, 'pgcode', None) == code
+
+    def signals_deadlock(self, e):
+        return self._pg_error(e, errorcodes.DEADLOCK_DETECTED)
+        
+    def signals_lock_not_available(self, e):
+        return self._pg_error(e, errorcodes.LOCK_NOT_AVAILABLE)
+    
     def last_executed_query(self, cursor, sql, params):
         # With psycopg2, cursor objects have a "query" attribute that is the
         # exact query sent to the database. See docs here:

tests/modeltests/select_for_update/tests.py

@@ +1 @@
+import threading
+import time
+from django.conf import settings
+from django.db import connection
+from django.db import transaction, connection
+from django.db.utils import ConnectionHandler, DEFAULT_DB_ALIAS, DatabaseError
+from django.test import TransactionTestCase, skipUnlessDBFeature
+
+from models import Person
+
+class SelectForUpdateTests(TransactionTestCase):
+
+    def setUp(self):
+        connection._rollback()
+        connection._enter_transaction_management(True)
+        self.new_connections = ConnectionHandler(settings.DATABASES)
+        self.person = Person.objects.create(name='Reinhardt')
+
+        # We need to set settings.DEBUG to True so we can capture
+        # the output SQL to examine.
+        self._old_debug = settings.DEBUG
+        settings.DEBUG = True
+
+    def tearDown(self):
+        connection._leave_transaction_management(True)
+        settings.DEBUG = self._old_debug
+        try:
+            self.end_blocking_transaction()
+        except (DatabaseError, AttributeError):
+            pass
+
+    def start_blocking_transaction(self):
+        self.new_connection = self.new_connections[DEFAULT_DB_ALIAS]
+        self.new_connection._enter_transaction_management(True)
+        self.cursor = self.new_connection.cursor()
+        sql = 'SELECT * FROM %(db_table)s %(for_update)s;' % {
+            'db_table': Person._meta.db_table,
+            'for_update': self.new_connection.ops.for_update_sql(),
+            }
+        self.cursor.execute(sql, ())
+        result = self.cursor.fetchone()
+
+    def end_blocking_transaction(self):
+        self.new_connection._rollback()
+        self.new_connection.close()
+        self.new_connection._leave_transaction_management(True)
+
+    def has_for_update_sql(self, tested_connection, nowait=False):
+        for_update_sql = tested_connection.ops.for_update_sql(nowait)
+        sql = tested_connection.queries[-1]['sql']
+        return bool(sql.find(for_update_sql) > -1)
+
+    def check_exc(self, exc):
+        self.failUnless(isinstance(exc, DatabaseError))
+
+    @skipUnlessDBFeature('has_select_for_update')
+    def test_for_update_sql_generated(self):
+        """
+        Test that the backend's FOR UPDATE variant appears in
+        generated SQL when select_for_update is invoked.
+        """
+        list(Person.objects.all().select_for_update())
+        self.assertTrue(self.has_for_update_sql(connection))
+
+    @skipUnlessDBFeature('has_select_for_update_nowait')
+    def test_for_update_sql_generated_nowait(self):
+        """
+        Test that the backend's FOR UPDATE NOWAIT variant appears in
+        generated SQL when select_for_update is invoked.
+        """
+        list(Person.objects.all().select_for_update(nowait=True))
+        self.assertTrue(self.has_for_update_sql(connection, nowait=True))
+
+    @skipUnlessDBFeature('has_select_for_update_nowait')
+    def test_nowait_raises_error_on_block(self):
+        """
+        If nowait is specified, we expect an error to be raised rather
+        than blocking.
+        """
+        self.start_blocking_transaction()
+        status = []
+        thread = threading.Thread(
+            target=self.run_select_for_update,
+            args=(status,),
+            kwargs={'nowait': True},
+        )
+
+        thread.start()
+        time.sleep(1)
+        thread.join()
+        self.end_blocking_transaction()
+        self.check_exc(status[-1])
+
+    def run_select_for_update(self, status, nowait=False):
+        status.append('started')
+        try:
+            connection._rollback()
+            people = list(Person.objects.all().select_for_update(nowait=nowait))
+            people[0].name = 'Fred'
+            people[0].save()
+            connection._commit()
+        except DatabaseError, e:
+            status.append(e)
+        except Exception, e:
+            raise
+
+    @skipUnlessDBFeature('has_select_for_update')
+    def test_block(self):
+        """
+        Check that a thread running a select_for_update that
+        accesses rows being touched by a similar operation
+        on another connection blocks correctly.
+        """
+        # First, let's start the transaction in our thread.
+        self.start_blocking_transaction()
+
+        # Now, try it again using the ORM's select_for_update
+        # facility. Do this in a separate thread.
+        status = []
+        thread = threading.Thread(target=self.run_select_for_update, args=(status,))
+
+        # The thread should immediately block, but we'll sleep
+        # for a bit to make sure
+        thread.start()
+        sanity_count = 0
+        while len(status) != 1 and sanity_count < 10:
+            sanity_count += 1
+            time.sleep(1)
+        if sanity_count >= 10:
+            raise ValueError, 'Thread did not run and block'
+
+        # Check the person hasn't been updated. Since this isn't
+        # using FOR UPDATE, it won't block.
+        p = Person.objects.get(pk=self.person.pk)
+        self.assertEqual('Reinhardt', p.name)
+
+        # When we end our blocking transaction, our thread should
+        # be able to continue.
+        self.end_blocking_transaction()
+        thread.join(5.0)
+
+        # Check the thread has finished. Assuming it has, we should
+        # find that it has updated the person's name.
+        self.failIf(thread.is_alive())
+        p = Person.objects.get(pk=self.person.pk)
+        self.assertEqual('Fred', p.name)
+
+    @skipUnlessDBFeature('has_select_for_update')
+    def test_raw_lock_not_available(self):
+        """
+        Check that running a raw query which can't obtain a FOR UPDATE lock
+        raises the correct exception
+        """
+        self.start_blocking_transaction()
+        def raw(status):
+            try:
+                list(
+                    Person.objects.raw(
+                        'SELECT * FROM %s %s' % (
+                            Person._meta.db_table,
+                            connection.ops.for_update_sql(nowait=True)
+                        )
+                    )
+                )
+            except DatabaseError, e:
+                status.append(e)
+        status = []
+        thread = threading.Thread(target=raw, kwargs={'status': status})
+        thread.start()
+        time.sleep(1)
+        thread.join()
+        self.end_blocking_transaction()
+        self.check_exc(status[-1])

tests/modeltests/select_for_update/models.py

@@ +1 @@
+from django.db import models
+
+class Person(models.Model):
+    name = models.CharField(max_length=30)
+ No newline at end of file

AUTHORS

@@ -522 +522 @@
     Gasper Zejn <zejn@kiberpipa.org>
     Jarek Zgoda <jarek.zgoda@gmail.com>
     Cheng Zhang
+    Dan Fairs <dan@fezconsulting.com>
 
 A big THANK YOU goes to:
 

docs/ref/models/querysets.txt

@@ -975 +975 @@
     # queries the database with the 'backup' alias
     >>> Entry.objects.using('backup')
 
+``select_for_update(nowait=False)`` 
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 
+         
+Returns a queryset that will lock rows until the end of the transaction,  
+generating a SELECT ... FOR UPDATE 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.models.LockNotAvailable`` will be raised when  
+the queryset is evaluated. 
+ 
+Using blocking locks on a database can lead to deadlocks. This occurs when two  
+concurrent transactions are both waiting on a lock the other transaction  
+already holds. To deal with deadlocks, wrap your views that use  
+``select_for_update(nowait=False)`` with the  
+``django.views.decorators.deadlock.handle_deadlocks`` decorator.  
+ 
+For example:: 
+ 
+    from django.db import transaction 
+    from django.views.decorators.deadlock import handle_deadlocks 
+ 
+    @handle_deadlocks(max_retries=2) 
+    @transaction.commit_on_success 
+    def my_view(request): 
+        ... 
+ 
+If the database engine detects a deadlock involving ``my_view`` and decides  
+to abort its transaction, it will be automatically retried. If deadlocks keep  
+occurring after two repeated attempts,  
+``django.views.decorators.DeadlockError`` will be raised, which can be  
+propagated to the user or handled in a middleware. 
+ 
+Currently the ``postgresql_psycopg2``, ``oracle``, and ``mysql`` 
+database backends support ``select_for_update()`` but MySQL has no 
+support for the ``nowait`` argument. Other backends will simply 
+generate queries as if ``select_for_update()`` had not been used. 
 
 Methods that do not return QuerySets
 ------------------------------------

docs/ref/databases.txt

@@ -362 +362 @@
 column types have a maximum length restriction of 255 characters, regardless
 of whether ``unique=True`` is specified or not.
 
+Row locking with ``QuerySet.select_for_update()`` 
+------------------------------------------------- 
+ 
+MySQL does not support the NOWAIT option to the SELECT ... FOR UPDATE  
+statement. However, you may call the ``select_for_update()`` method of a  
+queryset with ``nowait=True``. In that case, the argument will be silently  
+discarded and the generated query will block until the requested lock can be  
+acquired. 
+
 .. _sqlite-notes:
 
 SQLite notes