Fixed #30375 -- Added FOR NO KEY UPDATE support to QuerySet.select_for_update() on PostgreSQL.

This commit is contained in:
Manuel Weitzman 2020-05-10 12:25:06 -04:00 committed by Mariusz Felisiak
parent 0e893248b2
commit a4e6030904
10 changed files with 57 additions and 9 deletions

View File

@ -38,6 +38,7 @@ class BaseDatabaseFeatures:
has_select_for_update_nowait = False
has_select_for_update_skip_locked = False
has_select_for_update_of = False
has_select_for_no_key_update = False
# Does the database's SELECT FOR UPDATE OF syntax require a column rather
# than a table?
select_for_update_of_column = False

View File

@ -207,11 +207,12 @@ class BaseDatabaseOperations:
"""
return []
def for_update_sql(self, nowait=False, skip_locked=False, of=()):
def for_update_sql(self, nowait=False, skip_locked=False, of=(), no_key=False):
"""
Return the FOR UPDATE SQL clause to lock rows for an update operation.
"""
return 'FOR UPDATE%s%s%s' % (
return 'FOR%s UPDATE%s%s%s' % (
' NO KEY' if no_key else '',
' OF %s' % ', '.join(of) if of else '',
' NOWAIT' if nowait else '',
' SKIP LOCKED' if skip_locked else '',

View File

@ -18,6 +18,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
has_select_for_update_nowait = True
has_select_for_update_of = True
has_select_for_update_skip_locked = True
has_select_for_no_key_update = True
can_release_savepoints = True
supports_tablespaces = True
supports_transactions = True

View File

@ -1018,7 +1018,7 @@ class QuerySet:
return self
return self._combinator_query('difference', *other_qs)
def select_for_update(self, nowait=False, skip_locked=False, of=()):
def select_for_update(self, nowait=False, skip_locked=False, of=(), no_key=False):
"""
Return a new QuerySet instance that will select objects with a
FOR UPDATE lock.
@ -1031,6 +1031,7 @@ class QuerySet:
obj.query.select_for_update_nowait = nowait
obj.query.select_for_update_skip_locked = skip_locked
obj.query.select_for_update_of = of
obj.query.select_for_no_key_update = no_key
return obj
def select_related(self, *fields):

View File

@ -546,19 +546,26 @@ class SQLCompiler:
nowait = self.query.select_for_update_nowait
skip_locked = self.query.select_for_update_skip_locked
of = self.query.select_for_update_of
# If it's a NOWAIT/SKIP LOCKED/OF query but the backend
# doesn't support it, raise NotSupportedError to prevent a
# possible deadlock.
no_key = self.query.select_for_no_key_update
# If it's a NOWAIT/SKIP LOCKED/OF/NO KEY query but the
# backend doesn't support it, raise NotSupportedError to
# prevent a possible deadlock.
if nowait and not self.connection.features.has_select_for_update_nowait:
raise NotSupportedError('NOWAIT is not supported on this database backend.')
elif skip_locked and not self.connection.features.has_select_for_update_skip_locked:
raise NotSupportedError('SKIP LOCKED is not supported on this database backend.')
elif of and not self.connection.features.has_select_for_update_of:
raise NotSupportedError('FOR UPDATE OF is not supported on this database backend.')
elif no_key and not self.connection.features.has_select_for_no_key_update:
raise NotSupportedError(
'FOR NO KEY UPDATE is not supported on this '
'database backend.'
)
for_update_part = self.connection.ops.for_update_sql(
nowait=nowait,
skip_locked=skip_locked,
of=self.get_select_for_update_of_arguments(),
no_key=no_key,
)
if for_update_part and self.connection.features.for_update_after_from:

View File

@ -189,6 +189,7 @@ class Query(BaseExpression):
self.select_for_update_nowait = False
self.select_for_update_skip_locked = False
self.select_for_update_of = ()
self.select_for_no_key_update = False
self.select_related = False
# Arbitrary limit for select_related to prevents infinite recursion.

View File

@ -640,6 +640,7 @@ Option MariaDB MySQL
``SKIP LOCKED`` X (≥8.0.1)
``NOWAIT`` X (≥10.3) X (≥8.0.1)
``OF``
``NO KEY``
=============== ========= ==========
When using ``select_for_update()`` on MySQL, make sure you filter a queryset

