From f3700b687fdd5996fe6bcd9a53c7d43a045eb848 Mon Sep 17 00:00:00 2001 From: Denis Kopitsa Date: Mon, 20 May 2024 18:30:48 +0300 Subject: [PATCH 1/7] Add for_update clause --- docs/src/piccolo/query_clauses/for_update.rst | 62 +++++++++++++++++++ docs/src/piccolo/query_clauses/index.rst | 1 + piccolo/query/methods/objects.py | 10 +++ piccolo/query/methods/select.py | 17 +++++ piccolo/query/mixins.py | 48 ++++++++++++++ tests/table/test_select.py | 27 ++++++++ 6 files changed, 165 insertions(+) create mode 100644 docs/src/piccolo/query_clauses/for_update.rst diff --git a/docs/src/piccolo/query_clauses/for_update.rst b/docs/src/piccolo/query_clauses/for_update.rst new file mode 100644 index 000000000..f159605ae --- /dev/null +++ b/docs/src/piccolo/query_clauses/for_update.rst @@ -0,0 +1,62 @@ +.. _limit: + +for_update +===== + +You can use ``for_update`` clauses with the following queries: + +* :ref:`Objects` +* :ref:`Select` + +Returns a query that will lock rows until the end of the transaction, generating a SELECT ... FOR UPDATE SQL statement. + +.. note:: Postgres and CockroachDB only. + +------------------------------------------------------------------------------- + +default +~~~~~~~ +To use select for update without extra parameters. All matched rows will be locked until the end of transaction. + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').for_update() + + +equals to: + +.. code-block:: sql + + SELECT ... FOR UPDATE + + +nowait +~~~~~~~ +If another transaction has already acquired a lock on one or more selected rows, the exception will be raised instead of waiting for another transaction + + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').for_update(nowait=True) + + +skip_locked +~~~~~~~ +Ignore locked rows + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').for_update(skip_locked=True) + + + +of +~~~~~~~ +By default, if there are many tables in query (e.x. when joining), all tables will be locked. +with `of` you can specify tables, which should be locked. + + +.. code-block:: python + + await Band.select().where(Band.manager.name == 'Guido').for_update(of=(Band, )) + diff --git a/docs/src/piccolo/query_clauses/index.rst b/docs/src/piccolo/query_clauses/index.rst index f9167ff63..f578bf2e3 100644 --- a/docs/src/piccolo/query_clauses/index.rst +++ b/docs/src/piccolo/query_clauses/index.rst @@ -29,6 +29,7 @@ by modifying the return values. ./on_conflict ./output ./returning + ./for_update .. toctree:: :maxdepth: 1 diff --git a/piccolo/query/methods/objects.py b/piccolo/query/methods/objects.py index 7f2b5aaed..7e7300da9 100644 --- a/piccolo/query/methods/objects.py +++ b/piccolo/query/methods/objects.py @@ -12,6 +12,7 @@ AsOfDelegate, CallbackDelegate, CallbackType, + ForUpdateDelegate, LimitDelegate, OffsetDelegate, OrderByDelegate, @@ -194,6 +195,7 @@ class Objects( "callback_delegate", "prefetch_delegate", "where_delegate", + "for_update_delegate", ) def __init__( @@ -213,6 +215,7 @@ def __init__( self.prefetch_delegate = PrefetchDelegate() self.prefetch(*prefetch) self.where_delegate = WhereDelegate() + self.for_update_delegate = ForUpdateDelegate() def output(self: Self, load_json: bool = False) -> Self: self.output_delegate.output( @@ -272,6 +275,12 @@ def first(self) -> First[TableInstance]: self.limit_delegate.limit(1) return First[TableInstance](query=self) + def for_update( + self: Self, nowait: bool = False, skip_locked: bool = False, of=() + ) -> Self: + self.for_update_delegate.for_update(nowait, skip_locked, of) + return self + def get(self, where: Combinable) -> Get[TableInstance]: self.where_delegate.where(where) self.limit_delegate.limit(1) @@ -322,6 +331,7 @@ def default_querystrings(self) -> t.Sequence[QueryString]: "offset_delegate", "output_delegate", "order_by_delegate", + "for_update_delegate", ): setattr(select, attr, getattr(self, attr)) diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py index 0c590918b..2e4ea4922 100644 --- a/piccolo/query/methods/select.py +++ b/piccolo/query/methods/select.py @@ -17,6 +17,7 @@ CallbackType, ColumnsDelegate, DistinctDelegate, + ForUpdateDelegate, GroupByDelegate, LimitDelegate, OffsetDelegate, @@ -150,6 +151,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]): "output_delegate", "callback_delegate", "where_delegate", + "for_update_delegate", ) def __init__( @@ -174,6 +176,7 @@ def __init__( self.output_delegate = OutputDelegate() self.callback_delegate = CallbackDelegate() self.where_delegate = WhereDelegate() + self.for_update_delegate = ForUpdateDelegate() self.columns(*columns_list) @@ -219,6 +222,12 @@ def offset(self: Self, number: int) -> Self: self.offset_delegate.offset(number) return self + def for_update( + self: Self, nowait: bool = False, skip_locked: bool = False, of=() + ) -> Self: + self.for_update_delegate.for_update(nowait, skip_locked, of) + return self + async def _splice_m2m_rows( self, response: t.List[t.Dict[str, t.Any]], @@ -618,6 +627,14 @@ def default_querystrings(self) -> t.Sequence[QueryString]: query += "{}" args.append(self.offset_delegate._offset.querystring) + if engine_type == "sqlite" and self.for_update_delegate._for_update: + raise NotImplementedError( + "SQLite doesn't support SELECT .. FOR UPDATE" + ) + + if self.for_update_delegate._for_update: + args.append(self.for_update_delegate._for_update.querystring) + querystring = QueryString(query, *args) return [querystring] diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py index d9d5f84ca..ee6bfdf38 100644 --- a/piccolo/query/mixins.py +++ b/piccolo/query/mixins.py @@ -784,3 +784,51 @@ def on_conflict( target=target, action=action_, values=values, where=where ) ) + + +@dataclass +class ForUpdate: + __slots__ = ("nowait", "skip_locked", "of") + + nowait: bool + skip_locked: bool + of: tuple[Table] + + def __post_init__(self): + if not isinstance(self.nowait, bool): + raise TypeError("nowait must be an integer") + if not isinstance(self.skip_locked, bool): + raise TypeError("skip_locked must be an integer") + if not isinstance(self.of, tuple) or not all( + hasattr(x, "_meta") for x in self.of + ): + raise TypeError("of must be an tuple of Table") + if self.nowait and self.skip_locked: + raise TypeError( + "The nowait option cannot be used with skip_locked" + ) + + @property + def querystring(self) -> QueryString: + sql = " FOR UPDATE" + if self.of: + tables = ", ".join(x._meta.tablename for x in self.of) + sql += " OF " + tables + if self.nowait: + sql += " NOWAIT" + if self.skip_locked: + sql += " SKIP LOCKED" + + return QueryString(sql) + + def __str__(self) -> str: + return self.querystring.__str__() + + +@dataclass +class ForUpdateDelegate: + + _for_update: t.Optional[ForUpdate] = None + + def for_update(self, nowait=False, skip_locked=False, of=()): + self._for_update = ForUpdate(nowait, skip_locked, of) diff --git a/tests/table/test_select.py b/tests/table/test_select.py index ebf2c3ff8..12dfce299 100644 --- a/tests/table/test_select.py +++ b/tests/table/test_select.py @@ -1028,6 +1028,33 @@ def test_select_raw(self): response, [{"name": "Pythonistas", "popularity_log": 3.0}] ) + def test_for_update(self): + """ + Make sure the for_update clause works. + """ + self.insert_rows() + + query = Band.select() + self.assertNotIn("FOR UPDATE", query.__str__()) + + query = query.for_update() + self.assertTrue(query.__str__().endswith("FOR UPDATE")) + + query = query.for_update(skip_locked=True) + self.assertTrue(query.__str__().endswith("FOR UPDATE SKIP LOCKED")) + + query = query.for_update(nowait=True) + self.assertTrue(query.__str__().endswith("FOR UPDATE NOWAIT")) + + query = query.for_update(of=(Band,)) + self.assertTrue(query.__str__().endswith("FOR UPDATE OF band")) + + with self.assertRaises(TypeError): + query = query.for_update(skip_locked=True, nowait=True) + + response = query.run_sync() + assert response is not None + class TestSelectSecret(TestCase): def setUp(self): From fdd7f530188233d581e882835dae9b212a0347f8 Mon Sep 17 00:00:00 2001 From: Denis Kopitsa Date: Fri, 9 Aug 2024 09:34:18 +0300 Subject: [PATCH 2/7] Skip testing for_update for SQLite --- tests/table/test_select.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/table/test_select.py b/tests/table/test_select.py index 12dfce299..050614798 100644 --- a/tests/table/test_select.py +++ b/tests/table/test_select.py @@ -1028,6 +1028,10 @@ def test_select_raw(self): response, [{"name": "Pythonistas", "popularity_log": 3.0}] ) + @pytest.mark.skipif( + is_running_sqlite(), + reason="SQLite doesn't support SELECT .. FOR UPDATE.", + ) def test_for_update(self): """ Make sure the for_update clause works. From 5e3c4f0f7eef2802359d9bedb63ead4f1346467c Mon Sep 17 00:00:00 2001 From: Denis Kopitsa Date: Fri, 9 Aug 2024 12:12:19 +0300 Subject: [PATCH 3/7] Fix typos in exception messages --- piccolo/query/mixins.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py index ee6bfdf38..bffa7c7e4 100644 --- a/piccolo/query/mixins.py +++ b/piccolo/query/mixins.py @@ -796,13 +796,13 @@ class ForUpdate: def __post_init__(self): if not isinstance(self.nowait, bool): - raise TypeError("nowait must be an integer") + raise TypeError("nowait must be a bool") if not isinstance(self.skip_locked, bool): - raise TypeError("skip_locked must be an integer") + raise TypeError("skip_locked must be a bool") if not isinstance(self.of, tuple) or not all( hasattr(x, "_meta") for x in self.of ): - raise TypeError("of must be an tuple of Table") + raise TypeError("of must be a tuple of Table") if self.nowait and self.skip_locked: raise TypeError( "The nowait option cannot be used with skip_locked" From edc0bc04e968b67e98a9fff391f115d0c2add069 Mon Sep 17 00:00:00 2001 From: Denis Kopitsa Date: Thu, 19 Sep 2024 18:35:27 +0300 Subject: [PATCH 4/7] Rename method to 'lock_for' --- docs/src/piccolo/query_clauses/for_update.rst | 62 ------------ docs/src/piccolo/query_clauses/index.rst | 2 +- docs/src/piccolo/query_clauses/lock_for.rst | 94 +++++++++++++++++++ piccolo/query/methods/objects.py | 32 +++++-- piccolo/query/methods/select.py | 35 +++++-- piccolo/query/mixins.py | 58 ++++++++++-- tests/table/test_select.py | 15 +-- 7 files changed, 205 insertions(+), 93 deletions(-) delete mode 100644 docs/src/piccolo/query_clauses/for_update.rst create mode 100644 docs/src/piccolo/query_clauses/lock_for.rst diff --git a/docs/src/piccolo/query_clauses/for_update.rst b/docs/src/piccolo/query_clauses/for_update.rst deleted file mode 100644 index f159605ae..000000000 --- a/docs/src/piccolo/query_clauses/for_update.rst +++ /dev/null @@ -1,62 +0,0 @@ -.. _limit: - -for_update -===== - -You can use ``for_update`` clauses with the following queries: - -* :ref:`Objects` -* :ref:`Select` - -Returns a query that will lock rows until the end of the transaction, generating a SELECT ... FOR UPDATE SQL statement. - -.. note:: Postgres and CockroachDB only. - -------------------------------------------------------------------------------- - -default -~~~~~~~ -To use select for update without extra parameters. All matched rows will be locked until the end of transaction. - -.. code-block:: python - - await Band.select(Band.name == 'Pythonistas').for_update() - - -equals to: - -.. code-block:: sql - - SELECT ... FOR UPDATE - - -nowait -~~~~~~~ -If another transaction has already acquired a lock on one or more selected rows, the exception will be raised instead of waiting for another transaction - - -.. code-block:: python - - await Band.select(Band.name == 'Pythonistas').for_update(nowait=True) - - -skip_locked -~~~~~~~ -Ignore locked rows - -.. code-block:: python - - await Band.select(Band.name == 'Pythonistas').for_update(skip_locked=True) - - - -of -~~~~~~~ -By default, if there are many tables in query (e.x. when joining), all tables will be locked. -with `of` you can specify tables, which should be locked. - - -.. code-block:: python - - await Band.select().where(Band.manager.name == 'Guido').for_update(of=(Band, )) - diff --git a/docs/src/piccolo/query_clauses/index.rst b/docs/src/piccolo/query_clauses/index.rst index f578bf2e3..60a539cfc 100644 --- a/docs/src/piccolo/query_clauses/index.rst +++ b/docs/src/piccolo/query_clauses/index.rst @@ -29,7 +29,7 @@ by modifying the return values. ./on_conflict ./output ./returning - ./for_update + ./lock_for .. toctree:: :maxdepth: 1 diff --git a/docs/src/piccolo/query_clauses/lock_for.rst b/docs/src/piccolo/query_clauses/lock_for.rst new file mode 100644 index 000000000..c209748b5 --- /dev/null +++ b/docs/src/piccolo/query_clauses/lock_for.rst @@ -0,0 +1,94 @@ +.. _lock_for: + +lock_for +======== + +You can use ``lock_for`` clauses with the following queries: + +* :ref:`Objects` +* :ref:`Select` + +Returns a query that locks rows until the end of the transaction, generating a SELECT ... FOR UPDATE SQL statement or +similar with other lock strengths. + +.. note:: Postgres and CockroachDB only. + +------------------------------------------------------------------------------- + +Basic usage without parameters: + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').lock_for() + +Equivalent to: + +.. code-block:: sql + + SELECT ... FOR UPDATE + + +lock_strength +------------- + +The parameter ``lock_strength`` controls the strength of the row lock when performing an operation in PostgreSQL. +The value can be a predefined constant from the ``LockStrength`` enum or one of the following strings (case-insensitive): + + - ``UPDATE`` (default): Acquires an exclusive lock on the selected rows, preventing other transactions from modifying or locking them until the current transaction is complete. + - ``NO KEY UPDATE`` (Postgres only): Similar to UPDATE, but allows other transactions to insert or delete rows that do not affect the primary key or unique constraints. + - ``KEY SHARE`` (Postgres only): Permits other transactions to acquire key-share or share locks, allowing non-key modifications while preventing updates or deletes. + - ``SHARE``: Acquires a shared lock, allowing other transactions to read the rows but not modify or lock them. + + +You can specify a different lock strength: + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').lock_for('share') + +Which is equivalent to: + +.. code-block:: sql + + SELECT ... FOR SHARE + + + +nowait +------ + +If another transaction has already acquired a lock on one or more selected rows, an exception will be raised instead of +waiting for the other transaction to release the lock. + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').lock_for('update', nowait=True) + + +skip_locked +----------- + +Ignore locked rows. + +.. code-block:: python + + await Band.select(Band.name == 'Pythonistas').lock_for('update', skip_locked=True) + + + +of +-- + +By default, if there are many tables in a query (e.g., when joining), all tables will be locked. +Using ``of``, you can specify which tables should be locked. + +.. code-block:: python + + await Band.select().where(Band.manager.name == 'Guido').lock_for('update', of=(Band, )) + + +Learn more +---------- + +* `Postgres docs `_ +* `CockroachDB docs `_ diff --git a/piccolo/query/methods/objects.py b/piccolo/query/methods/objects.py index 7e7300da9..8cd66b408 100644 --- a/piccolo/query/methods/objects.py +++ b/piccolo/query/methods/objects.py @@ -12,8 +12,9 @@ AsOfDelegate, CallbackDelegate, CallbackType, - ForUpdateDelegate, LimitDelegate, + LockForDelegate, + LockStrength, OffsetDelegate, OrderByDelegate, OrderByRaw, @@ -28,6 +29,7 @@ if t.TYPE_CHECKING: # pragma: no cover from piccolo.columns import Column + from piccolo.table import Table ############################################################################### @@ -195,7 +197,7 @@ class Objects( "callback_delegate", "prefetch_delegate", "where_delegate", - "for_update_delegate", + "lock_for_delegate", ) def __init__( @@ -215,7 +217,7 @@ def __init__( self.prefetch_delegate = PrefetchDelegate() self.prefetch(*prefetch) self.where_delegate = WhereDelegate() - self.for_update_delegate = ForUpdateDelegate() + self.lock_for_delegate = LockForDelegate() def output(self: Self, load_json: bool = False) -> Self: self.output_delegate.output( @@ -275,10 +277,26 @@ def first(self) -> First[TableInstance]: self.limit_delegate.limit(1) return First[TableInstance](query=self) - def for_update( - self: Self, nowait: bool = False, skip_locked: bool = False, of=() + def lock_for( + self: Self, + lock_strength: t.Union[ + LockStrength, + t.Literal[ + "UPDATE", + "NO KEY UPDATE", + "KEY SHARE", + "SHARE", + "update", + "no key update", + "key share", + "share", + ], + ] = LockStrength.update, + nowait: bool = False, + skip_locked: bool = False, + of: t.Tuple[type[Table], ...] = (), ) -> Self: - self.for_update_delegate.for_update(nowait, skip_locked, of) + self.lock_for_delegate.lock_for(lock_strength, nowait, skip_locked, of) return self def get(self, where: Combinable) -> Get[TableInstance]: @@ -331,7 +349,7 @@ def default_querystrings(self) -> t.Sequence[QueryString]: "offset_delegate", "output_delegate", "order_by_delegate", - "for_update_delegate", + "lock_for_delegate", ): setattr(select, attr, getattr(self, attr)) diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py index 2e4ea4922..4d5788938 100644 --- a/piccolo/query/methods/select.py +++ b/piccolo/query/methods/select.py @@ -17,9 +17,10 @@ CallbackType, ColumnsDelegate, DistinctDelegate, - ForUpdateDelegate, GroupByDelegate, LimitDelegate, + LockForDelegate, + LockStrength, OffsetDelegate, OrderByDelegate, OrderByRaw, @@ -151,7 +152,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]): "output_delegate", "callback_delegate", "where_delegate", - "for_update_delegate", + "lock_for_delegate", ) def __init__( @@ -176,7 +177,7 @@ def __init__( self.output_delegate = OutputDelegate() self.callback_delegate = CallbackDelegate() self.where_delegate = WhereDelegate() - self.for_update_delegate = ForUpdateDelegate() + self.lock_for_delegate = LockForDelegate() self.columns(*columns_list) @@ -222,10 +223,26 @@ def offset(self: Self, number: int) -> Self: self.offset_delegate.offset(number) return self - def for_update( - self: Self, nowait: bool = False, skip_locked: bool = False, of=() + def lock_for( + self: Self, + lock_strength: t.Union[ + LockStrength, + t.Literal[ + "UPDATE", + "NO KEY UPDATE", + "KEY SHARE", + "SHARE", + "update", + "no key update", + "key share", + "share", + ], + ] = LockStrength.update, + nowait: bool = False, + skip_locked: bool = False, + of: t.Tuple[type[Table], ...] = (), ) -> Self: - self.for_update_delegate.for_update(nowait, skip_locked, of) + self.lock_for_delegate.lock_for(lock_strength, nowait, skip_locked, of) return self async def _splice_m2m_rows( @@ -627,13 +644,13 @@ def default_querystrings(self) -> t.Sequence[QueryString]: query += "{}" args.append(self.offset_delegate._offset.querystring) - if engine_type == "sqlite" and self.for_update_delegate._for_update: + if engine_type == "sqlite" and self.lock_for_delegate._lock_for: raise NotImplementedError( "SQLite doesn't support SELECT .. FOR UPDATE" ) - if self.for_update_delegate._for_update: - args.append(self.for_update_delegate._for_update.querystring) + if self.lock_for_delegate._lock_for: + args.append(self.lock_for_delegate._lock_for.querystring) querystring = QueryString(query, *args) diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py index bffa7c7e4..b11946cb7 100644 --- a/piccolo/query/mixins.py +++ b/piccolo/query/mixins.py @@ -786,15 +786,31 @@ def on_conflict( ) +class LockStrength(str, Enum): + """ + Specify lock strength + + https://www.postgresql.org/docs/current/sql-select.html#SQL-FOR-UPDATE-SHARE + """ + + update = "UPDATE" + no_key_update = "NO KEY UPDATE" + share = "SHARE" + key_share = "KEY SHARE" + + @dataclass -class ForUpdate: - __slots__ = ("nowait", "skip_locked", "of") +class LockFor: + __slots__ = ("lock_strength", "nowait", "skip_locked", "of") + lock_strength: LockStrength nowait: bool skip_locked: bool - of: tuple[Table] + of: tuple[type[Table], ...] def __post_init__(self): + if not isinstance(self.lock_strength, LockStrength): + raise TypeError("lock_strength must be a LockStrength") if not isinstance(self.nowait, bool): raise TypeError("nowait must be a bool") if not isinstance(self.skip_locked, bool): @@ -810,7 +826,7 @@ def __post_init__(self): @property def querystring(self) -> QueryString: - sql = " FOR UPDATE" + sql = f" FOR {self.lock_strength.value}" if self.of: tables = ", ".join(x._meta.tablename for x in self.of) sql += " OF " + tables @@ -826,9 +842,35 @@ def __str__(self) -> str: @dataclass -class ForUpdateDelegate: +class LockForDelegate: + + _lock_for: t.Optional[LockFor] = None - _for_update: t.Optional[ForUpdate] = None + def lock_for( + self, + lock_strength: t.Union[ + LockStrength, + t.Literal[ + "UPDATE", + "NO KEY UPDATE", + "KEY SHARE", + "SHARE", + "update", + "no key update", + "key share", + "share", + ], + ] = LockStrength.update, + nowait=False, + skip_locked=False, + of: t.Tuple[type[Table], ...] = (), + ): + lock_strength_: LockStrength + if isinstance(lock_strength, LockStrength): + lock_strength_ = lock_strength + elif isinstance(lock_strength, str): + lock_strength_ = LockStrength(lock_strength.upper()) + else: + raise ValueError("Unrecognised `lock_strength` value.") - def for_update(self, nowait=False, skip_locked=False, of=()): - self._for_update = ForUpdate(nowait, skip_locked, of) + self._lock_for = LockFor(lock_strength_, nowait, skip_locked, of) diff --git a/tests/table/test_select.py b/tests/table/test_select.py index 050614798..fa99c6181 100644 --- a/tests/table/test_select.py +++ b/tests/table/test_select.py @@ -1032,7 +1032,7 @@ def test_select_raw(self): is_running_sqlite(), reason="SQLite doesn't support SELECT .. FOR UPDATE.", ) - def test_for_update(self): + def test_lock_for(self): """ Make sure the for_update clause works. """ @@ -1041,20 +1041,23 @@ def test_for_update(self): query = Band.select() self.assertNotIn("FOR UPDATE", query.__str__()) - query = query.for_update() + query = query.lock_for() self.assertTrue(query.__str__().endswith("FOR UPDATE")) - query = query.for_update(skip_locked=True) + query = query.lock_for(lock_strength='key share') + self.assertTrue(query.__str__().endswith('FOR KEY SHARE')) + + query = query.lock_for(skip_locked=True) self.assertTrue(query.__str__().endswith("FOR UPDATE SKIP LOCKED")) - query = query.for_update(nowait=True) + query = query.lock_for(nowait=True) self.assertTrue(query.__str__().endswith("FOR UPDATE NOWAIT")) - query = query.for_update(of=(Band,)) + query = query.lock_for(of=(Band,)) self.assertTrue(query.__str__().endswith("FOR UPDATE OF band")) with self.assertRaises(TypeError): - query = query.for_update(skip_locked=True, nowait=True) + query = query.lock_for(skip_locked=True, nowait=True) response = query.run_sync() assert response is not None From ce779e1574da36db823a327c9b3f6966d8cfa658 Mon Sep 17 00:00:00 2001 From: Denis Kopitsa Date: Thu, 19 Sep 2024 18:54:33 +0300 Subject: [PATCH 5/7] Format 'test_select' --- docs/src/piccolo/query_clauses/lock_for.rst | 1 - tests/table/test_select.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/src/piccolo/query_clauses/lock_for.rst b/docs/src/piccolo/query_clauses/lock_for.rst index c209748b5..f7a855e95 100644 --- a/docs/src/piccolo/query_clauses/lock_for.rst +++ b/docs/src/piccolo/query_clauses/lock_for.rst @@ -53,7 +53,6 @@ Which is equivalent to: SELECT ... FOR SHARE - nowait ------ diff --git a/tests/table/test_select.py b/tests/table/test_select.py index fa99c6181..0e2ac2396 100644 --- a/tests/table/test_select.py +++ b/tests/table/test_select.py @@ -1044,8 +1044,8 @@ def test_lock_for(self): query = query.lock_for() self.assertTrue(query.__str__().endswith("FOR UPDATE")) - query = query.lock_for(lock_strength='key share') - self.assertTrue(query.__str__().endswith('FOR KEY SHARE')) + query = query.lock_for(lock_strength="key share") + self.assertTrue(query.__str__().endswith("FOR KEY SHARE")) query = query.lock_for(skip_locked=True) self.assertTrue(query.__str__().endswith("FOR UPDATE SKIP LOCKED")) From 15226339d4efdca6e48b3734ebc8eb0c4d43514b Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Mon, 23 Sep 2024 16:02:18 +0100 Subject: [PATCH 6/7] wip doc changes --- docs/src/piccolo/query_clauses/index.rst | 2 +- docs/src/piccolo/query_clauses/lock_for.rst | 19 +++++++++++-------- docs/src/piccolo/query_types/objects.rst | 5 +++++ docs/src/piccolo/query_types/select.rst | 6 ++++++ 4 files changed, 23 insertions(+), 9 deletions(-) diff --git a/docs/src/piccolo/query_clauses/index.rst b/docs/src/piccolo/query_clauses/index.rst index 60a539cfc..553f5f96d 100644 --- a/docs/src/piccolo/query_clauses/index.rst +++ b/docs/src/piccolo/query_clauses/index.rst @@ -25,11 +25,11 @@ by modifying the return values. ./distinct ./freeze ./group_by + ./lock_for ./offset ./on_conflict ./output ./returning - ./lock_for .. toctree:: :maxdepth: 1 diff --git a/docs/src/piccolo/query_clauses/lock_for.rst b/docs/src/piccolo/query_clauses/lock_for.rst index f7a855e95..ef05f1b27 100644 --- a/docs/src/piccolo/query_clauses/lock_for.rst +++ b/docs/src/piccolo/query_clauses/lock_for.rst @@ -8,13 +8,17 @@ You can use ``lock_for`` clauses with the following queries: * :ref:`Objects` * :ref:`Select` -Returns a query that locks rows until the end of the transaction, generating a SELECT ... FOR UPDATE SQL statement or -similar with other lock strengths. +Returns a query that locks rows until the end of the transaction, generating a +``SELECT ... FOR UPDATE`` SQL statement or similar with other lock strengths. .. note:: Postgres and CockroachDB only. ------------------------------------------------------------------------------- +Basic Usage +----------- + + Basic usage without parameters: .. code-block:: python @@ -34,11 +38,10 @@ lock_strength The parameter ``lock_strength`` controls the strength of the row lock when performing an operation in PostgreSQL. The value can be a predefined constant from the ``LockStrength`` enum or one of the following strings (case-insensitive): - - ``UPDATE`` (default): Acquires an exclusive lock on the selected rows, preventing other transactions from modifying or locking them until the current transaction is complete. - - ``NO KEY UPDATE`` (Postgres only): Similar to UPDATE, but allows other transactions to insert or delete rows that do not affect the primary key or unique constraints. - - ``KEY SHARE`` (Postgres only): Permits other transactions to acquire key-share or share locks, allowing non-key modifications while preventing updates or deletes. - - ``SHARE``: Acquires a shared lock, allowing other transactions to read the rows but not modify or lock them. - +* ``UPDATE`` (default): Acquires an exclusive lock on the selected rows, preventing other transactions from modifying or locking them until the current transaction is complete. +* ``NO KEY UPDATE`` (Postgres only): Similar to ``UPDATE``, but allows other transactions to insert or delete rows that do not affect the primary key or unique constraints. +* ``KEY SHARE`` (Postgres only): Permits other transactions to acquire key-share or share locks, allowing non-key modifications while preventing updates or deletes. +* ``SHARE``: Acquires a shared lock, allowing other transactions to read the rows but not modify or lock them. You can specify a different lock strength: @@ -78,7 +81,7 @@ Ignore locked rows. of -- -By default, if there are many tables in a query (e.g., when joining), all tables will be locked. +By default, if there are many tables in a query (e.g. when joining), all tables will be locked. Using ``of``, you can specify which tables should be locked. .. code-block:: python diff --git a/docs/src/piccolo/query_types/objects.rst b/docs/src/piccolo/query_types/objects.rst index 968aefcab..4133d4113 100644 --- a/docs/src/piccolo/query_types/objects.rst +++ b/docs/src/piccolo/query_types/objects.rst @@ -335,6 +335,11 @@ limit See :ref:`limit`. +lock_for +~~~~~~~~ + +See :ref:`lock_for`. + offset ~~~~~~ diff --git a/docs/src/piccolo/query_types/select.rst b/docs/src/piccolo/query_types/select.rst index 092291e4c..ef6c7bcba 100644 --- a/docs/src/piccolo/query_types/select.rst +++ b/docs/src/piccolo/query_types/select.rst @@ -376,6 +376,12 @@ limit See :ref:`limit`. + +lock_for +~~~~~~~~ + +See :ref:`lock_for`. + offset ~~~~~~ From da2316991e856f9c3af6b9fdf4a5752de0b1f16a Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Tue, 24 Sep 2024 00:29:58 +0100 Subject: [PATCH 7/7] rename `lock_for` to `lock_rows` --- docs/src/piccolo/query_clauses/index.rst | 2 +- .../{lock_for.rst => lock_rows.rst} | 23 +++++++------- docs/src/piccolo/query_types/objects.rst | 4 +-- docs/src/piccolo/query_types/select.rst | 4 +-- piccolo/query/methods/objects.py | 18 +++++------ piccolo/query/methods/select.py | 30 +++++++++---------- piccolo/query/mixins.py | 20 ++++++------- tests/table/test_select.py | 18 +++++------ 8 files changed, 56 insertions(+), 63 deletions(-) rename docs/src/piccolo/query_clauses/{lock_for.rst => lock_rows.rst} (78%) diff --git a/docs/src/piccolo/query_clauses/index.rst b/docs/src/piccolo/query_clauses/index.rst index 553f5f96d..feac07b0f 100644 --- a/docs/src/piccolo/query_clauses/index.rst +++ b/docs/src/piccolo/query_clauses/index.rst @@ -25,7 +25,7 @@ by modifying the return values. ./distinct ./freeze ./group_by - ./lock_for + ./lock_rows ./offset ./on_conflict ./output diff --git a/docs/src/piccolo/query_clauses/lock_for.rst b/docs/src/piccolo/query_clauses/lock_rows.rst similarity index 78% rename from docs/src/piccolo/query_clauses/lock_for.rst rename to docs/src/piccolo/query_clauses/lock_rows.rst index ef05f1b27..54d435326 100644 --- a/docs/src/piccolo/query_clauses/lock_for.rst +++ b/docs/src/piccolo/query_clauses/lock_rows.rst @@ -1,15 +1,14 @@ -.. _lock_for: +.. _lock_rows: -lock_for -======== +lock_rows +========= -You can use ``lock_for`` clauses with the following queries: +You can use the ``lock_rows`` clause with the following queries: * :ref:`Objects` * :ref:`Select` -Returns a query that locks rows until the end of the transaction, generating a -``SELECT ... FOR UPDATE`` SQL statement or similar with other lock strengths. +It returns a query that locks rows until the end of the transaction, generating a ``SELECT ... FOR UPDATE`` SQL statement or similar with other lock strengths. .. note:: Postgres and CockroachDB only. @@ -18,12 +17,11 @@ Returns a query that locks rows until the end of the transaction, generating a Basic Usage ----------- - Basic usage without parameters: .. code-block:: python - await Band.select(Band.name == 'Pythonistas').lock_for() + await Band.select(Band.name == 'Pythonistas').lock_rows() Equivalent to: @@ -47,7 +45,7 @@ You can specify a different lock strength: .. code-block:: python - await Band.select(Band.name == 'Pythonistas').lock_for('share') + await Band.select(Band.name == 'Pythonistas').lock_rows('SHARE') Which is equivalent to: @@ -64,7 +62,7 @@ waiting for the other transaction to release the lock. .. code-block:: python - await Band.select(Band.name == 'Pythonistas').lock_for('update', nowait=True) + await Band.select(Band.name == 'Pythonistas').lock_rows('UPDATE', nowait=True) skip_locked @@ -74,8 +72,7 @@ Ignore locked rows. .. code-block:: python - await Band.select(Band.name == 'Pythonistas').lock_for('update', skip_locked=True) - + await Band.select(Band.name == 'Pythonistas').lock_rows('UPDATE', skip_locked=True) of @@ -86,7 +83,7 @@ Using ``of``, you can specify which tables should be locked. .. code-block:: python - await Band.select().where(Band.manager.name == 'Guido').lock_for('update', of=(Band, )) + await Band.select().where(Band.manager.name == 'Guido').lock_rows('UPDATE', of=(Band, )) Learn more diff --git a/docs/src/piccolo/query_types/objects.rst b/docs/src/piccolo/query_types/objects.rst index 4133d4113..66acf9413 100644 --- a/docs/src/piccolo/query_types/objects.rst +++ b/docs/src/piccolo/query_types/objects.rst @@ -335,10 +335,10 @@ limit See :ref:`limit`. -lock_for +lock_rows ~~~~~~~~ -See :ref:`lock_for`. +See :ref:`lock_rows`. offset ~~~~~~ diff --git a/docs/src/piccolo/query_types/select.rst b/docs/src/piccolo/query_types/select.rst index ef6c7bcba..96c8b06d0 100644 --- a/docs/src/piccolo/query_types/select.rst +++ b/docs/src/piccolo/query_types/select.rst @@ -377,10 +377,10 @@ limit See :ref:`limit`. -lock_for +lock_rows ~~~~~~~~ -See :ref:`lock_for`. +See :ref:`lock_rows`. offset ~~~~~~ diff --git a/piccolo/query/methods/objects.py b/piccolo/query/methods/objects.py index 8cd66b408..41d105b9f 100644 --- a/piccolo/query/methods/objects.py +++ b/piccolo/query/methods/objects.py @@ -13,7 +13,7 @@ CallbackDelegate, CallbackType, LimitDelegate, - LockForDelegate, + LockRowsDelegate, LockStrength, OffsetDelegate, OrderByDelegate, @@ -197,7 +197,7 @@ class Objects( "callback_delegate", "prefetch_delegate", "where_delegate", - "lock_for_delegate", + "lock_rows_delegate", ) def __init__( @@ -217,7 +217,7 @@ def __init__( self.prefetch_delegate = PrefetchDelegate() self.prefetch(*prefetch) self.where_delegate = WhereDelegate() - self.lock_for_delegate = LockForDelegate() + self.lock_rows_delegate = LockRowsDelegate() def output(self: Self, load_json: bool = False) -> Self: self.output_delegate.output( @@ -277,7 +277,7 @@ def first(self) -> First[TableInstance]: self.limit_delegate.limit(1) return First[TableInstance](query=self) - def lock_for( + def lock_rows( self: Self, lock_strength: t.Union[ LockStrength, @@ -286,17 +286,15 @@ def lock_for( "NO KEY UPDATE", "KEY SHARE", "SHARE", - "update", - "no key update", - "key share", - "share", ], ] = LockStrength.update, nowait: bool = False, skip_locked: bool = False, of: t.Tuple[type[Table], ...] = (), ) -> Self: - self.lock_for_delegate.lock_for(lock_strength, nowait, skip_locked, of) + self.lock_rows_delegate.lock_rows( + lock_strength, nowait, skip_locked, of + ) return self def get(self, where: Combinable) -> Get[TableInstance]: @@ -349,7 +347,7 @@ def default_querystrings(self) -> t.Sequence[QueryString]: "offset_delegate", "output_delegate", "order_by_delegate", - "lock_for_delegate", + "lock_rows_delegate", ): setattr(select, attr, getattr(self, attr)) diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py index 4d5788938..5d7856c5a 100644 --- a/piccolo/query/methods/select.py +++ b/piccolo/query/methods/select.py @@ -19,7 +19,7 @@ DistinctDelegate, GroupByDelegate, LimitDelegate, - LockForDelegate, + LockRowsDelegate, LockStrength, OffsetDelegate, OrderByDelegate, @@ -152,7 +152,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]): "output_delegate", "callback_delegate", "where_delegate", - "lock_for_delegate", + "lock_rows_delegate", ) def __init__( @@ -177,7 +177,7 @@ def __init__( self.output_delegate = OutputDelegate() self.callback_delegate = CallbackDelegate() self.where_delegate = WhereDelegate() - self.lock_for_delegate = LockForDelegate() + self.lock_rows_delegate = LockRowsDelegate() self.columns(*columns_list) @@ -223,7 +223,7 @@ def offset(self: Self, number: int) -> Self: self.offset_delegate.offset(number) return self - def lock_for( + def lock_rows( self: Self, lock_strength: t.Union[ LockStrength, @@ -232,17 +232,15 @@ def lock_for( "NO KEY UPDATE", "KEY SHARE", "SHARE", - "update", - "no key update", - "key share", - "share", ], ] = LockStrength.update, nowait: bool = False, skip_locked: bool = False, of: t.Tuple[type[Table], ...] = (), ) -> Self: - self.lock_for_delegate.lock_for(lock_strength, nowait, skip_locked, of) + self.lock_rows_delegate.lock_rows( + lock_strength, nowait, skip_locked, of + ) return self async def _splice_m2m_rows( @@ -644,13 +642,15 @@ def default_querystrings(self) -> t.Sequence[QueryString]: query += "{}" args.append(self.offset_delegate._offset.querystring) - if engine_type == "sqlite" and self.lock_for_delegate._lock_for: - raise NotImplementedError( - "SQLite doesn't support SELECT .. FOR UPDATE" - ) + if self.lock_rows_delegate._lock_rows: + if engine_type == "sqlite": + raise NotImplementedError( + "SQLite doesn't support row locking e.g. SELECT ... FOR " + "UPDATE" + ) - if self.lock_for_delegate._lock_for: - args.append(self.lock_for_delegate._lock_for.querystring) + query += "{}" + args.append(self.lock_rows_delegate._lock_rows.querystring) querystring = QueryString(query, *args) diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py index b11946cb7..b1b726cf3 100644 --- a/piccolo/query/mixins.py +++ b/piccolo/query/mixins.py @@ -800,13 +800,13 @@ class LockStrength(str, Enum): @dataclass -class LockFor: +class LockRows: __slots__ = ("lock_strength", "nowait", "skip_locked", "of") lock_strength: LockStrength nowait: bool skip_locked: bool - of: tuple[type[Table], ...] + of: t.Tuple[t.Type[Table], ...] def __post_init__(self): if not isinstance(self.lock_strength, LockStrength): @@ -828,7 +828,9 @@ def __post_init__(self): def querystring(self) -> QueryString: sql = f" FOR {self.lock_strength.value}" if self.of: - tables = ", ".join(x._meta.tablename for x in self.of) + tables = ", ".join( + i._meta.get_formatted_tablename() for i in self.of + ) sql += " OF " + tables if self.nowait: sql += " NOWAIT" @@ -842,11 +844,11 @@ def __str__(self) -> str: @dataclass -class LockForDelegate: +class LockRowsDelegate: - _lock_for: t.Optional[LockFor] = None + _lock_rows: t.Optional[LockRows] = None - def lock_for( + def lock_rows( self, lock_strength: t.Union[ LockStrength, @@ -855,10 +857,6 @@ def lock_for( "NO KEY UPDATE", "KEY SHARE", "SHARE", - "update", - "no key update", - "key share", - "share", ], ] = LockStrength.update, nowait=False, @@ -873,4 +871,4 @@ def lock_for( else: raise ValueError("Unrecognised `lock_strength` value.") - self._lock_for = LockFor(lock_strength_, nowait, skip_locked, of) + self._lock_rows = LockRows(lock_strength_, nowait, skip_locked, of) diff --git a/tests/table/test_select.py b/tests/table/test_select.py index 0e2ac2396..74b876d5f 100644 --- a/tests/table/test_select.py +++ b/tests/table/test_select.py @@ -1030,9 +1030,9 @@ def test_select_raw(self): @pytest.mark.skipif( is_running_sqlite(), - reason="SQLite doesn't support SELECT .. FOR UPDATE.", + reason="SQLite doesn't support SELECT ... FOR UPDATE.", ) - def test_lock_for(self): + def test_lock_rows(self): """ Make sure the for_update clause works. """ @@ -1041,23 +1041,23 @@ def test_lock_for(self): query = Band.select() self.assertNotIn("FOR UPDATE", query.__str__()) - query = query.lock_for() + query = query.lock_rows() self.assertTrue(query.__str__().endswith("FOR UPDATE")) - query = query.lock_for(lock_strength="key share") + query = query.lock_rows(lock_strength="KEY SHARE") self.assertTrue(query.__str__().endswith("FOR KEY SHARE")) - query = query.lock_for(skip_locked=True) + query = query.lock_rows(skip_locked=True) self.assertTrue(query.__str__().endswith("FOR UPDATE SKIP LOCKED")) - query = query.lock_for(nowait=True) + query = query.lock_rows(nowait=True) self.assertTrue(query.__str__().endswith("FOR UPDATE NOWAIT")) - query = query.lock_for(of=(Band,)) - self.assertTrue(query.__str__().endswith("FOR UPDATE OF band")) + query = query.lock_rows(of=(Band,)) + self.assertTrue(query.__str__().endswith('FOR UPDATE OF "band"')) with self.assertRaises(TypeError): - query = query.lock_for(skip_locked=True, nowait=True) + query = query.lock_rows(skip_locked=True, nowait=True) response = query.run_sync() assert response is not None