can now reverse deleting tables and columns

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -123,6 +123,7 @@ class MigrationManager:
     """
 
     migration_id: str = ""
+    app_name: str = ""
     add_tables: t.List[DiffableTable] = field(default_factory=list)
     drop_tables: t.List[DiffableTable] = field(default_factory=list)
     rename_tables: t.List[RenameTable] = field(default_factory=list)
@@ -324,8 +325,8 @@ def deserialise_params(
 
     def get_table_from_snaphot(
         self,
-        app_name: str,
         table_class_name: str,
+        app_name: t.Optional[str],
         offset: int = 0,
         migration_id: t.Optional[str] = None,
     ) -> t.Type[Table]:
@@ -343,6 +344,9 @@ def get_table_from_snaphot(
         if migration_id is None:
             migration_id = self.migration_id
 
+        if app_name is None:
+            app_name = self.app_name
+
         return (
             BaseMigrationManager()
             .get_table_from_snaphot(
@@ -397,14 +401,31 @@ async def _run_alter_columns(self, backwards=False):
 
     async def _run_drop_tables(self, backwards=False):
         if backwards:
-            print("Dropped tables can't currently be reversed.")
+            for diffable_table in self.drop_tables:
+                _Table = self.get_table_from_snaphot(
+                    table_class_name=diffable_table.class_name,
+                    app_name=self.app_name,
+                    offset=-1,
+                )
+                await _Table.create_table().run()
         else:
-            for table in self.drop_tables:
-                await table.to_table_class().alter().drop_table().run()
+            for diffable_table in self.drop_tables:
+                await diffable_table.to_table_class().alter().drop_table().run()
 
     async def _run_drop_columns(self, backwards=False):
         if backwards:
-            print("Dropped columns can't currently be reversed.")
+            for drop_column in self.drop_columns.drop_columns:
+                _Table = self.get_table_from_snaphot(
+                    table_class_name=drop_column.table_class_name,
+                    app_name=self.app_name,
+                    offset=-1,
+                )
+                column_to_restore = _Table._meta.get_column_by_name(
+                    drop_column.column_name
+                )
+                await _Table.alter().add_column(
+                    name=drop_column.column_name, column=column_to_restore
+                ).run()
         else:
             for table_class_name in self.drop_columns.table_class_names:
                 columns = self.drop_columns.for_table_class_name(

piccolo/apps/migrations/commands/new.py

@@ -66,7 +66,10 @@ def _create_new_migration(app_config: AppConfig, auto=False) -> None:
             sys.exit(1)
 
         file_contents = render_template(
-            migration_id=_id, auto=True, alter_statements=alter_statements,
+            migration_id=_id,
+            auto=True,
+            alter_statements=alter_statements,
+            app_name=app_config.app_name,
         )
     else:
         file_contents = render_template(migration_id=_id, auto=False)

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

@@ -5,7 +5,7 @@ ID = '{{ migration_id }}'
 
 
 async def forwards():
-    manager = MigrationManager(migration_id=ID)
+    manager = MigrationManager(migration_id=ID, app_name="{{ app_name }}")
     {% if auto %}
     {% for alter_statement in alter_statements %}
     {{ alter_statement }}

tests/apps/migrations/auto/test_migration_manager.py

@@ -1,7 +1,9 @@
 import asyncio
+from unittest.mock import patch, MagicMock
 
 from asyncpg.exceptions import UniqueViolationError
 from piccolo.apps.migrations.auto import MigrationManager
+from piccolo.apps.migrations.commands.base import BaseMigrationManager
 from piccolo.columns import Varchar
 
 from tests.base import DBTestCase
@@ -193,3 +195,33 @@ def test_alter_column(self):
         )
         response = self.run_sync("SELECT name FROM manager;")
         self.assertEqual(response, [{"name": "Dave"}, {"name": "Dave"}])
+
+    @patch.object(BaseMigrationManager, "get_migration_managers")
+    def test_drop_table(self, get_migration_managers: MagicMock):
+        self.run_sync("DROP TABLE IF EXISTS musician;")
+
+        name_column = Varchar()
+        name_column._meta.name = "name"
+
+        manager_1 = MigrationManager(migration_id="1", app_name="music")
+        manager_1.add_table(
+            class_name="Musician", tablename="musician", columns=[name_column]
+        )
+        asyncio.run(manager_1.run())
+
+        manager_2 = MigrationManager(migration_id="2", app_name="music")
+        manager_2.drop_table(class_name="Musician", tablename="musician")
+        asyncio.run(manager_2.run())
+
+        get_migration_managers.return_value = [manager_1]
+
+        self.assertTrue(not self.table_exists("musician"))
+
+        asyncio.run(manager_2.run_backwards())
+
+        get_migration_managers.assert_called_with(
+            app_name="music", max_migration_id="2", offset=-1
+        )
+        self.assertTrue(self.table_exists("musician"))
+
+        self.run_sync("DROP TABLE IF EXISTS musician;")

tests/base.py

@@ -33,6 +33,7 @@ def run_sync(self, query):
 
     def table_exists(self, tablename: str):
         _Table: t.Type[Table] = type(tablename.upper(), (Table,), {})
+        _Table._meta.tablename = tablename
         return _Table.table_exists().run_sync()
 
     def create_table(self):