added set_default method to Alter

Author: dantownsend

piccolo/columns/base.py

@@ -245,6 +245,37 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
         """
         return self._meta.get_full_name(just_alias=just_alias)
 
+    def get_sql_value(self, value: t.Any) -> str:
+        """
+        When using DDL statements, we can't parameterise the values. An example
+        is when setting the default for a column. So we have to convert from
+        the Python type to a string representation which we can include in our
+        DDL statements.
+
+        :param value:
+            The Python value to convert to a string usable in a DDL statement
+            e.g. 1.
+        :returns:
+            The string usable in the DDL statement e.g. '1'.
+
+        """
+        if isinstance(value, Default):
+            output = getattr(value, self._meta.engine_type)
+        elif value is None:
+            output = "null"
+        elif isinstance(value, (float, decimal.Decimal)):
+            output = str(value)
+        elif isinstance(value, str):
+            output = f"'{value}'"
+        elif isinstance(value, bool):
+            output = str(value).lower()
+        elif isinstance(value, datetime.datetime):
+            output = f"'{value.isoformat().replace('T', '')}'"
+        else:
+            output = value
+
+        return output
+
     @property
     def querystring(self) -> QueryString:
         """
@@ -278,22 +309,8 @@ def querystring(self) -> QueryString:
 
         if not self._meta.primary:
             default = self.get_default_value()
-            if isinstance(default, Default):
-                default_value = getattr(default, self._meta.engine_type)
-            elif default is None:
-                default_value = "null"
-            elif isinstance(default, (float, decimal.Decimal)):
-                default_value = str(default)
-            elif isinstance(default, str):
-                default_value = f"'{default}'"
-            elif isinstance(default, bool):
-                default_value = str(default).lower()
-            elif isinstance(default, datetime.datetime):
-                default_value = f"'{default.isoformat().replace('T', '')}'"
-            else:
-                default_value = default
-
-            query += f" DEFAULT {default_value}"
+            sql_value = self.get_sql_value(value=default)
+            query += f" DEFAULT {sql_value}"
 
         return QueryString(query)
 

piccolo/query/methods/alter.py

@@ -91,6 +91,21 @@ def querystring(self) -> QueryString:
         return QueryString("DROP DEFAULT {}", self.column_name)
 
 
+@dataclass
+class SetDefault(AlterColumnStatement):
+    __slots__ = ("column", "value")
+
+    column: Column
+    value: t.Any
+
+    @property
+    def querystring(self) -> QueryString:
+        sql_value = self.column.get_sql_value(self.value)
+        return QueryString(
+            f"ALTER COLUMN {self.column_name} SET DEFAULT {sql_value}"
+        )
+
+
 @dataclass
 class Unique(AlterColumnStatement):
     __slots__ = ("boolean",)
@@ -224,9 +239,10 @@ class Alter(Query):
         "_null",
         "_rename_columns",
         "_rename_table",
+        "_set_default",
+        "_set_digits",
         "_set_length",
         "_unique",
-        "_set_digits",
     )
 
     def __init__(self, table: t.Type[Table]):
@@ -240,9 +256,10 @@ def __init__(self, table: t.Type[Table]):
         self._null: t.List[Null] = []
         self._rename_columns: t.List[RenameColumn] = []
         self._rename_table: t.List[RenameTable] = []
+        self._set_default: t.List[SetDefault] = []
+        self._set_digits: t.List[SetDigits] = []
         self._set_length: t.List[SetLength] = []
         self._unique: t.List[Unique] = []
-        self._set_digits: t.List[SetDigits] = []
 
     def add_column(self, name: str, column: Column) -> Alter:
         """
@@ -284,6 +301,15 @@ def rename_column(
         self._rename_columns.append(RenameColumn(column, new_name))
         return self
 
+    def set_default(self, column: Column, value: t.Any) -> Alter:
+        """
+        Set the default for a column.
+
+        Band.alter().set_default(Band.popularity, 0)
+        """
+        self._set_default.append(SetDefault(column=column, value=value))
+        return self
+
     def set_null(
         self, column: t.Union[str, Column], boolean: bool = True
     ) -> Alter:
@@ -463,6 +489,7 @@ def querystrings(self) -> t.Sequence[QueryString]:
                 self._unique,
                 self._null,
                 self._set_length,
+                self._set_default,
                 self._set_digits,
             )
         ]

tests/table/test_alter.py

@@ -141,6 +141,20 @@ def test_null(self):
         pass
 
 
+@postgres_only
+class TestSetDefault(DBTestCase):
+    def test_set_default(self):
+        Manager.alter().set_default(Manager.name, "Pending").run_sync()
+
+        # Bypassing the ORM to make sure the database default is present.
+        Band.raw(
+            "INSERT INTO manager (id, name) VALUES (DEFAULT, DEFAULT);"
+        ).run_sync()
+
+        manager = Manager.objects().first().run_sync()
+        self.assertTrue(manager.name == "Pending")
+
+
 ###############################################################################