MigrationManager handles changes in column defaults

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -399,6 +399,19 @@ async def _run_alter_columns(self, backwards=False):
                         column=column, boolean=unique
                     ).run()
 
+                if "default" in params.keys():
+                    # Don't compare to None or ..., as both as valid values.
+                    default = params.get("default")
+                    column = Column()
+                    column._meta._table = _Table
+                    column._meta._name = column_name
+                    if default is ...:
+                        await _Table.alter().drop_default(column=column).run()
+                    else:
+                        await _Table.alter().set_default(
+                            column=column, value=default
+                        ).run()
+
                 # None is a valid value, so retrieve ellipsis if not found.
                 digits = params.get("digits", ...)
                 if digits is not ...:

piccolo/columns/base.py

@@ -172,6 +172,24 @@ def __init__(
             params=kwargs,
         )
 
+    def _validate_arguments(
+        default: t.Any, allowed_types: t.Tuple[t.Type[t.Any]], null: bool
+    ) -> bool:
+        """
+        Make sure that the default values are of the allowed types. Also
+        make sure that a value of None isn't the default if the column is not
+        nullable.
+        """
+        if default is None:
+            if not null:
+                raise ValueError(
+                    "A default value of None isn't allowed if the column is "
+                    "null = False."
+                )
+            return None in allowed_types
+        else:
+            return type(default) in allowed_types
+
     def is_in(self, values: Iterable) -> Where:
         return Where(column=self, values=values, operator=In)
 

piccolo/columns/column_types.py

@@ -7,7 +7,12 @@
 
 from piccolo.columns.base import Column, OnDelete, OnUpdate, ForeignKeyMeta
 from piccolo.columns.operators.string import ConcatPostgres, ConcatSQLite
-from piccolo.custom_types import UUIDDefault
+from piccolo.custom_types import (
+    DateDefault,
+    TimeDefault,
+    TimestampDefault,
+    UUIDDefault,
+)
 from piccolo.querystring import Unquoted, QueryString
 
 if t.TYPE_CHECKING:
@@ -120,7 +125,9 @@ class Varchar(Column):
     value_type = str
     concat_delegate: ConcatDelegate = ConcatDelegate()
 
-    def __init__(self, length: int = 255, default: str = "", **kwargs) -> None:
+    def __init__(
+        self, length: int = 255, default: t.Union[str, None] = "", **kwargs
+    ) -> None:
         self.length = length
         self.default = default
         kwargs.update({"length": length, "default": default})
@@ -168,7 +175,7 @@ class Text(Column):
     value_type = str
     concat_delegate: ConcatDelegate = ConcatDelegate()
 
-    def __init__(self, default: str = "", **kwargs) -> None:
+    def __init__(self, default: t.Union[str, None] = "", **kwargs) -> None:
         self.default = default
         super().__init__(**kwargs)
 
@@ -205,7 +212,7 @@ class Integer(Column):
 
     math_delegate = MathDelegate()
 
-    def __init__(self, default: int = None, **kwargs) -> None:
+    def __init__(self, default: t.Union[str, None] = 0, **kwargs) -> None:
         self.default = default
         kwargs.update({"default": default})
         super().__init__(**kwargs)
@@ -351,7 +358,9 @@ class Timestamp(Column):
 
     value_type = datetime
 
-    def __init__(self, default: TimestampArg = None, **kwargs) -> None:
+    def __init__(
+        self, default: TimestampArg = TimestampDefault.now, **kwargs
+    ) -> None:
         self.default = default
         kwargs.update({"default": default})
         super().__init__(**kwargs)
@@ -360,7 +369,7 @@ def __init__(self, default: TimestampArg = None, **kwargs) -> None:
 class Date(Column):
     value_type = date
 
-    def __init__(self, default: DateArg = None, **kwargs) -> None:
+    def __init__(self, default: DateArg = DateDefault.now, **kwargs) -> None:
         self.default = default
         kwargs.update({"default": default})
         super().__init__(**kwargs)
@@ -369,7 +378,7 @@ def __init__(self, default: DateArg = None, **kwargs) -> None:
 class Time(Column):
     value_type = time
 
-    def __init__(self, default: TimeArg = None, **kwargs) -> None:
+    def __init__(self, default: TimeArg = TimeDefault.now, **kwargs) -> None:
         self.default = default
         kwargs.update({"default": default})
         super().__init__(**kwargs)
