Index: django/db/models/sql/compiler.py =================================================================== --- django/db/models/sql/compiler.py (revision 14778) +++ django/db/models/sql/compiler.py (working copy) @@ -117,6 +117,10 @@ 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): Index: django/db/models/sql/query.py =================================================================== --- django/db/models/sql/query.py (revision 14778) +++ django/db/models/sql/query.py (working copy) @@ -11,7 +11,7 @@ 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,8 +23,16 @@ 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,10 +91,15 @@ return "" % (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,6 +144,8 @@ 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,6 +275,8 @@ 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,6 +383,7 @@ query.clear_ordering(True) query.clear_limits() + query.select_for_update = False query.select_related = False query.related_select_cols = [] query.related_select_fields = [] Index: django/db/models/manager.py =================================================================== --- django/db/models/manager.py (revision 14778) +++ django/db/models/manager.py (working copy) @@ -164,6 +164,9 @@ 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) Index: django/db/models/query.py =================================================================== --- django/db/models/query.py (revision 14778) +++ django/db/models/query.py (working copy) @@ -432,6 +432,7 @@ 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,6 +581,18 @@ 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. Index: django/db/backends/mysql/base.py =================================================================== --- django/db/backends/mysql/base.py (revision 14778) +++ django/db/backends/mysql/base.py (working copy) @@ -23,7 +23,7 @@ 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,6 +124,8 @@ 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,6 +137,7 @@ 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 Index: django/db/backends/oracle/base.py =================================================================== --- django/db/backends/oracle/base.py (revision 14778) +++ django/db/backends/oracle/base.py (working copy) @@ -48,6 +48,8 @@ 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,6 +287,12 @@ '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 '' Index: django/db/backends/__init__.py =================================================================== --- django/db/backends/__init__.py (revision 14778) +++ django/db/backends/__init__.py (working copy) @@ -103,6 +103,8 @@ # 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,6 +284,16 @@ """ 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 Index: django/db/backends/postgresql_psycopg2/base.py =================================================================== --- django/db/backends/postgresql_psycopg2/base.py (revision 14778) +++ django/db/backends/postgresql_psycopg2/base.py (working copy) @@ -19,6 +19,7 @@ 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,8 +71,21 @@ 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: Index: tests/modeltests/select_for_update/__init__.py =================================================================== Index: tests/modeltests/select_for_update/tests.py =================================================================== --- tests/modeltests/select_for_update/tests.py (revision 0) +++ tests/modeltests/select_for_update/tests.py (revision 0) @@ -0,0 +1,173 @@ +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]) Index: tests/modeltests/select_for_update/models.py =================================================================== --- tests/modeltests/select_for_update/models.py (revision 0) +++ tests/modeltests/select_for_update/models.py (revision 0) @@ -0,0 +1,4 @@ +from django.db import models + +class Person(models.Model): + name = models.CharField(max_length=30) \ No newline at end of file Index: AUTHORS =================================================================== --- AUTHORS (revision 14778) +++ AUTHORS (working copy) @@ -522,6 +522,7 @@ Gasper Zejn Jarek Zgoda Cheng Zhang + Dan Fairs A big THANK YOU goes to: Index: docs/ref/models/querysets.txt =================================================================== --- docs/ref/models/querysets.txt (revision 14778) +++ docs/ref/models/querysets.txt (working copy) @@ -975,6 +975,53 @@ # 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 ------------------------------------ Index: docs/ref/databases.txt =================================================================== --- docs/ref/databases.txt (revision 14778) +++ docs/ref/databases.txt (working copy) @@ -362,6 +362,15 @@ 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