Add SELECT ... FOR NO KEY UPDATE support for Postgres (#1949)

Allow slightly reducing the lock level when using `select_for_update` on
Postgres DBs. This lock level allows concurrent inserts and deletes on
tables that have foreign key references to the locked rows.

The parameters follow the Django ORM design:

https://code.djangoproject.com/ticket/30375

* Update pypika-tortoise version requirement

Author: Pankrat

CHANGELOG.rst

@@ -15,6 +15,10 @@ Changed
 ^^^^^
 - Force async task switch every 2000 rows when converting db objects to python objects to avoid blocking the event loop (#1939)
 
+Added
+^^^^^
+- Add `no_key` parameter to `queryset.select_for_update`.
+
 0.25.0
 ------
 Fixed

CONTRIBUTORS.rst

@@ -64,6 +64,7 @@ Contributors
 * Lance Moe ``@lancemoe``
 * Markus Beckschulte ``@markus-96``
 * Frederic Aoustin ``@fraoustin``
+* Ludwig Hähne ``@pankrat``
 
 Special Thanks
 ==============

poetry.lock

@@ -11,7 +11,7 @@ include = ["CHANGELOG.rst", "LICENSE", "README.rst"]
 dynamic = [ "classifiers" ]
 requires-python = ">=3.9"
 dependencies = [
-    "pypika-tortoise (>=0.5.0,<1.0.0); python_version < '4.0'",
+    "pypika-tortoise (>=0.6.0,<1.0.0); python_version < '4.0'",
     "iso8601 (>=2.1.0,<3.0.0); python_version < '4.0'",
     "aiosqlite (>=0.16.0,<1.0.0)",
     "pytz",

pyproject.toml

@@ -674,6 +674,7 @@ async def test_select_for_update(self):
         sql2 = IntFields.filter(pk=1).only("id").select_for_update(nowait=True).sql()
         sql3 = IntFields.filter(pk=1).only("id").select_for_update(skip_locked=True).sql()
         sql4 = IntFields.filter(pk=1).only("id").select_for_update(of=("intfields",)).sql()
+        sql5 = IntFields.filter(pk=1).only("id").select_for_update(no_key=True).sql()
 
         dialect = self.db.schema_generator.DIALECT
         if dialect == "postgres":
@@ -694,6 +695,10 @@ async def test_select_for_update(self):
                     sql4,
                     'SELECT "id" "id" FROM "intfields" WHERE "id"=%s FOR UPDATE OF "intfields"',
                 )
+                self.assertEqual(
+                    sql5,
+                    'SELECT "id" "id" FROM "intfields" WHERE "id"=%s FOR NO KEY UPDATE',
+                )
             else:
                 self.assertEqual(
                     sql1,
@@ -711,6 +716,10 @@ async def test_select_for_update(self):
                     sql4,
                     'SELECT "id" "id" FROM "intfields" WHERE "id"=$1 FOR UPDATE OF "intfields"',
                 )
+                self.assertEqual(
+                    sql5,
+                    'SELECT "id" "id" FROM "intfields" WHERE "id"=$1 FOR NO KEY UPDATE',
+                )
         elif dialect == "mysql":
             self.assertEqual(
                 sql1,
@@ -728,6 +737,10 @@ async def test_select_for_update(self):
                 sql4,
                 "SELECT `id` `id` FROM `intfields` WHERE `id`=%s FOR UPDATE OF `intfields`",
             )
+            self.assertEqual(
+                sql5,
+                "SELECT `id` `id` FROM `intfields` WHERE `id`=%s FOR UPDATE",
+            )
 
     async def test_select_related(self):
         tournament = await Tournament.create(name="1")

tests/test_queryset.py

@@ -34,6 +34,7 @@ class Capabilities:
         DDL statement, and not as a separate statement.
     :param supports_transactions: Indicates that this DB supports transactions.
     :param support_for_update: Indicates that this DB supports SELECT ... FOR UPDATE SQL statement.
+    :param support_for_update_no_key: Indicates that this DB supports SELECT ... FOR NO KEY UPDATE SQL statement.
     :param support_index_hint: Support force index or use index.
     :param support_update_limit_order_by: support update/delete with limit and order by.
     :param: support_for_posix_regex_queries: indicated if the db supports posix regex queries
@@ -50,6 +51,7 @@ def __init__(
         inline_comment: bool = False,
         supports_transactions: bool = True,
         support_for_update: bool = True,
+        support_for_no_key_update: bool = False,
         # Support force index or use index?
         support_index_hint: bool = False,
         # support update/delete with limit and order by
@@ -64,6 +66,7 @@ def __init__(
         self.inline_comment = inline_comment
         self.supports_transactions = supports_transactions
         self.support_for_update = support_for_update
+        self.support_for_no_key_update = support_for_no_key_update
         self.support_index_hint = support_index_hint
         self.support_update_limit_order_by = support_update_limit_order_by
         self.support_for_posix_regex_queries = support_for_posix_regex_queries

tortoise/backends/base/client.py

@@ -45,7 +45,10 @@ class BasePostgresClient(BaseDBAsyncClient, abc.ABC):
     executor_class: type[BasePostgresExecutor] = BasePostgresExecutor
     schema_generator: type[BasePostgresSchemaGenerator] = BasePostgresSchemaGenerator
     capabilities = Capabilities(
-        "postgres", support_update_limit_order_by=False, support_for_posix_regex_queries=True
+        "postgres",
+        support_update_limit_order_by=False,
+        support_for_posix_regex_queries=True,
+        support_for_no_key_update=True,
     )
     connection_class: AsyncConnection | Connection | None = None
     loop: AbstractEventLoop | None = None

tortoise/backends/base_postgres/client.py

@@ -1147,14 +1147,17 @@ def select_for_update(
         skip_locked: bool = False,
         of: tuple[str, ...] = (),
         using_db: BaseDBAsyncClient | None = None,
+        no_key: bool = False,
     ) -> QuerySet[Self]:
         """
         Make QuerySet select for update.
 
         Returns a queryset that will lock rows until the end of the transaction,
         generating a SELECT ... FOR UPDATE SQL statement on supported databases.
         """
-        return cls._db_queryset(using_db, for_write=True).select_for_update(nowait, skip_locked, of)
+        return cls._db_queryset(using_db, for_write=True).select_for_update(
+            nowait, skip_locked, of, no_key=no_key
+        )
 
     @classmethod
     async def update_or_create(

tortoise/models.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from collections import defaultdict
 import types
+from collections import defaultdict
 from collections.abc import AsyncIterator, Callable, Collection, Generator, Iterable
 from copy import copy
 from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, cast, overload
@@ -326,6 +326,7 @@ class QuerySet(AwaitableQuery[MODEL]):
         "_select_for_update_nowait",
         "_select_for_update_skip_locked",
         "_select_for_update_of",
+        "_select_for_update_no_key",
         "_select_related",
         "_select_related_idx",
         "_use_indexes",
@@ -351,6 +352,7 @@ def __init__(self, model: type[MODEL]) -> None:
         self._select_for_update_nowait: bool = False
         self._select_for_update_skip_locked: bool = False
         self._select_for_update_of: set[str] = set()
+        self._select_for_update_no_key: bool = False
         self._select_related: set[str] = set()
         self._select_related_idx: list[
             tuple[type[Model], int, Table | str, type[Model], Iterable[str | None]]
@@ -385,6 +387,7 @@ def _clone(self) -> QuerySet[MODEL]:
         queryset._select_for_update_nowait = self._select_for_update_nowait
         queryset._select_for_update_skip_locked = self._select_for_update_skip_locked
         queryset._select_for_update_of = self._select_for_update_of
+        queryset._select_for_update_no_key = self._select_for_update_no_key
         queryset._select_related = self._select_related
         queryset._select_related_idx = self._select_related_idx
         queryset._force_indexes = self._force_indexes
@@ -572,20 +575,40 @@ def distinct(self) -> QuerySet[MODEL]:
         return queryset
 
     def select_for_update(
-        self, nowait: bool = False, skip_locked: bool = False, of: tuple[str, ...] = ()
+        self,
+        nowait: bool = False,
+        skip_locked: bool = False,
+        of: tuple[str, ...] = (),
+        no_key: bool = False,
     ) -> QuerySet[MODEL]:
         """
         Make QuerySet select for update.
 
         Returns a queryset that will lock rows until the end of the transaction,
         generating a SELECT ... FOR UPDATE SQL statement on supported databases.
+
+        :param nowait:
+            If `True`, raise an error if the lock cannot be obtained immediately.
+        :param skip_locked:
+            If `True`, skip rows that are already locked by other transactions instead of waiting.
+        :param of:
+            Specify the tables to lock when dealing with multiple related tables, e.g. when using `select_related`.
+            Provide a tuple of table names to indicate which tables' rows should be locked. By default, all fetched
+            rows are locked.
+        :param no_key:
+            If `True`, use the lower SELECT ... FOR NO KEY UPDATE lock strength on PostgreSQL to allow creating or
+            deleting rows in other tables that reference the locked rows via foreign keys. The parameter is ignored
+            on other backends.
         """
         if self.capabilities.support_for_update:
             queryset = self._clone()
             queryset._select_for_update = True
             queryset._select_for_update_nowait = nowait
             queryset._select_for_update_skip_locked = skip_locked
             queryset._select_for_update_of = set(of)
+            queryset._select_for_update_no_key = (
+                no_key and self.capabilities.support_for_no_key_update
+            )
             return queryset
         return self
 
@@ -1186,6 +1209,7 @@ def _make_query(self) -> None:
                 self._select_for_update_nowait,
                 self._select_for_update_skip_locked,
                 self._select_for_update_of,
+                self._select_for_update_no_key,
             )
         if self._select_related:
             for select_related in self._select_related: