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

10 files changed, 57 insertions, 9 deletions

diff --git a/django/db/backends/base/features.py b/django/db/backends/base/features.py
index 35ce3ba299..31113e1c7a 100644
--- a/django/db/backends/base/features.py
+++ b/django/db/backends/base/features.py
@@ -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
diff --git a/django/db/backends/base/operations.py b/django/db/backends/base/operations.py
index 6d0f5c68b3..2e283a3193 100644
--- a/django/db/backends/base/operations.py
+++ b/django/db/backends/base/operations.py
@@ -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 '',
diff --git a/django/db/backends/postgresql/features.py b/django/db/backends/postgresql/features.py
index a4ba0b99fc..f8d2ea1286 100644
--- a/django/db/backends/postgresql/features.py
+++ b/django/db/backends/postgresql/features.py
@@ -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
diff --git a/django/db/models/query.py b/django/db/models/query.py
index 5c70229263..07d6ffd4ca 100644
--- a/django/db/models/query.py
+++ b/django/db/models/query.py
@@ -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):
diff --git a/django/db/models/sql/compiler.py b/django/db/models/sql/compiler.py
index c368a59226..29e33c075f 100644
--- a/django/db/models/sql/compiler.py
+++ b/django/db/models/sql/compiler.py
@@ -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:
diff --git a/django/db/models/sql/query.py b/django/db/models/sql/query.py
index 375e22c4de..ce18098fd2 100644
--- a/django/db/models/sql/query.py
+++ b/django/db/models/sql/query.py
@@ -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.
diff --git a/docs/ref/databases.txt b/docs/ref/databases.txt
index d129a19536..b4e23e76c0 100644
--- a/docs/ref/databases.txt
+++ b/docs/ref/databases.txt
@@ -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
diff --git a/docs/ref/models/querysets.txt b/docs/ref/models/querysets.txt
index b96cd67e98..2e862dbbb7 100644
--- a/docs/ref/models/querysets.txt
+++ b/docs/ref/models/querysets.txt
@@ -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()``
 ~~~~~~~~~
 
diff --git a/docs/releases/3.2.txt b/docs/releases/3.2.txt
index 54ffbbfd27..e646fe4681 100644
--- a/docs/releases/3.2.txt
+++ b/docs/releases/3.2.txt
@@ -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
 ~~~~~~~~~~~~~~~~~~~~~~
diff --git a/tests/select_for_update/tests.py b/tests/select_for_update/tests.py
index 3622a95c11..2197596a16 100644
--- a/tests/select_for_update/tests.py
+++ b/tests/select_for_update/tests.py
@@ -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):
         """