piccolo 1.17.1__py3-none-any.whl → 1.19.0__py3-none-any.whl

piccolo/__init__.py

@@ -1 +1 @@
-__VERSION__ = "1.17.1"
+__VERSION__ = "1.19.0"

piccolo/query/methods/objects.py

@@ -13,6 +13,8 @@ from piccolo.query.mixins import (
     CallbackDelegate,
     CallbackType,
     LimitDelegate,
+    LockRowsDelegate,
+    LockStrength,
     OffsetDelegate,
     OrderByDelegate,
     OrderByRaw,
@@ -27,6 +29,7 @@ from piccolo.utils.sync import run_sync
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.columns import Column
+    from piccolo.table import Table
 
 
 ###############################################################################
@@ -173,6 +176,61 @@ class Create(t.Generic[TableInstance]):
         return run_sync(self.run(*args, **kwargs))
 
 
+class UpdateSelf:
+
+    def __init__(
+        self,
+        row: Table,
+        values: t.Dict[t.Union[Column, str], t.Any],
+    ):
+        self.row = row
+        self.values = values
+
+    async def run(
+        self,
+        node: t.Optional[str] = None,
+        in_pool: bool = True,
+    ) -> None:
+        if not self.row._exists_in_db:
+            raise ValueError("This row doesn't exist in the database.")
+
+        TableClass = self.row.__class__
+
+        primary_key = TableClass._meta.primary_key
+        primary_key_value = getattr(self.row, primary_key._meta.name)
+
+        if primary_key_value is None:
+            raise ValueError("The primary key is None")
+
+        columns = [
+            TableClass._meta.get_column_by_name(i) if isinstance(i, str) else i
+            for i in self.values.keys()
+        ]
+
+        response = (
+            await TableClass.update(self.values)
+            .where(primary_key == primary_key_value)
+            .returning(*columns)
+            .run(
+                node=node,
+                in_pool=in_pool,
+            )
+        )
+
+        for key, value in response[0].items():
+            setattr(self.row, key, value)
+
+    def __await__(self) -> t.Generator[None, None, None]:
+        """
+        If the user doesn't explicity call .run(), proxy to it as a
+        convenience.
+        """
+        return self.run().__await__()
+
+    def run_sync(self, *args, **kwargs) -> None:
+        return run_sync(self.run(*args, **kwargs))
+
+
 ###############################################################################
 
 
@@ -194,6 +252,7 @@ class Objects(
         "callback_delegate",
         "prefetch_delegate",
         "where_delegate",
+        "lock_rows_delegate",
     )
 
     def __init__(
@@ -213,6 +272,7 @@ class Objects(
         self.prefetch_delegate = PrefetchDelegate()
         self.prefetch(*prefetch)
         self.where_delegate = WhereDelegate()
+        self.lock_rows_delegate = LockRowsDelegate()
 
     def output(self: Self, load_json: bool = False) -> Self:
         self.output_delegate.output(
@@ -272,6 +332,26 @@ class Objects(
         self.limit_delegate.limit(1)
         return First[TableInstance](query=self)
 
+    def lock_rows(
+        self: Self,
+        lock_strength: t.Union[
+            LockStrength,
+            t.Literal[
+                "UPDATE",
+                "NO KEY UPDATE",
+                "KEY SHARE",
+                "SHARE",
+            ],
+        ] = LockStrength.update,
+        nowait: bool = False,
+        skip_locked: bool = False,
+        of: t.Tuple[type[Table], ...] = (),
+    ) -> Self:
+        self.lock_rows_delegate.lock_rows(
+            lock_strength, 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 +402,7 @@ class Objects(
             "offset_delegate",
             "output_delegate",
             "order_by_delegate",
+            "lock_rows_delegate",
         ):
             setattr(select, attr, getattr(self, attr))
 

piccolo/query/methods/raw.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import typing as t
 
+from piccolo.engine.base import BaseBatch
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
 
@@ -21,6 +22,18 @@ class Raw(Query):
         super().__init__(table, **kwargs)
         self.querystring = querystring
 
+    async def batch(
+        self,
+        batch_size: t.Optional[int] = None,
+        node: t.Optional[str] = None,
+        **kwargs,
+    ) -> BaseBatch:
+        if batch_size:
+            kwargs.update(batch_size=batch_size)
+        if node:
+            kwargs.update(node=node)
+        return await self.table._meta.db.batch(self, **kwargs)
+
     @property
     def default_querystrings(self) -> t.Sequence[QueryString]:
         return [self.querystring]

piccolo/query/methods/select.py

@@ -19,6 +19,8 @@ from piccolo.query.mixins import (
     DistinctDelegate,
     GroupByDelegate,
     LimitDelegate,
+    LockRowsDelegate,
+    LockStrength,
     OffsetDelegate,
     OrderByDelegate,
     OrderByRaw,
@@ -150,6 +152,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         "output_delegate",
         "callback_delegate",
         "where_delegate",
+        "lock_rows_delegate",
     )
 
     def __init__(
@@ -174,6 +177,7 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         self.output_delegate = OutputDelegate()
         self.callback_delegate = CallbackDelegate()
         self.where_delegate = WhereDelegate()
+        self.lock_rows_delegate = LockRowsDelegate()
 
         self.columns(*columns_list)
 
@@ -219,6 +223,26 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
         self.offset_delegate.offset(number)
         return self
 
+    def lock_rows(
+        self: Self,
+        lock_strength: t.Union[
+            LockStrength,
+            t.Literal[
+                "UPDATE",
+                "NO KEY UPDATE",
+                "KEY SHARE",
+                "SHARE",
+            ],
+        ] = LockStrength.update,
+        nowait: bool = False,
+        skip_locked: bool = False,
+        of: t.Tuple[type[Table], ...] = (),
+    ) -> Self:
+        self.lock_rows_delegate.lock_rows(
+            lock_strength, nowait, skip_locked, of
+        )
+        return self
+
     async def _splice_m2m_rows(
         self,
         response: t.List[t.Dict[str, t.Any]],
@@ -618,6 +642,16 @@ class Select(Query[TableInstance, t.List[t.Dict[str, t.Any]]]):
             query += "{}"
             args.append(self.offset_delegate._offset.querystring)
 
+        if self.lock_rows_delegate._lock_rows:
+            if engine_type == "sqlite":
+                raise NotImplementedError(
+                    "SQLite doesn't support row locking e.g. SELECT ... FOR "
+                    "UPDATE"
+                )
+
+            query += "{}"
+            args.append(self.lock_rows_delegate._lock_rows.querystring)
+
         querystring = QueryString(query, *args)
 
         return [querystring]

piccolo/query/mixins.py

@@ -784,3 +784,91 @@ class OnConflictDelegate:
                 target=target, action=action_, values=values, where=where
             )
         )
+
+
+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 LockRows:
+    __slots__ = ("lock_strength", "nowait", "skip_locked", "of")
+
+    lock_strength: LockStrength
+    nowait: bool
+    skip_locked: bool
+    of: t.Tuple[t.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):
+            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 a 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 = f" FOR {self.lock_strength.value}"
+        if self.of:
+            tables = ", ".join(
+                i._meta.get_formatted_tablename() for i 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 LockRowsDelegate:
+
+    _lock_rows: t.Optional[LockRows] = None
+
+    def lock_rows(
+        self,
+        lock_strength: t.Union[
+            LockStrength,
+            t.Literal[
+                "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.")
+
+        self._lock_rows = LockRows(lock_strength_, nowait, skip_locked, of)

piccolo/table.py

@@ -46,7 +46,7 @@ from piccolo.query import (
 )
 from piccolo.query.methods.create_index import CreateIndex
 from piccolo.query.methods.indexes import Indexes
-from piccolo.query.methods.objects import First
+from piccolo.query.methods.objects import First, UpdateSelf
 from piccolo.query.methods.refresh import Refresh
 from piccolo.querystring import QueryString
 from piccolo.utils import _camel_to_snake
@@ -525,6 +525,43 @@ class Table(metaclass=TableMetaclass):
             == getattr(self, self._meta.primary_key._meta.name)
         )
 
+    def update_self(
+        self, values: t.Dict[t.Union[Column, str], t.Any]
+    ) -> UpdateSelf:
+        """
+        This allows the user to update a single object - useful when the values
+        are derived from the database in some way.
+
+        For example, if we have the following table::
+
+            class Band(Table):
+                name = Varchar()
+                popularity = Integer()
+
+        And we fetch an object::
+
+            >>> band = await Band.objects().get(name="Pythonistas")
+
+        We could use the typical syntax for updating the object::
+
+            >>> band.popularity += 1
+            >>> await band.save()
+
+        The problem with this, is what if another object has already
+        incremented ``popularity``? It would overide the value.
+
+        Instead we can do this:
+
+            >>> await band.update_self({
+            ...     Band.popularity: Band.popularity + 1
+            ... })
+
+        This updates ``popularity`` in the database, and also sets the new
+        value for ``popularity`` on the object.
+
+        """
+        return UpdateSelf(row=self, values=values)
+
     def remove(self) -> Delete:
         """
         A proxy to a delete query.

{piccolo-1.17.1.dist-info → piccolo-1.19.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: piccolo
-Version: 1.17.1
+Version: 1.19.0
 Summary: A fast, user friendly ORM and query builder which supports asyncio.
 Home-page: https://github.com/piccolo-orm/piccolo
 Author: Daniel Townsend

{piccolo-1.17.1.dist-info → piccolo-1.19.0.dist-info}/RECORD

@@ -1,10 +1,10 @@
-piccolo/__init__.py,sha256=LHM-IUVB0dp8WEkvdMoZPznVSxU-UEwiRDYV3iBGA4w,23
+piccolo/__init__.py,sha256=P8Wjra0Mf_3SPYYqkD2tAWe0dnPdqTSizVJGqA8al_Q,23
 piccolo/custom_types.py,sha256=7HMQAze-5mieNLfbQ5QgbRQgR2abR7ol0qehv2SqROY,604
 piccolo/main.py,sha256=1VsFV67FWTUikPTysp64Fmgd9QBVa_9wcwKfwj2UCEA,5117
 piccolo/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 piccolo/querystring.py,sha256=yZdURtiVSlxkkEVJoFKAmL2OfNxrUr8xfsfuBBB7IuY,9662
 piccolo/schema.py,sha256=qNNy4tG_HqnXR9t3hHMgYXtGxHabwQAhUpc6RKLJ_gE,7960
-piccolo/table.py,sha256=nS3zuhGNPZ4H9s_E3ieFbrE_u1Tr6TenBVw_nrrtmRQ,49766
+piccolo/table.py,sha256=rwLkom0-4ahOTVon9JDoZ92cp5d3nuwo24ES8attxCM,50896
 piccolo/table_reflection.py,sha256=jrN1nHerDJ4tU09GtNN3hz7ap-7rXnSUjljFO6LB2H0,7094
 piccolo/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 piccolo/apps/app/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -147,7 +147,7 @@ piccolo/engine/postgres.py,sha256=DekL3KafCdzSAEQ6_EgOiUB1ERXh2xpePYwI9QvmN-c,18
 piccolo/engine/sqlite.py,sha256=KwJc3UttBP_8qSREbLJshqEfROF17ENf0Ju9BwI5_so,25236
 piccolo/query/__init__.py,sha256=bcsMV4813rMRAIqGv4DxI4eyO4FmpXkDv9dfTk5pt3A,699
 piccolo/query/base.py,sha256=iI9Fv3oOw7T4ZWZvRKRwdtClvQtSaAepslH24vwxZVA,14616
-piccolo/query/mixins.py,sha256=EFEFb9It4y1mR6_JXLn139h5M9KgeP750STYy5M4MLs,21951
+piccolo/query/mixins.py,sha256=X9HEYnj6uOjgTkGr4vgqTwN_dokJPzVagwbFx385atQ,24468
 piccolo/query/proxy.py,sha256=Yq4jNc7IWJvdeO3u7_7iPyRy2WhVj8KsIUcIYHBIi9Q,1839
 piccolo/query/functions/__init__.py,sha256=pZkzOIh7Sg9HPNOeegOwAS46Oxt31ATlSVmwn-lxCbc,605
 piccolo/query/functions/aggregate.py,sha256=OdjDjr_zyD4S9UbrZ2C3V5mz4OT2sIfAFAdTGr4WL54,4248
@@ -166,10 +166,10 @@ piccolo/query/methods/drop_index.py,sha256=5x3vHpoOmQ1SMhj6L7snKXX6M9l9j1E1PFSO6
 piccolo/query/methods/exists.py,sha256=lTMjtrFPFygZmaPV3sfQKXc3K0sVqJ2S6PDc3fRK6YQ,1203
 piccolo/query/methods/indexes.py,sha256=J-QUqaBJwpgahskUH0Cu0Mq7zEKcfVAtDsUVIVX-C4c,943
 piccolo/query/methods/insert.py,sha256=ssLJ_wn08KnOwwr7t-VILyn1P4hrvM63CfPIcAJWT5k,4701
-piccolo/query/methods/objects.py,sha256=iahDUziUtlx7pJ2uBAhdm3hCTmg2AS9C8cal1my5KR0,11705
-piccolo/query/methods/raw.py,sha256=VhYpCB52mZk4zqFTsqK5CHKTDGskUjISXTBV7UjohmA,600
+piccolo/query/methods/objects.py,sha256=kWHSgnXFVYe-x6tIHstbZvZic6mQCCGES-L_2Icbg-I,13910
+piccolo/query/methods/raw.py,sha256=wQWR8b-yA_Gr-5lqRMZe9BOAAMBAw8CqTx37qVYvM1A,987
 piccolo/query/methods/refresh.py,sha256=wg1zghKfwz-VmqK4uWa4GNMiDtK-skTqow591Hb3ONM,5854
-piccolo/query/methods/select.py,sha256=UH-y2g3Ub7bEowfLObrrhw0W-HepTXWCmuMPhk13roE,21406
+piccolo/query/methods/select.py,sha256=jNR7CsEPUarffYMsXytKm7iScrQ_nqpRX-5mTPrXfjg,22414
 piccolo/query/methods/table_exists.py,sha256=0yb3n6Jd2ovSBWlZ-gl00K4E7Jnbj7J8qAAX5d7hvNk,1259
 piccolo/query/methods/update.py,sha256=LfWqIXEl1aecc0rkVssTFmwyD6wXGhlKcTrUVhtlEsw,3705
 piccolo/testing/__init__.py,sha256=pRFSqRInfx95AakOq54atmvqoB-ue073q2aR8u8zR40,83
@@ -316,7 +316,7 @@ tests/query/mixins/test_order_by_delegate.py,sha256=mOV3Gxs0XeliONxjWSOniI1z6lbZ
 tests/table/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/table/test_all_columns.py,sha256=wZ7i9mTT9wKWLE2BoQ9jDbPaqnBHfV-ZlsROp7SZq7k,667
 tests/table/test_alter.py,sha256=pMD38BIFfta1vxFqp8YoaRfMxdwxhQSwcxYO4erpUi8,12394
-tests/table/test_batch.py,sha256=l3kSHOT2DQ8FJjbr_bFyMGgQ2BI0Ed5Gs4GeJdEsjVo,3795
+tests/table/test_batch.py,sha256=fhBfC5P_vsk1wlmVLeVCDQ8tz9nrny57aIgRnPJ6tlg,4944
 tests/table/test_callback.py,sha256=mNp4NOQ4PM_PIKn-WLG4HyvN-YHk7w0Ka_Bljg5Ugd4,6381
 tests/table/test_constructor.py,sha256=zPbzhKQWzzsQQUK7P9WM8OgCi1ndhXedP6rU0tg6XIM,832
 tests/table/test_count.py,sha256=qm4dwlQJ5gv8FPSsgYTS-3Gsd_KLgvWlFnmXweKydxw,2297
@@ -339,10 +339,11 @@ tests/table/test_raw.py,sha256=9PTvYngQi41nYd5lKzkJdTqsEcwrdOXcvZjq-W26CwQ,1683
 tests/table/test_ref.py,sha256=eYNRnYHzNMXuMbV3B1ca5EidpIg4500q6hr1ccuVaso,269
 tests/table/test_refresh.py,sha256=-BaLS6fZiR2RtQaFa7D9WGBjrbrss1-tt5xz1NE_m8E,9250
 tests/table/test_repr.py,sha256=uahz3_GffGQrf2mDE-4-Pu4AmSLBAyso6-9rbohCl58,446
-tests/table/test_select.py,sha256=jgeiahIlNFVijxYb3a54g1sJWVfH3llaYrsTBmdicrs,40390
+tests/table/test_select.py,sha256=PcLRtHGEZfYNB0Qm2b9P6NFJOH2lWvviRjsr1cEmsMQ,41503
 tests/table/test_str.py,sha256=eztWNULcjARR1fr9X5n4tojhDNgDfatVyNHwuYrzHAo,1731
 tests/table/test_table_exists.py,sha256=upv2e9UD32V2QZOShzmcw0reMqRbYiX_jxWx57p25jg,1082
 tests/table/test_update.py,sha256=Cqi0xX3kEuJ0k-x_emPGB3arXuGWZ9e3CJ3HPFnw9Zw,20505
+tests/table/test_update_self.py,sha256=im6HcM-WLkEhZP0vTL42tYEJZyAZG6gDOxCnbikCBD4,907
 tests/table/instance/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/table/instance/test_create.py,sha256=JD0l7L9dDK1FKPhUs6WC_B2bruPR1qQ8aIqXpEbfiUg,1105
 tests/table/instance/test_get_related.py,sha256=dAT89KHuO_XEMY6h0ngDHrXp_nJsMoSToTbH_IRclTE,1275
@@ -367,9 +368,9 @@ tests/utils/test_sql_values.py,sha256=vzxRmy16FfLZPH-sAQexBvsF9MXB8n4smr14qoEOS5
 tests/utils/test_sync.py,sha256=9ytVo56y2vPQePvTeIi9lHIouEhWJbodl1TmzkGFrSo,799
 tests/utils/test_table_reflection.py,sha256=SIzuat-IpcVj1GCFyOWKShI8YkhdOPPFH7qVrvfyPNE,3794
 tests/utils/test_warnings.py,sha256=NvSC_cvJ6uZcwAGf1m-hLzETXCqprXELL8zg3TNLVMw,269
-piccolo-1.17.1.dist-info/LICENSE,sha256=zFIpi-16uIJ420UMIG75NU0JbDBykvrdnXcj5U_EYBI,1059
-piccolo-1.17.1.dist-info/METADATA,sha256=TZvCpeVtvGuvBWlyeNvBGPXN7dFZNSrRVuvUwyCOGRs,5178
-piccolo-1.17.1.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-piccolo-1.17.1.dist-info/entry_points.txt,sha256=SJPHET4Fi1bN5F3WqcKkv9SClK3_F1I7m4eQjk6AFh0,46
-piccolo-1.17.1.dist-info/top_level.txt,sha256=-SR74VGbk43VoPy1HH-mHm97yoGukLK87HE5kdBW6qM,24
-piccolo-1.17.1.dist-info/RECORD,,
+piccolo-1.19.0.dist-info/LICENSE,sha256=zFIpi-16uIJ420UMIG75NU0JbDBykvrdnXcj5U_EYBI,1059
+piccolo-1.19.0.dist-info/METADATA,sha256=sjcIxHmO7mxcywJ7dGiQtleQjBW1uIIBL-ld4CJeVoU,5178
+piccolo-1.19.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+piccolo-1.19.0.dist-info/entry_points.txt,sha256=SJPHET4Fi1bN5F3WqcKkv9SClK3_F1I7m4eQjk6AFh0,46
+piccolo-1.19.0.dist-info/top_level.txt,sha256=-SR74VGbk43VoPy1HH-mHm97yoGukLK87HE5kdBW6qM,24
+piccolo-1.19.0.dist-info/RECORD,,

{piccolo-1.17.1.dist-info → piccolo-1.19.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.1.2)
+Generator: setuptools (75.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/table/test_batch.py

@@ -95,6 +95,47 @@ class TestBatchObjects(DBTestCase):
         self.assertEqual(iterations, _iterations)
 
 
+class TestBatchRaw(DBTestCase):
+    def _check_results(self, batch):
+        """
+        Make sure the data is returned in the correct format.
+        """
+        self.assertEqual(type(batch), list)
+        if len(batch) > 0:
+            row = batch[0]
+            self.assertIsInstance(row, Manager)
+
+    async def run_batch(self, batch_size):
+        row_count = 0
+        iterations = 0
+
+        async with await Manager.raw("SELECT * FROM manager").batch(
+            batch_size=batch_size
+        ) as batch:
+            async for _batch in batch:
+                self._check_results(_batch)
+                _row_count = len(_batch)
+                row_count += _row_count
+                iterations += 1
+
+        return row_count, iterations
+
+    async def test_batch(self):
+        row_count = 1000
+        self.insert_many_rows(row_count)
+
+        batch_size = 10
+
+        _row_count, iterations = asyncio.run(
+            self.run_batch(batch_size=batch_size), debug=True
+        )
+
+        _iterations = math.ceil(row_count / batch_size)
+
+        self.assertEqual(_row_count, row_count)
+        self.assertEqual(iterations, _iterations)
+
+
 @engines_only("postgres", "cockroach")
 class TestBatchNodeArg(TestCase):
     def test_batch_extra_node(self):

tests/table/test_select.py

@@ -1028,6 +1028,40 @@ class TestSelect(DBTestCase):
             response, [{"name": "Pythonistas", "popularity_log": 3.0}]
         )
 
+    @pytest.mark.skipif(
+        is_running_sqlite(),
+        reason="SQLite doesn't support SELECT ... FOR UPDATE.",
+    )
+    def test_lock_rows(self):
+        """
+        Make sure the for_update clause works.
+        """
+        self.insert_rows()
+
+        query = Band.select()
+        self.assertNotIn("FOR UPDATE", query.__str__())
+
+        query = query.lock_rows()
+        self.assertTrue(query.__str__().endswith("FOR UPDATE"))
+
+        query = query.lock_rows(lock_strength="KEY SHARE")
+        self.assertTrue(query.__str__().endswith("FOR KEY SHARE"))
+
+        query = query.lock_rows(skip_locked=True)
+        self.assertTrue(query.__str__().endswith("FOR UPDATE SKIP LOCKED"))
+
+        query = query.lock_rows(nowait=True)
+        self.assertTrue(query.__str__().endswith("FOR UPDATE NOWAIT"))
+
+        query = query.lock_rows(of=(Band,))
+        self.assertTrue(query.__str__().endswith('FOR UPDATE OF "band"'))
+
+        with self.assertRaises(TypeError):
+            query = query.lock_rows(skip_locked=True, nowait=True)
+
+        response = query.run_sync()
+        assert response is not None
+
 
 class TestSelectSecret(TestCase):
     def setUp(self):

tests/table/test_update_self.py

@@ -0,0 +1,27 @@
+from piccolo.testing.test_case import AsyncTableTest
+from tests.example_apps.music.tables import Band, Manager
+
+
+class TestUpdateSelf(AsyncTableTest):
+
+    tables = [Band, Manager]
+
+    async def test_update_self(self):
+        band = Band({Band.name: "Pythonistas", Band.popularity: 1000})
+
+        # Make sure we get a ValueError if it's not in the database yet.
+        with self.assertRaises(ValueError):
+            await band.update_self({Band.popularity: Band.popularity + 1})
+
+        # Save it, so it's in the database
+        await band.save()
+
+        # Make sure we can successfully update the object
+        await band.update_self({Band.popularity: Band.popularity + 1})
+
+        # Make sure the value was updated on the object
+        assert band.popularity == 1001
+
+        # Make sure the value was updated in the database
+        await band.refresh()
+        assert band.popularity == 1001