1073 Add update_self method (#1081)

* prototype for `UpdateSelf`

* fleshed out implementation

* add tests

* update docstring - use `Band` table as an example

* improve docs

* finish docs

Author: dantownsend

docs/src/piccolo/query_types/objects.rst

@@ -78,6 +78,9 @@ We also have this shortcut which combines the above into a single line:
 Updating objects
 ----------------
 
+``save``
+~~~~~~~~
+
 Objects have a :meth:`save <piccolo.table.Table.save>` method, which is
 convenient for updating values:
 
@@ -95,6 +98,36 @@ convenient for updating values:
     # Or specify specific columns to save:
     await band.save([Band.popularity])
 
+``update_self``
+~~~~~~~~~~~~~~~
+
+The :meth:`save <piccolo.table.Table.save>` method is fine in the majority of
+cases, but there are some situations where the :meth:`update_self <piccolo.table.Table.update_self>`
+method is preferable.
+
+For example, if we want to increment the ``popularity`` value, we can do this:
+
+.. code-block:: python
+
+    await band.update_self({
+        Band.popularity: Band.popularity + 1
+    })
+
+Which does the following:
+
+* Increments the popularity in the database
+* Assigns the new value to the object
+
+This is safer than:
+
+.. code-block:: python
+
+    band.popularity += 1
+    await band.save()
+
+Because ``update_self`` increments the current ``popularity`` value in the
+database, not the one on the object, which might be out of date.
+
 -------------------------------------------------------------------------------
 
 Deleting objects
@@ -115,8 +148,8 @@ Similarly, we can delete objects, using the ``remove`` method.
 Fetching related objects
 ------------------------
 
-get_related
-~~~~~~~~~~~
+``get_related``
+~~~~~~~~~~~~~~~
 
 If you have an object from a table with a :class:`ForeignKey <piccolo.columns.column_types.ForeignKey>`
 column, and you want to fetch the related row as an object, you can do so
@@ -195,8 +228,8 @@ prefer.
 
 -------------------------------------------------------------------------------
 
-get_or_create
--------------
+``get_or_create``
+-----------------
 
 With ``get_or_create`` you can get an existing record matching the criteria,
 or create a new one with the ``defaults`` arguments:
@@ -239,8 +272,8 @@ Complex where clauses are supported, but only within reason. For example:
 
 -------------------------------------------------------------------------------
 
-to_dict
--------
+``to_dict``
+-----------
 
 If you need to convert an object into a dictionary, you can do so using the
 ``to_dict`` method.
@@ -264,8 +297,8 @@ the columns:
 
 -------------------------------------------------------------------------------
 
-refresh
--------
+``refresh``
+-----------
 
 If you have an object which has gotten stale, and want to refresh it, so it
 has the latest data from the database, you can use the

piccolo/query/methods/objects.py

@@ -27,6 +27,7 @@
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.columns import Column
+    from piccolo.table import Table
 
 
 ###############################################################################
@@ -173,6 +174,61 @@ def run_sync(self, *args, **kwargs) -> 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))
+
+
 ###############################################################################
 
 

piccolo/table.py

@@ -46,7 +46,7 @@
 )
 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 @@ def save(
             == 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.

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