Merge pull request piccolo-orm#27 from piccolo-orm/migration_improvements

Migration improvements

Author: dantownsend

piccolo/apps/migrations/auto/migration_manager.py

@@ -265,7 +265,7 @@ def add_raw_backwards(self, raw: t.Union[t.Callable, t.Coroutine]):
 
     ###########################################################################
 
-    def get_table_from_snaphot(
+    async def get_table_from_snaphot(
         self,
         table_class_name: str,
         app_name: t.Optional[str],
@@ -289,16 +289,13 @@ def get_table_from_snaphot(
         if app_name is None:
             app_name = self.app_name
 
-        return (
-            BaseMigrationManager()
-            .get_table_from_snaphot(
-                app_name=app_name,
-                table_class_name=table_class_name,
-                max_migration_id=migration_id,
-                offset=offset,
-            )
-            .to_table_class()
+        diffable_table = await BaseMigrationManager().get_table_from_snaphot(
+            app_name=app_name,
+            table_class_name=table_class_name,
+            max_migration_id=migration_id,
+            offset=offset,
         )
+        return diffable_table.to_table_class()
 
     ###########################################################################
 
@@ -375,7 +372,7 @@ async def _run_alter_columns(self, backwards=False):
     async def _run_drop_tables(self, backwards=False):
         if backwards:
             for diffable_table in self.drop_tables:
-                _Table = self.get_table_from_snaphot(
+                _Table = await self.get_table_from_snaphot(
                     table_class_name=diffable_table.class_name,
                     app_name=self.app_name,
                     offset=-1,
@@ -390,7 +387,7 @@ async def _run_drop_tables(self, backwards=False):
     async def _run_drop_columns(self, backwards=False):
         if backwards:
             for drop_column in self.drop_columns.drop_columns:
-                _Table = self.get_table_from_snaphot(
+                _Table = await self.get_table_from_snaphot(
                     table_class_name=drop_column.table_class_name,
                     app_name=self.app_name,
                     offset=-1,
@@ -477,24 +474,21 @@ async def _run_add_tables(self, backwards=False):
                 ).run()
         else:
             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(
                     add_table.class_name,
                     (Table,),
-                    {column._meta.name: column for column in columns},
+                    {
+                        column._meta.name: column
+                        for column in add_table.columns
+                    },
                 )
                 _Table._meta.tablename = add_table.tablename
 
                 await _Table.create_table().run()
 
     async def _run_add_columns(self, backwards=False):
         """
-        Add columns, which belong to existing tables
+        Add columns, which belong to new and existing tables
         """
         if backwards:
             for add_column in self.add_columns.add_columns:
@@ -512,12 +506,7 @@ async def _run_add_columns(self, backwards=False):
 
                 await _Table.alter().drop_column(add_column.column).run()
         else:
-            new_table_class_names = [i.class_name for i in self.add_tables]
-
             for table_class_name in self.add_columns.table_class_names:
-                if table_class_name in new_table_class_names:
-                    continue
-
                 add_columns: t.List[
                     AddColumnClass
                 ] = self.add_columns.for_table_class_name(table_class_name)
@@ -554,8 +543,8 @@ async def run(self):
             await self._run_add_tables()
             await self._run_rename_tables()
             await self._run_add_columns()
-            await self._run_drop_tables()
             await self._run_drop_columns()
+            await self._run_drop_tables()
             await self._run_rename_columns()
             await self._run_alter_columns()
 
@@ -575,8 +564,8 @@ async def run_backwards(self):
                 else:
                     raw()
 
-            await self._run_add_tables(backwards=True)
             await self._run_add_columns(backwards=True)
+            await self._run_add_tables(backwards=True)
             await self._run_drop_tables(backwards=True)
             await self._run_rename_tables(backwards=True)
             await self._run_drop_columns(backwards=True)

piccolo/apps/migrations/auto/schema_differ.py

@@ -129,12 +129,31 @@ def check_rename_tables(self) -> RenameTableCollection:
                     drop_column_names
                 )
                 if len(same_column_names) > 0:
+                    if (
+                        drop_table.class_name == new_table.class_name
+                        and drop_table.tablename != new_table.tablename
+                    ):
+                        # The class names are the same, but the tablename
+                        # has changed - we can assume this is a deliberate
+                        # rename.
+                        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,
+                            )
+                        )
+                        continue
+
                     user_response = (
                         self.auto_input
                         if self.auto_input
                         else input(
-                            f"Did you rename {drop_table.class_name} to "
-                            f"{new_table.class_name}? (y/N)"
+                            f"Did you rename {drop_table.class_name} "
+                            f"(tablename: {drop_table.tablename}) to "
+                            f"{new_table.class_name} "
+                            f"(tablename: {new_table.tablename})? (y/N)"
                         )
                     )
                     if user_response.lower() == "y":

piccolo/apps/migrations/commands/backwards.py

@@ -5,7 +5,6 @@
 from piccolo.apps.migrations.auto import MigrationManager
 from piccolo.apps.migrations.commands.base import BaseMigrationManager
 from piccolo.apps.migrations.tables import Migration
-from piccolo.utils.sync import run_sync
 
 
 class BackwardsMigrationManager(BaseMigrationManager):
@@ -22,8 +21,8 @@ def __init__(
         self.clean = clean
         super().__init__()
 
-    def run(self):
-        self.create_migration_table()
+    async def run(self):
+        await self.create_migration_table()
 
         app_modules = self.get_app_modules()
 
@@ -37,7 +36,7 @@ def run(self):
                 )
                 break
 
-        ran_migration_ids = Migration.get_migrations_which_ran(
+        ran_migration_ids = await Migration.get_migrations_which_ran(
             app_name=self.app_name
         )
         if len(ran_migration_ids) == 0:
@@ -85,14 +84,14 @@ def run(self):
             for migration_id in reversed_migration_ids:
                 print(f"Reversing {migration_id}")
                 migration_module = migration_modules[migration_id]
-                response = run_sync(migration_module.forwards())
+                response = await migration_module.forwards()
 
                 if isinstance(response, MigrationManager):
-                    run_sync(response.run_backwards())
+                    await response.run_backwards()
 
-                Migration.delete().where(
+                await Migration.delete().where(
                     Migration.name == migration_id
-                ).run_sync()
+                ).run()
 
                 if self.clean:
                     os.unlink(migration_module.__file__)
@@ -101,7 +100,7 @@ def run(self):
             sys.exit("Not proceeding.")
 
 
-def backwards(
+async def backwards(
     app_name: str,
     migration_id: str = "1",
     auto_agree: bool = False,
@@ -118,7 +117,7 @@ def backwards(
         Specify a value of 'all' to undo all of the migrations. Specify a
         value of '1' to undo the most recent migration.
     :param auto_agree:
-        Automatically agree to any input prompts.
+        If true, automatically agree to any input prompts.
     :param clean:
         If true, the migration files which have been run backwards are deleted
         from the disk after completing.
@@ -146,12 +145,12 @@ def backwards(
                     migration_id="all",
                     auto_agree=auto_agree,
                 )
-                manager.run()
+                await manager.run()
     else:
         manager = BackwardsMigrationManager(
             app_name=app_name,
             migration_id=migration_id,
             auto_agree=auto_agree,
             clean=clean,
         )
-        manager.run()
+        await manager.run()

piccolo/apps/migrations/commands/base.py

@@ -12,25 +12,26 @@
 from piccolo.apps.migrations.auto.diffable_table import DiffableTable
 from piccolo.apps.migrations.auto.schema_snapshot import SchemaSnapshot
 from piccolo.apps.migrations.tables import Migration
-from piccolo.utils.sync import run_sync
 
 
 class BaseMigrationManager(Finder):
-    def create_migration_table(self) -> bool:
+    async def create_migration_table(self) -> bool:
         """
         Creates the migration table in the database. Returns True/False
         depending on whether it was created or not.
         """
-        if not Migration.table_exists().run_sync():
-            Migration.create_table().run_sync()
+        if not await Migration.table_exists().run():
+            await Migration.create_table().run()
             return True
         return False
 
     def get_migration_modules(
         self, folder_path: str
     ) -> t.Dict[str, MigrationModule]:
         """
-        Import the migration modules and store them in a dictionary.
+        Imports the migration modules in the given folder path, and returns
+        a mapping of migration ID to the corresponding migration module.
+
         """
         migration_modules = {}
 
@@ -63,13 +64,17 @@ def get_migration_ids(
         """
         return sorted(list(migration_module_dict.keys()))
 
-    def get_migration_managers(
+    async def get_migration_managers(
         self,
         app_name: str,
         max_migration_id: t.Optional[str] = None,
         offset: int = 0,
     ) -> t.List[MigrationManager]:
         """
+        Call the forwards coroutine in each migration module. Each one should
+        return a `MigrationManger`. Combine all of the results, and return in
+        a list.
+
         :param max_migration_id:
             If set, only MigrationManagers up to and including the given
             migration ID will be returned.
@@ -84,16 +89,18 @@ def get_migration_managers(
             str, MigrationModule
         ] = self.get_migration_modules(migrations_folder)
 
-        for _, migration_module in migration_modules.items():
-            response = run_sync(migration_module.forwards())
+        migration_ids = sorted(migration_modules.keys())
+
+        for migration_id in migration_ids:
+            migration_module = migration_modules[migration_id]
+            response = await migration_module.forwards()
             if isinstance(response, MigrationManager):
-                if max_migration_id:
-                    if response.migration_id == max_migration_id:
-                        break
-                    else:
-                        migration_managers.append(response)
-                else:
-                    migration_managers.append(response)
+                migration_managers.append(response)
+                if (
+                    max_migration_id
+                    and response.migration_id == max_migration_id
+                ):
+                    break
 
         if offset > 0:
             raise Exception(
@@ -104,7 +111,7 @@ def get_migration_managers(
         else:
             return migration_managers
 
-    def get_table_from_snaphot(
+    async def get_table_from_snaphot(
         self,
         app_name: str,
         table_class_name: str,
@@ -115,7 +122,7 @@ def get_table_from_snaphot(
         This will generate a SchemaSnapshot up to the given migration ID, and
         will return a DiffableTable class from that snapshot.
         """
-        migration_managers = self.get_migration_managers(
+        migration_managers = await self.get_migration_managers(
             app_name=app_name, max_migration_id=max_migration_id, offset=offset
         )
         schema_snapshot = SchemaSnapshot(managers=migration_managers)

piccolo/apps/migrations/commands/check.py

@@ -8,11 +8,11 @@ def __init__(self, app_name: str):
         self.app_name = app_name
         super().__init__()
 
-    def run(self):
+    async def run(self):
         print("Listing migrations ...")
 
         # Make sure the migration table exists, otherwise we'll get an error.
-        self.create_migration_table()
+        await self.create_migration_table()
 
         print(
             f'{get_fixed_length_string("APP NAME")} | '
@@ -37,25 +37,25 @@ def run(self):
             ids = self.get_migration_ids(migration_modules)
             for _id in ids:
                 has_ran = (
-                    Migration.exists()
+                    await Migration.exists()
                     .where(
                         (Migration.name == _id)
                         & (Migration.app_name == app_name)
                     )
-                    .run_sync()
+                    .run()
                 )
                 fixed_length_id = get_fixed_length_string(_id)
                 print(
                     f"{fixed_length_app_name} | {fixed_length_id} | {has_ran}"
                 )
 
 
-def check(app_name: str = "all"):
+async def check(app_name: str = "all"):
     """
     Lists all migrations which have and haven't ran.
 
     :param app_name:
         The name of the app to check. Specify a value of 'all' to check
         the migrations for all apps.
     """
-    CheckMigrationManager(app_name=app_name).run()
+    await CheckMigrationManager(app_name=app_name).run()