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_nowait = False
has_select_for_update_skip_locked = False has_select_for_update_skip_locked = False
has_select_for_update_of = 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 # Does the database's SELECT FOR UPDATE OF syntax require a column rather
# than a table? # than a table?
select_for_update_of_column = False select_for_update_of_column = False

View File

@ -207,11 +207,12 @@ class BaseDatabaseOperations:
""" """
return [] 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 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 '', ' OF %s' % ', '.join(of) if of else '',
' NOWAIT' if nowait else '', ' NOWAIT' if nowait else '',
' SKIP LOCKED' if skip_locked 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_nowait = True
has_select_for_update_of = True has_select_for_update_of = True
has_select_for_update_skip_locked = True has_select_for_update_skip_locked = True
has_select_for_no_key_update = True
can_release_savepoints = True can_release_savepoints = True
supports_tablespaces = True supports_tablespaces = True
supports_transactions = True supports_transactions = True

View File

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

View File

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

View File

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

View File

@ -1663,7 +1663,7 @@ For example::
``select_for_update()`` ``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, Returns a queryset that will lock rows until the end of the transaction,
generating a ``SELECT ... FOR UPDATE`` SQL statement on supported databases. 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')) 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:: You can't use ``select_for_update()`` on nullable relations::
>>> Person.objects.select_related('hometown').select_for_update() >>> 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 backends support ``select_for_update()``. However, MariaDB 10.3+ supports only
the ``nowait`` argument and MySQL 8.0.1+ supports the ``nowait`` and the ``nowait`` argument and MySQL 8.0.1+ supports the ``nowait`` and
``skip_locked`` arguments. MySQL and MariaDB don't support the ``of`` argument. ``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 ``select_for_update()`` using database backends that do not support these
options, such as MySQL, raises a :exc:`~django.db.NotSupportedError`. This options, such as MySQL, raises a :exc:`~django.db.NotSupportedError`. This
prevents code from unexpectedly blocking. 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 PostgreSQL doesn't support ``select_for_update()`` with
:class:`~django.db.models.expressions.Window` expressions. :class:`~django.db.models.expressions.Window` expressions.
.. versionchanged:: 3.2
The ``no_key`` argument was added.
``raw()`` ``raw()``
~~~~~~~~~ ~~~~~~~~~

View File

@ -169,7 +169,9 @@ Migrations
Models 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 Requests and Responses
~~~~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~

View File

@ -97,6 +97,16 @@ class SelectForUpdateTests(TransactionTestCase):
list(Person.objects.all().select_for_update(skip_locked=True)) list(Person.objects.all().select_for_update(skip_locked=True))
self.assertTrue(self.has_for_update_sql(ctx.captured_queries, 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') @skipUnlessDBFeature('has_select_for_update_of')
def test_for_update_sql_generated_of(self): def test_for_update_sql_generated_of(self):
""" """
@ -291,6 +301,18 @@ class SelectForUpdateTests(TransactionTestCase):
with transaction.atomic(): with transaction.atomic():
Person.objects.select_for_update(of=('self',)).get() 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') @skipUnlessDBFeature('has_select_for_update', 'has_select_for_update_of')
def test_unrelated_of_argument_raises_error(self): def test_unrelated_of_argument_raises_error(self):
""" """