refactoring update command

Author: dantownsend

piccolo/commands/playground.py

@@ -8,7 +8,7 @@
 
 from piccolo import table
 from piccolo import columns
-from piccolo.engine import SQLiteEngine
+from piccolo.engine import SQLiteEngine, PostgresEngine
 
 
 # DB = PostgresEngine({})

piccolo/query/methods/update.py

@@ -1,15 +1,48 @@
 from __future__ import annotations
 
 from piccolo.query.base import Query
-from piccolo.query.mixins import WhereMixin
+from piccolo.query.mixins import ValuesMixin, WhereMixin
 from piccolo.querystring import QueryString
 
 
-class Update(Query, WhereMixin):
+class Update(Query, ValuesMixin, WhereMixin):
 
     @property
-    def querystring(self) -> QueryString:
-        querystring = self.base
+    def sqlite_querystring(self) -> QueryString:
+        raise NotImplementedError()
+
+    def validate(self):
+        if len(self._values) == 0:
+            raise ValueError(
+                "No values were specified to update - please use .values"
+            )
+
+        for column, value in self._values.items():
+            if len(column.call_chain) > 0:
+                raise ValueError(
+                    "Related values can't be updated via an update"
+                )
+
+    @property
+    def postgres_querystring(self) -> QueryString:
+        self.validate()
+
+        columns_str = ', '.join([
+            f'{column_name} = {{}}' for column_name, _ in self._values.items()
+        ])
+
+        query = f'UPDATE {self.table.Meta.tablename} SET ' + columns_str
+
+        querystring = QueryString(
+            query,
+            self._values.values()
+        )
+
+        # TODO - need to support joins
+        # UPDATE band SET name='bob' FROM manager WHERE band.manager = manager.id AND manager.name = 'Guido';
+        # I think the joins should work differently from select i.e. no need
+        # to shared underlying code.
+
         if self._where:
             where_querystring = QueryString(
                 '{} WHERE {}',
@@ -20,5 +53,3 @@ def querystring(self) -> QueryString:
         else:
             return querystring
 
-    def __str__(self) -> str:
-        return self.querystring.__str__()

piccolo/query/mixins.py

@@ -198,3 +198,21 @@ def __init__(self):
     def columns(self, *columns: Column):
         self.selected_columns += columns
         return self
+
+
+class ValuesMixin():
+    """
+    Used to specify new column values - primarily used in update queries.
+
+    Example usage:
+
+    .values({MyTable.column_a: 1})
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._values: t.Dict[Column, t.Any] = {}
+
+    def values(self, values: t.Dict[Column, t.Any]):
+        self._values.update(values)
+        return self

piccolo/table.py

@@ -177,6 +177,19 @@ def select(cls) -> Select:
             table=cls,
         )
 
+    @property
+    def update(cls) -> Update:
+        """
+        Update rows.
+
+        await Band.update.values(
+            {Band.name: "Spamalot"}
+        ).where(Band.name=="Pythonistas")
+        """
+        return Update(
+            table=cls,
+        )
+
 
 class Table(metaclass=TableMeta):
 
@@ -355,30 +368,6 @@ def insert(cls, *rows: 'Table') -> Insert:
             query.add(*rows)
         return query
 
-    # TODO - needs refactoring into Band.update.columns()
-    @classmethod
-    def update(cls, **columns) -> Update:
-        """
-        await Band.update(name='Spamalot').where(
-            Band.name='Pythonistas'
-        ).run()
-        """
-        columns_str = ', '.join([
-            f'{column_name} = {{}}' for column_name in columns.keys()
-        ])
-
-        query = f'UPDATE {cls.Meta.tablename} SET ' + columns_str
-
-        querystring = QueryString(
-            query,
-            *columns.values()
-        )
-
-        return Update(
-            table=cls,
-            base=querystring
-        )
-
     @classmethod
     def raw(cls, sql: str) -> Raw:
         """