making auto migrations work in reverse

Author: dantownsend

docs/src/piccolo/migrations/create.rst

@@ -38,10 +38,6 @@ The contents of an empty migration file looks like this:
         manager.add_raw(run)
         return manager
 
-
-    async def backwards():
-        pass
-
 Replace the `run` function with whatever you want the migration to do -
 typically running some SQL. It can be a function or a coroutine.
 

piccolo/apps/migrations/auto/migration_manager.py

@@ -140,6 +140,9 @@ class MigrationManager:
         default_factory=AlterColumnCollection
     )
     raw: t.List[t.Union[t.Callable, t.Coroutine]] = field(default_factory=list)
+    raw_backwards: t.List[t.Union[t.Callable, t.Coroutine]] = field(
+        default_factory=list
+    )
 
     def add_table(
         self,
@@ -162,11 +165,16 @@ def drop_table(self, class_name: str, tablename: str):
         )
 
     def rename_table(
-        self, old_class_name: str, new_class_name: str, new_tablename: str
+        self,
+        old_class_name: str,
+        old_tablename: str,
+        new_class_name: str,
+        new_tablename: str,
     ):
         self.rename_tables.append(
             RenameTable(
                 old_class_name=old_class_name,
+                old_tablename=old_tablename,
                 new_class_name=new_class_name,
                 new_tablename=new_tablename,
             )
@@ -245,6 +253,13 @@ def add_raw(self, raw: t.Union[t.Callable, t.Coroutine]):
         """
         self.raw.append(raw)
 
+    def add_raw_backwards(self, raw: t.Union[t.Callable, t.Coroutine]):
+        """
+        When reversing a migration, you may want to run extra code to help
+        clean up.
+        """
+        self.raw_backwards.append(raw)
+
     ###########################################################################
 
     def deserialise_params(
@@ -325,16 +340,19 @@ async def run(self):
             ###################################################################
             # Add tables
 
-            for table in self.add_tables:
-                columns = self.add_columns.columns_for_table_class_name(
-                    table.class_name
+            for add_table in self.add_tables:
+                columns = (
+                    self.add_columns.columns_for_table_class_name(
+                        add_table.class_name
+                    )
+                    + add_table.columns
                 )
                 _Table: t.Type[Table] = type(
-                    table.class_name,
+                    add_table.class_name,
                     (Table,),
                     {column._meta.name: column for column in columns},
                 )
-                _Table._meta.tablename = table.tablename
+                _Table._meta.tablename = add_table.tablename
 
                 await _Table.create_table().run()
 
@@ -351,6 +369,7 @@ async def run(self):
                 _Table: t.Type[Table] = type(
                     rename_table.old_class_name, (Table,), {}
                 )
+                _Table._meta.tablename = rename_table.old_tablename
                 await _Table.alter().rename_table(
                     new_name=rename_table.new_tablename
                 ).run()
@@ -453,3 +472,108 @@ async def run(self):
                         await _Table.alter().set_unique(
                             column=row_name, boolean=unique
                         ).run()
+
+    ###########################################################################
+
+    async def run_backwards(self):
+        print("Reversing MigrationManager ...")
+
+        engine = engine_finder()
+
+        if not engine:
+            raise Exception("Can't find engine")
+
+        async with engine.transaction():
+
+            for raw in self.raw_backwards:
+                if inspect.iscoroutinefunction(raw):
+                    await raw()
+                else:
+                    raw()
+
+            ###################################################################
+            # Reverse add tables
+
+            for add_table in self.add_tables:
+                await add_table.to_table_class().alter().drop_table().run()
+
+            ###################################################################
+            # Reverse drop tables
+
+            if self.drop_tables:
+                print("Dropped tables can't currently be reversed.")
+
+            ###################################################################
+            # Reverse rename tables
+
+            for rename_table in self.rename_tables:
+                _Table: t.Type[Table] = type(
+                    rename_table.new_class_name, (Table,), {}
+                )
+                _Table._meta.tablename = rename_table.new_tablename
+
+                await _Table.alter().rename_table(
+                    new_name=rename_table.old_tablename
+                ).run()
+
+            ###################################################################
+            # Reverse add columns
+
+            for add_column in self.add_columns.add_columns:
+                _Table: t.Type[Table] = type(
+                    add_column.table_class_name, (Table,), {}
+                )
+                _Table._meta.tablename = add_column.tablename
+
+                await _Table.alter().drop_column(add_column.column).run()
+
+            ###################################################################
+            # Reverse drop columns
+
+            if self.drop_columns:
+                print("Dropped columns can't currently be reversed.")
+
+            ###################################################################
+            # Reverse rename columns
+
+            for rename_column in self.rename_columns.rename_columns:
+                _Table: t.Type[Table] = type(
+                    rename_column.table_class_name, (Table,), {}
+                )
+                _Table._meta.tablename = rename_column.tablename
+
+                await _Table.alter().rename_column(
+                    column=rename_column.new_column_name,
+                    new_name=rename_column.old_column_name,
+                ).run()
+
+            ###################################################################
+            # Alter columns
+
+            # TODO - need to find what the old values are.
+            # for alter_column in self.alter_columns.alter_columns:
+            #     _Table: t.Type[Table] = type(
+            #         alter_column.table_class_name, (Table,), {}
+            #     )
+            #     _Table._meta.tablename = alter_column.tablename
+
+            #     params = alter_column.params
+            #     row_name = alter_column.row_name
+
+            #     null = params.get("null")
+            #     if null is not None:
+            #         await _Table.alter().set_null(
+            #             column=row_name, boolean=null
+            #         ).run()
+
+            #     length = params.get("length")
+            #     if length is not None:
+            #         await _Table.alter().set_length(
+            #             column=row_name, length=length
+            #         ).run()
+
+            #     unique = params.get("unique")
+            #     if unique is not None:
+            #         await _Table.alter().set_unique(
+            #             column=row_name, boolean=unique
+            #         ).run()

piccolo/apps/migrations/auto/operations.py

@@ -5,6 +5,7 @@
 @dataclass
 class RenameTable:
     old_class_name: str
+    old_tablename: str
     new_class_name: str
     new_tablename: str
 

piccolo/apps/migrations/auto/schema_differ.py

@@ -134,6 +134,7 @@ def check_rename_tables(self) -> RenameTableCollection:
                         collection.append(
                             RenameTable(
                                 old_class_name=drop_table.class_name,
+                                old_tablename=drop_table.tablename,
                                 new_class_name=new_table.class_name,
                                 new_tablename=new_table.tablename,
                             )
@@ -237,7 +238,7 @@ def drop_tables(self) -> t.List[str]:
     @property
     def rename_tables(self) -> t.List[str]:
         return [
-            f"manager.rename_table(old_class_name='{renamed_table.old_class_name}', new_class_name='{renamed_table.new_class_name}', new_tablename='{renamed_table.new_tablename}')"  # noqa
+            f"manager.rename_table(old_class_name='{renamed_table.old_class_name}', old_tablename='{renamed_table.old_tablename}', new_class_name='{renamed_table.new_class_name}', new_tablename='{renamed_table.new_tablename}')"  # noqa
             for renamed_table in self.rename_tables_collection.rename_tables
         ]
 

piccolo/apps/migrations/commands/backwards.py

@@ -57,11 +57,11 @@ def run(self):
                 print(f"Reversing {migration_name}")
                 migration_module = migration_modules[migration_name]
                 response = asyncio.run(
-                    migration_module.backwards()
+                    migration_module.forwards()
                 )  # type: ignore
 
                 if isinstance(response, MigrationManager):
-                    asyncio.run(response.run())
+                    asyncio.run(response.run_backwards())
 
                 Migration.delete().where(
                     Migration.name == migration_name

piccolo/apps/migrations/commands/base.py

@@ -18,10 +18,6 @@ class MigrationModule(ModuleType):
     async def forwards() -> None:
         pass
 
-    @staticmethod
-    async def backwards() -> None:
-        pass
-
 
 class PiccoloAppModule(ModuleType):
     APP_CONFIG: AppConfig

piccolo/apps/migrations/commands/templates/migration.py.jinja

@@ -17,7 +17,3 @@ async def forwards():
     manager.add_raw(run)
     {% endif %}
     return manager
-
-
-async def backwards():
-    pass

piccolo/apps/user/piccolo_migrations/2019-11-14T21:52:21.py

@@ -82,7 +82,3 @@ async def forwards():
     )
 
     return manager
-
-
-async def backwards():
-    pass

tests/apps/migrations/auto/test_migration_manager.py

@@ -1,6 +1,7 @@
 import asyncio
 
 from piccolo.apps.migrations.auto import MigrationManager
+from piccolo.columns import Varchar
 
 from tests.base import DBTestCase
 
@@ -26,6 +27,12 @@ def test_rename_column(self):
         self.assertTrue("title" in response[0].keys())
         self.assertTrue("name" not in response[0].keys())
 
+        # Now reverse it
+        asyncio.run(manager.run_backwards())
+        response = self.run_sync("SELECT * FROM band;")
+        self.assertTrue("title" not in response[0].keys())
+        self.assertTrue("name" in response[0].keys())
+
     def test_raw_function(self):
         """
         Test adding raw functions to a MigrationManager.
@@ -39,10 +46,15 @@ def run():
 
         manager = MigrationManager()
         manager.add_raw(run)
+        manager.add_raw_backwards(run)
 
         with self.assertRaises(HasRun):
             asyncio.run(manager.run())
 
+        # Reverse
+        with self.assertRaises(HasRun):
+            asyncio.run(manager.run_backwards())
+
     def test_raw_coroutine(self):
         """
         Test adding raw coroutines to a MigrationManager.
@@ -56,6 +68,96 @@ async def run():
 
         manager = MigrationManager()
         manager.add_raw(run)
+        manager.add_raw_backwards(run)
 
         with self.assertRaises(HasRun):
             asyncio.run(manager.run())
+
+        # Reverse
+        with self.assertRaises(HasRun):
+            asyncio.run(manager.run_backwards())
+
+    def test_add_table(self):
+        """
+        Test adding a table to a MigrationManager.
+        """
+        self.run_sync("DROP TABLE IF EXISTS musician;")
+
+        manager = MigrationManager()
+        name_column = Varchar()
+        name_column._meta.name = "name"
+        manager.add_table(
+            class_name="Musician", tablename="musician", columns=[name_column]
+        )
+        asyncio.run(manager.run())
+
+        self.run_sync("INSERT INTO musician VALUES (default, 'Bob Jones');")
+        response = self.run_sync("SELECT * FROM musician;")
+
+        self.assertEqual(response, [{"id": 1, "name": "Bob Jones"}])
+
+        # Reverse
+        asyncio.run(manager.run_backwards())
+        self.assertEqual(self.table_exists("musician"), False)
+        self.run_sync("DROP TABLE IF EXISTS musician;")
+
+    def test_add_column(self):
+        """
+        Test adding a column to a MigrationManager.
+        """
+        manager = MigrationManager()
+        manager.add_column(
+            table_class_name="Manager",
+            tablename="manager",
+            column_name="email",
+            column_class_name="Varchar",
+            params={
+                "length": 100,
+                "default": "",
+                "null": True,
+                "primary": False,
+                "key": False,
+                "unique": True,
+                "index": False,
+            },
+        )
+        asyncio.run(manager.run())
+
+        self.run_sync(
+            "INSERT INTO manager VALUES (default, 'Dave', 'dave@me.com');"
+        )
+
+        response = self.run_sync("SELECT * FROM manager;")
+        self.assertEqual(
+            response, [{"id": 1, "name": "Dave", "email": "dave@me.com"}]
+        )
+
+        # Reverse
+        asyncio.run(manager.run_backwards())
+        response = self.run_sync("SELECT * FROM manager;")
+        self.assertEqual(response, [{"id": 1, "name": "Dave"}])
+
+    def test_rename_table(self):
+        """
+        Test renaming a table with MigrationManager.
+        """
+        manager = MigrationManager()
+
+        manager.rename_table(
+            old_class_name="Manager",
+            old_tablename="manager",
+            new_class_name="Director",
+            new_tablename="director",
+        )
+
+        asyncio.run(manager.run())
+
+        self.run_sync("INSERT INTO director VALUES (default, 'Dave');")
+
+        response = self.run_sync("SELECT * FROM director;")
+        self.assertEqual(response, [{"id": 1, "name": "Dave"}])
+
+        # Reverse
+        asyncio.run(manager.run_backwards())
+        response = self.run_sync("SELECT * FROM manager;")
+        self.assertEqual(response, [{"id": 1, "name": "Dave"}])