View File

@ -1663,7 +1663,7 @@ For example::
``select_for_update()``
~~~~~~~~~~~~~~~~~~~~~~~
.. method:: select_for_update(nowait=False, skip_locked=False, of=())
.. method:: select_for_update(nowait=False, skip_locked=False, of=(), no_key=False)
Returns a queryset that will lock rows until the end of the transaction,
generating a ``SELECT ... FOR UPDATE`` SQL statement on supported databases.
@ -1708,6 +1708,12 @@ to refer to the queryset's model.
Restaurant.objects.select_for_update(of=('self', 'place_ptr'))
On PostgreSQL only, you can pass ``no_key=True`` in order to acquire a weaker
lock, that still allows creating rows that merely reference locked rows
(through a foreign key, for example) whilst the lock is in place. The
PostgreSQL documentation has more details about `row-level lock modes
<https://www.postgresql.org/docs/current/explicit-locking.html#LOCKING-ROWS>`_.
You can't use ``select_for_update()`` on nullable relations::
>>> Person.objects.select_related('hometown').select_for_update()
@ -1725,8 +1731,9 @@ Currently, the ``postgresql``, ``oracle``, and ``mysql`` database
backends support ``select_for_update()``. However, MariaDB 10.3+ supports only
the ``nowait`` argument and MySQL 8.0.1+ supports the ``nowait`` and
``skip_locked`` arguments. MySQL and MariaDB don't support the ``of`` argument.
The ``no_key`` argument is supported only on PostgreSQL.
Passing ``nowait=True``, ``skip_locked=True``, or ``of`` to
Passing ``nowait=True``, ``skip_locked=True``, ``no_key=True``, or ``of`` to
``select_for_update()`` using database backends that do not support these
options, such as MySQL, raises a :exc:`~django.db.NotSupportedError`. This
prevents code from unexpectedly blocking.
@ -1758,6 +1765,10 @@ raised if ``select_for_update()`` is used in autocommit mode.
PostgreSQL doesn't support ``select_for_update()`` with
:class:`~django.db.models.expressions.Window` expressions.
.. versionchanged:: 3.2
The ``no_key`` argument was added.
``raw()``
~~~~~~~~~

View File

@ -169,7 +169,9 @@ Migrations
Models
~~~~~~
* ...
* The new ``no_key`` parameter for :meth:`.QuerySet.select_for_update()`,
supported on PostgreSQL, allows acquiring weaker locks that don't block the
creation of rows that reference locked rows through a foreign key.
Requests and Responses
~~~~~~~~~~~~~~~~~~~~~~

View File

@ -97,6 +97,16 @@ class SelectForUpdateTests(TransactionTestCase):
list(Person.objects.all().select_for_update(skip_locked=True))
self.assertTrue(self.has_for_update_sql(ctx.captured_queries, skip_locked=True))
@skipUnlessDBFeature('has_select_for_no_key_update')
def test_update_sql_generated_no_key(self):
"""
The backend's FOR NO KEY UPDATE variant appears in generated SQL when
select_for_update() is invoked.
"""
with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
list(Person.objects.all().select_for_update(no_key=True))
self.assertIs(self.has_for_update_sql(ctx.captured_queries, no_key=True), True)
@skipUnlessDBFeature('has_select_for_update_of')
def test_for_update_sql_generated_of(self):
"""
@ -291,6 +301,18 @@ class SelectForUpdateTests(TransactionTestCase):
with transaction.atomic():
Person.objects.select_for_update(of=('self',)).get()
@skipIfDBFeature('has_select_for_no_key_update')
@skipUnlessDBFeature('has_select_for_update')
def test_unsuported_no_key_raises_error(self):
"""
NotSupportedError is raised if a SELECT...FOR NO KEY UPDATE... is run
on a database backend that supports FOR UPDATE but not NO KEY.
"""
msg = 'FOR NO KEY UPDATE is not supported on this database backend.'
with self.assertRaisesMessage(NotSupportedError, msg):
with transaction.atomic():
Person.objects.select_for_update(no_key=True).get()
@skipUnlessDBFeature('has_select_for_update', 'has_select_for_update_of')
def test_unrelated_of_argument_raises_error(self):
"""