Alter.set_column_type now supports a using_expression argument

Author: dantownsend

piccolo/query/methods/alter.py

@@ -93,20 +93,30 @@ def querystring(self) -> QueryString:
 
 @dataclass
 class SetColumnType(AlterStatement):
-    __slots__ = ("old_column", "new_column")
+    """
+    :param using_expression:
+        Postgres can't automatically convert between certain column types. You
+        can tell Postgres which action to take. For example
+        `my_column_name::integer`.
+
+    """
 
     old_column: Column
     new_column: Column
+    using_expression: t.Optional[str] = None
 
     @property
     def querystring(self) -> QueryString:
         if self.new_column._meta._table is None:
             self.new_column._meta._table = self.old_column._meta.table
 
         column_name = self.old_column._meta.name
-        return QueryString(
+        query = (
             f"ALTER COLUMN {column_name} TYPE {self.new_column.column_type}"
         )
+        if self.using_expression is not None:
+            query += f" USING {self.using_expression}"
+        return QueryString(query)
 
 
 @dataclass
@@ -353,12 +363,21 @@ def rename_column(
         self._rename_columns.append(RenameColumn(column, new_name))
         return self
 
-    def set_column_type(self, old_column: Column, new_column: Column) -> Alter:
+    def set_column_type(
+        self,
+        old_column: Column,
+        new_column: Column,
+        using_expression: t.Optional[str] = None,
+    ) -> Alter:
         """
         Change the type of a column.
         """
         self._set_column_type.append(
-            SetColumnType(old_column=old_column, new_column=new_column)
+            SetColumnType(
+                old_column=old_column,
+                new_column=new_column,
+                using_expression=using_expression,
+            )
         )
         return self
 

tests/table/test_alter.py

@@ -1,7 +1,6 @@
-from piccolo.columns.column_types import Varchar
 from unittest import TestCase
 
-from piccolo.columns import BigInt, Integer, Numeric
+from piccolo.columns import BigInt, Integer, Numeric, Varchar
 from piccolo.table import Table
 
 from ..base import DBTestCase, postgres_only
@@ -186,6 +185,23 @@ def test_integer_to_varchar(self):
         )
         self.assertEqual(popularity, "1000")
 
+    def test_using_expression(self):
+        """
+        Test the `using_expression` option, which can be used to tell Postgres
+        how to convert certain column types.
+        """
+        Band(name="1").save().run_sync()
+
+        alter_query = Band.alter().set_column_type(
+            old_column=Band.name,
+            new_column=Integer(),
+            using_expression="name::integer",
+        )
+        alter_query.run_sync()
+
+        popularity = Band.select(Band.name).first().run_sync()["name"]
+        self.assertEqual(popularity, 1)
+
 
 @postgres_only
 class TestSetNull(DBTestCase):