Skip to content

Commit

Permalink
rename lock_for to lock_rows
Browse files Browse the repository at this point in the history
  • Loading branch information
dantownsend committed Sep 23, 2024
1 parent 1522633 commit da23169
Show file tree
Hide file tree
Showing 8 changed files with 56 additions and 63 deletions.
2 changes: 1 addition & 1 deletion docs/src/piccolo/query_clauses/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ by modifying the return values.
./distinct
./freeze
./group_by
./lock_for
./lock_rows
./offset
./on_conflict
./output
Expand Down
Original file line number Diff line number Diff line change
@@ -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.

Expand All @@ -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:

Expand All @@ -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:

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions docs/src/piccolo/query_types/objects.rst
Original file line number Diff line number Diff line change
Expand Up @@ -335,10 +335,10 @@ limit

See :ref:`limit`.

lock_for
lock_rows
~~~~~~~~

See :ref:`lock_for`.
See :ref:`lock_rows`.

offset
~~~~~~
Expand Down
4 changes: 2 additions & 2 deletions docs/src/piccolo/query_types/select.rst
Original file line number Diff line number Diff line change
Expand Up @@ -377,10 +377,10 @@ limit
See :ref:`limit`.


lock_for
lock_rows
~~~~~~~~

See :ref:`lock_for`.
See :ref:`lock_rows`.

offset
~~~~~~
Expand Down
18 changes: 8 additions & 10 deletions piccolo/query/methods/objects.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
CallbackDelegate,
CallbackType,
LimitDelegate,
LockForDelegate,
LockRowsDelegate,
LockStrength,
OffsetDelegate,
OrderByDelegate,
Expand Down Expand Up @@ -197,7 +197,7 @@ class Objects(
"callback_delegate",
"prefetch_delegate",
"where_delegate",
"lock_for_delegate",
"lock_rows_delegate",
)

def __init__(
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand All @@ -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]:
Expand Down Expand Up @@ -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))

Expand Down
30 changes: 15 additions & 15 deletions piccolo/query/methods/select.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
DistinctDelegate,
GroupByDelegate,
LimitDelegate,
LockForDelegate,
LockRowsDelegate,
LockStrength,
OffsetDelegate,
OrderByDelegate,
Expand Down Expand Up @@ -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__(
Expand All @@ -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)

Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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)

Expand Down
20 changes: 9 additions & 11 deletions piccolo/query/mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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"
Expand All @@ -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,
Expand All @@ -855,10 +857,6 @@ def lock_for(
"NO KEY UPDATE",
"KEY SHARE",
"SHARE",
"update",
"no key update",
"key share",
"share",
],
] = LockStrength.update,
nowait=False,
Expand All @@ -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)
18 changes: 9 additions & 9 deletions tests/table/test_select.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand All @@ -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
Expand Down

0 comments on commit da23169

Please sign in to comment.