@@ -503,6 +512,8 @@ class ForeignKey(Integer):
     def __init__(
         self,
         references: t.Union[t.Type[Table], str],
+        default: t.Union[int, None] = None,
+        null: bool = True,
         on_delete: OnDelete = OnDelete.cascade,
         on_update: OnUpdate = OnUpdate.cascade,
         **kwargs,
@@ -522,7 +533,7 @@ def __init__(
                 "on_update": on_update,
             }
         )
-        super().__init__(**kwargs)
+        super().__init__(default=default, null=null, **kwargs)
 
         if t.TYPE_CHECKING:
             # This is here just for type inference - the actual value is set by

piccolo/query/methods/alter.py

@@ -276,6 +276,13 @@ def drop_column(self, column: t.Union[str, Column]) -> Alter:
         self._drop.append(DropColumn(column))
         return self
 
+    def drop_default(self, column: t.Union[str, Column]) -> Alter:
+        """
+        Band.alter().drop_default(Band.popularity)
+        """
+        self._drop_default.append(DropDefault(column=column))
+        return self
+
     def drop_table(self) -> Alter:
         """
         Band.alter().drop_table()
@@ -486,6 +493,7 @@ def querystrings(self) -> t.Sequence[QueryString]:
                 self._rename_columns,
                 self._rename_table,
                 self._drop,
+                self._drop_default,
                 self._unique,
                 self._null,
                 self._set_length,

tests/apps/migrations/auto/test_migration_manager.py

@@ -228,11 +228,22 @@ def test_alter_column_unique(self):
         response = self.run_sync("SELECT name FROM manager;")
         self.assertEqual(response, [{"name": "Dave"}, {"name": "Dave"}])
 
-    def _get_precision_and_scale(self):
+    def _get_column_precision_and_scale(
+        self, tablename="ticket", column_name="price"
+    ):
         return self.run_sync(
             "SELECT numeric_precision, numeric_scale "
             "FROM information_schema.COLUMNS "
-            "WHERE TABLE_NAME = 'ticket' AND column_name = 'price';"
+            f"WHERE table_name = '{tablename}' AND "
+            f"column_name = '{column_name}';"
+        )
+
+    def _get_column_default(self, tablename="manager", column_name="name"):
+        return self.run_sync(
+            "SELECT column_default "
+            "FROM information_schema.COLUMNS "
+            f"WHERE table_name = '{tablename}' "
+            f"AND column_name = '{column_name}';"
         )
 
     @postgres_only
@@ -251,18 +262,44 @@ def test_alter_column_digits(self):
         )
 
         asyncio.run(manager.run())
-
         self.assertEqual(
-            self._get_precision_and_scale(),
+            self._get_column_precision_and_scale(),
             [{"numeric_precision": 6, "numeric_scale": 2}],
         )
 
         asyncio.run(manager.run_backwards())
         self.assertEqual(
-            self._get_precision_and_scale(),
+            self._get_column_precision_and_scale(),
             [{"numeric_precision": 5, "numeric_scale": 2}],
         )
 
+    @postgres_only
+    def test_alter_column_set_default(self):
+        """
+        Test altering a column default with MigrationManager.
+        """
+        manager = MigrationManager()
+
+        manager.alter_column(
+            table_class_name="Manager",
+            tablename="manager",
+            column_name="name",
+            params={"default": "Unknown"},
+            old_params={"default": ""},
+        )
+
+        asyncio.run(manager.run())
+        self.assertEqual(
+            self._get_column_default(),
+            [{"column_default": "'Unknown'::character varying"}],
+        )
+
+        asyncio.run(manager.run_backwards())
+        self.assertEqual(
+            self._get_column_default(),
+            [{"column_default": "''::character varying"}],
+        )
+
     @postgres_only
     @patch.object(BaseMigrationManager, "get_migration_managers")
     def test_drop_table(self, get_migration_managers: MagicMock):

tests/apps/migrations/auto/test_schema_differ.py

@@ -176,3 +176,6 @@ def test_alter_column_precision(self):
             schema_differ.alter_columns[0],
             "manager.alter_column(table_class_name='Ticket', tablename='ticket', column_name='price', params={'digits': (5, 2)}, old_params={'digits': (4, 2)})",  # noqa
         )
+
+    def test_alter_default(self):
+        pass

tests/table/test_alter.py

@@ -74,7 +74,9 @@ def test_add(self):
         """
         self.insert_row()
 
-        add_query = Band.alter().add_column("weight", Integer(null=True))
+        add_query = Band.alter().add_column(
+            "weight", Integer(null=True, default=None)
+        )
         add_query.run_sync()
 
         response = Band.raw("SELECT * FROM band").run_sync()