integrated SetColumnType into migrations

Author: dantownsend

piccolo/apps/migrations/auto/diffable_table.py

@@ -110,6 +110,8 @@ def __sub__(self, value: DiffableTable) -> TableDelta:
                         column_name=column._meta.name,
                         params=deserialise_params(delta),
                         old_params=old_params,
+                        column_class=column.__class__,
+                        old_column_class=existing_column.__class__,
                     )
                 )
 

piccolo/apps/migrations/auto/migration_manager.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 from dataclasses import dataclass, field
 import inspect
+from piccolo.query.methods import alter
 import typing as t
 
 from piccolo.columns import Column, column_types
@@ -252,6 +253,8 @@ def alter_column(
         column_name: str,
         params: t.Dict[str, t.Any],
         old_params: t.Dict[str, t.Any],
+        column_class: t.Optional[t.Type[Column]] = None,
+        old_column_class: t.Optional[t.Type[Column]] = None,
     ):
         """
         All possible alterations aren't currently supported.
@@ -263,6 +266,8 @@ def alter_column(
                 column_name=column_name,
                 params=params,
                 old_params=old_params,
+                column_class=column_class,
+                old_column_class=old_column_class,
             )
         )
 
@@ -328,9 +333,34 @@ async def _run_alter_columns(self, backwards=False):
             _Table: t.Type[Table] = type(table_class_name, (Table,), {})
             _Table._meta.tablename = alter_columns[0].tablename
 
-            for column in alter_columns:
-                params = column.old_params if backwards else column.params
-                column_name = column.column_name
+            for alter_column in alter_columns:
+                params = (
+                    alter_column.old_params
+                    if backwards
+                    else alter_column.params
+                )
+                column_name = alter_column.column_name
+
+                # Change the column type if possible
+                column_class = alter_column.column_class
+                old_column_class = alter_column.old_column_class
+                if (old_column_class is not None) and (
+                    column_class is not None
+                ):
+                    if old_column_class != column_class:
+                        old_column = old_column_class(
+                            **alter_column.old_params
+                        )
+                        old_column._meta._table = _Table
+                        old_column._meta._name = alter_column.column_name
+
+                        new_column = column_class(**alter_column.params)
+                        new_column._meta._table = _Table
+                        new_column._meta._name = alter_column.column_name
+
+                        await _Table.alter().set_column_type(
+                            old_column=old_column, new_column=new_column
+                        )
 
                 null = params.get("null")
                 if null is not None:
@@ -383,7 +413,7 @@ async def _run_alter_columns(self, backwards=False):
                 digits = params.get("digits", ...)
                 if digits is not ...:
                     await _Table.alter().set_digits(
-                        column=column.column_name, digits=digits,
+                        column=alter_column.column_name, digits=digits,
                     ).run()
 
     async def _run_drop_tables(self, backwards=False):

piccolo/apps/migrations/auto/operations.py

@@ -26,6 +26,8 @@ class AlterColumn:
     tablename: str
     params: t.Dict[str, t.Any]
     old_params: t.Dict[str, t.Any]
+    column_class: t.Optional[t.Type[Column]] = None
+    old_column_class: t.Optional[t.Type[Column]] = None
 
 
 @dataclass

piccolo/apps/migrations/auto/schema_differ.py

@@ -308,17 +308,45 @@ def alter_columns(self) -> AlterStatements:
             else:
                 continue
 
-            for i in delta.alter_columns:
-                new_params = serialise_params(i.params)
+            for alter_column in delta.alter_columns:
+                new_params = serialise_params(alter_column.params)
                 extra_imports.extend(new_params.extra_imports)
                 extra_definitions.extend(new_params.extra_definitions)
 
-                old_params = serialise_params(i.old_params)
+                old_params = serialise_params(alter_column.old_params)
                 extra_imports.extend(old_params.extra_imports)
                 extra_definitions.extend(old_params.extra_definitions)
 
+                column_class = (
+                    alter_column.column_class.__name__
+                    if alter_column.column_class
+                    else "None"
+                )
+
+                old_column_class = (
+                    alter_column.old_column_class.__name__
+                    if alter_column.old_column_class
+                    else "None"
+                )
+
+                if alter_column.column_class is not None:
+                    extra_imports.append(
+                        Import(
+                            module=alter_column.column_class.__module__,
+                            target=alter_column.column_class.__name__,
+                        )
+                    )
+
+                if alter_column.old_column_class is not None:
+                    extra_imports.append(
+                        Import(
+                            module=alter_column.old_column_class.__module__,
+                            target=alter_column.old_column_class.__name__,
+                        )
+                    )
+
                 response.append(
-                    f"manager.alter_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{i.column_name}', params={new_params.params}, old_params={old_params.params})"  # noqa: E501
+                    f"manager.alter_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{alter_column.column_name}', params={new_params.params}, old_params={old_params.params}, column_class={column_class}, old_column_class={old_column_class})"  # noqa: E501
                 )
 
         return AlterStatements(

piccolo/query/methods/alter.py

@@ -100,7 +100,9 @@ class SetColumnType(AlterStatement):
 
     @property
     def querystring(self) -> QueryString:
-        self.new_column._meta._table = self.old_column._meta.table
+        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(
             f"ALTER COLUMN {column_name} TYPE {self.new_column.column_type}"

tests/apps/migrations/auto/test_schema_differ.py

@@ -187,7 +187,7 @@ def test_alter_column_precision(self):
         self.assertTrue(len(schema_differ.alter_columns.statements) == 1)
         self.assertEqual(
             schema_differ.alter_columns.statements[0],
-            "manager.alter_column(table_class_name='Ticket', tablename='ticket', column_name='price', params={'digits': (4, 2)}, old_params={'digits': (5, 2)})",  # noqa
+            "manager.alter_column(table_class_name='Ticket', tablename='ticket', column_name='price', params={'digits': (4, 2)}, old_params={'digits': (5, 2)}, column_class=Numeric, old_column_class=Numeric)",  # noqa
         )
 
     def test_alter_default(self):

tests/table/test_alter.py

@@ -137,6 +137,7 @@ def test_multiple(self):
         self.assertTrue("column_b" in column_names)
 
 
+# TODO - test more conversions.
 @postgres_only
 class TestSetColumnType(DBTestCase):
     def test_integer_to_bigint(self):