changed the migration commands to be top level async

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/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,
@@ -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,17 +12,16 @@
 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
 
@@ -65,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.
@@ -86,8 +89,11 @@ 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):
                 migration_managers.append(response)
                 if (
@@ -105,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,
@@ -116,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()

piccolo/apps/migrations/commands/clean.py

@@ -39,11 +39,11 @@ def get_migration_ids_to_remove(self) -> t.List[str]:
         migration_ids_to_remove = query.run_sync()
         return migration_ids_to_remove
 
-    def run(self):
+    async def run(self):
         print("Checking the migration table ...")
 
         # Make sure the migration table exists, otherwise we'll get an error.
-        self.create_migration_table()
+        await self.create_migration_table()
 
         migration_ids_to_remove = self.get_migration_ids_to_remove()
 
@@ -60,9 +60,9 @@ def run(self):
                 else input("Would you like to delete these rows? (y/N)")
             )
             if confirm == "y":
-                Migration.delete().where(
+                await Migration.delete().where(
                     Migration.name.is_in(migration_ids_to_remove)
-                ).run_sync()
+                ).run()
                 print("Deleted")
             else:
                 print("Cancelled")
@@ -73,7 +73,7 @@ def run(self):
             )
 
 
-def clean(app_name: str, auto_agree: bool = False):
+async def clean(app_name: str, auto_agree: bool = False):
     """
     Identifies any rows in the migration table which have no corresponding
     migration module on disk, and then optionally deletes those rows.
@@ -84,4 +84,4 @@ def clean(app_name: str, auto_agree: bool = False):
         If true, automatically agree to any input prompts.
 
     """
-    CleanMigrationManager(app_name=app_name, auto_agree=auto_agree).run()
+    await CleanMigrationManager(app_name=app_name, auto_agree=auto_agree).run()