Merge pull request piccolo-orm#90 from piccolo-orm/programmatic_migrations

dantownsend · web-flow · commit 0511b9acf085 · 2021-04-26T21:51:41.000+01:00
Programmatic migrations
diff --git a/piccolo/apps/migrations/commands/backwards.py b/piccolo/apps/migrations/commands/backwards.py
@@ -3,7 +3,10 @@
 import sys
 
 from piccolo.apps.migrations.auto import MigrationManager
-from piccolo.apps.migrations.commands.base import BaseMigrationManager
+from piccolo.apps.migrations.commands.base import (
+    BaseMigrationManager,
+    MigrationResult,
+)
 from piccolo.apps.migrations.tables import Migration
 
 
@@ -21,12 +24,12 @@ def __init__(
         self.clean = clean
         super().__init__()
 
-    async def run(self):
+    async def run(self) -> MigrationResult:
         await self.create_migration_table()
 
         app_modules = self.get_app_modules()
 
-        migration_modules = []
+        migration_modules = {}
 
         for app_module in app_modules:
             app_config = getattr(app_module, "APP_CONFIG")
@@ -40,10 +43,11 @@ async def run(self):
             app_name=self.app_name
         )
         if len(ran_migration_ids) == 0:
-            # Make sure a status of 0 is returned, as we don't want this
+            # Make sure a success is returned, as we don't want this
             # to appear as an error in automated scripts.
-            print("No migrations to reverse!")
-            sys.exit(0)
+            message = "No migrations to reverse!"
+            print(message)
+            return MigrationResult(success=True, message=message)
 
         #######################################################################
 
@@ -55,10 +59,12 @@ async def run(self):
             earliest_migration_id = self.migration_id
 
         if earliest_migration_id not in ran_migration_ids:
-            sys.exit(
+            message = (
                 "Unrecognized migration name - must be one of "
                 f"{ran_migration_ids}"
             )
+            print(message, file=sys.stderr)
+            return MigrationResult(success=False, message=message)
 
         #######################################################################
 
@@ -99,33 +105,20 @@ async def run(self):
                 if self.clean:
                     os.unlink(migration_module.__file__)
 
+            return MigrationResult(success=True)
+
         else:  # pragma: no cover
-            sys.exit("Not proceeding.")
+            message = "Not proceeding."
+            print(message, file=sys.stderr)
+            return MigrationResult(success=False, message=message)
 
 
-async def backwards(
+async def run_backwards(
     app_name: str,
     migration_id: str = "1",
     auto_agree: bool = False,
     clean: bool = False,
-):
-    """
-    Undo migrations up to a specific migration.
-
-    :param app_name:
-        The app to reverse migrations for. Specify a value of 'all' to reverse
-        migrations for all apps.
-    :param migration_id:
-        Migrations will be reversed up to and including this migration_id.
-        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:
-        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.
-
-    """
+) -> MigrationResult:
     if app_name == "all":
         sorted_app_names = BaseMigrationManager().get_sorted_app_names()
         sorted_app_names.reverse()
@@ -149,11 +142,48 @@ async def backwards(
                     auto_agree=auto_agree,
                 )
                 await manager.run()
+            return MigrationResult(success=True)
+        else:
+            return MigrationResult(success=False, message="User cancelled")
     else:
         manager = BackwardsMigrationManager(
             app_name=app_name,
             migration_id=migration_id,
             auto_agree=auto_agree,
             clean=clean,
         )
-        await manager.run()
+        return await manager.run()
+
+
+async def backwards(
+    app_name: str,
+    migration_id: str = "1",
+    auto_agree: bool = False,
+    clean: bool = False,
+):
+    """
+    Undo migrations up to a specific migration.
+
+    :param app_name:
+        The app to reverse migrations for. Specify a value of 'all' to reverse
+        migrations for all apps.
+    :param migration_id:
+        Migrations will be reversed up to and including this migration_id.
+        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:
+        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.
+
+    """
+    response = await run_backwards(
+        app_name=app_name,
+        migration_id=migration_id,
+        auto_agree=auto_agree,
+        clean=clean,
+    )
+
+    if not response.success:
+        sys.exit(1)
diff --git a/piccolo/apps/migrations/commands/base.py b/piccolo/apps/migrations/commands/base.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+from dataclasses import dataclass
 import importlib
 import os
 import sys
@@ -14,6 +15,12 @@
 from piccolo.apps.migrations.tables import Migration
 
 
+@dataclass
+class MigrationResult:
+    success: bool
+    message: t.Optional[str] = None
+
+
 class BaseMigrationManager(Finder):
     async def create_migration_table(self) -> bool:
         """
diff --git a/piccolo/apps/migrations/commands/forwards.py b/piccolo/apps/migrations/commands/forwards.py
@@ -3,26 +3,22 @@
 import typing as t
 
 from piccolo.apps.migrations.auto import MigrationManager
-from piccolo.apps.migrations.commands.base import BaseMigrationManager
+from piccolo.apps.migrations.commands.base import (
+    BaseMigrationManager,
+    MigrationResult,
+)
 from piccolo.apps.migrations.tables import Migration
 from piccolo.conf.apps import AppConfig, MigrationModule
 
 
 class ForwardsMigrationManager(BaseMigrationManager):
-    def __init__(
-        self,
-        app_name: str,
-        migration_id: str,
-        fake: bool = False,
-        *args,
-        **kwargs,
-    ):
+    def __init__(self, app_name: str, migration_id: str, fake: bool = False):
         self.app_name = app_name
         self.migration_id = migration_id
         self.fake = fake
         super().__init__()
 
-    async def run_migrations(self, app_config: AppConfig) -> None:
+    async def run_migrations(self, app_config: AppConfig) -> MigrationResult:
         already_ran = await Migration.get_migrations_which_ran(
             app_name=self.app_name
         )
@@ -38,10 +34,11 @@ async def run_migrations(self, app_config: AppConfig) -> None:
         print(f"Haven't run = {havent_run}")
 
         if len(havent_run) == 0:
-            # Make sure a status of 0 is returned, as we don't want this
+            # Make sure this still appears successful, as we don't want this
             # to appear as an error in automated scripts.
-            print("No migrations left to run!")
-            sys.exit(0)
+            message = "No migrations left to run!"
+            print(message)
+            return MigrationResult(success=True, message=message)
 
         if self.migration_id == "all":
             subset = havent_run
@@ -51,7 +48,9 @@ async def run_migrations(self, app_config: AppConfig) -> None:
             try:
                 index = havent_run.index(self.migration_id)
             except ValueError:
-                sys.exit(f"{self.migration_id} is unrecognised")
+                message = f"{self.migration_id} is unrecognised"
+                print(message, file=sys.stderr)
+                return MigrationResult(success=False, message=message)
             else:
                 subset = havent_run[: index + 1]
 
@@ -71,13 +70,43 @@ async def run_migrations(self, app_config: AppConfig) -> None:
                 Migration(name=_id, app_name=self.app_name)
             ).run()
 
-    async def run(self):
+        return MigrationResult(success=True, message="Ran successfully")
+
+    async def run(self) -> MigrationResult:
         print("Running migrations ...")
         await self.create_migration_table()
 
         app_config = self.get_app_config(app_name=self.app_name)
 
-        await self.run_migrations(app_config)
+        return await self.run_migrations(app_config)
+
+
+async def run_forwards(
+    app_name: str, migration_id: str = "all", fake: bool = False
+) -> MigrationResult:
+    """
+    Run the migrations. This function can be used to programatically run
+    migrations - for example, in a unit test.
+    """
+    if app_name == "all":
+        sorted_app_names = BaseMigrationManager().get_sorted_app_names()
+        for _app_name in sorted_app_names:
+            print(f"\nMigrating {_app_name}")
+            print("------------------------------------------------")
+            manager = ForwardsMigrationManager(
+                app_name=_app_name, migration_id="all", fake=fake
+            )
+            response = await manager.run()
+            if not response.success:
+                return response
+
+        return MigrationResult(success=True)
+
+    else:
+        manager = ForwardsMigrationManager(
+            app_name=app_name, migration_id=migration_id, fake=fake
+        )
+        return await manager.run()
 
 
 async def forwards(
@@ -97,17 +126,9 @@ async def forwards(
         If set, will record the migrations as being run without actually
         running them.
     """
-    if app_name == "all":
-        sorted_app_names = BaseMigrationManager().get_sorted_app_names()
-        for _app_name in sorted_app_names:
-            print(f"\nMigrating {_app_name}")
-            print("------------------------------------------------")
-            manager = ForwardsMigrationManager(
-                app_name=_app_name, migration_id="all", fake=fake
-            )
-            await manager.run()
-    else:
-        manager = ForwardsMigrationManager(
-            app_name=app_name, migration_id=migration_id, fake=fake
-        )
-        await manager.run()
+    response = await run_forwards(
+        app_name=app_name, migration_id=migration_id, fake=fake
+    )
+
+    if not response.success:
+        sys.exit(1)
diff --git a/tests/apps/app/commands/test_show_all.py b/tests/apps/app/commands/test_show_all.py
@@ -1,12 +1,12 @@
 from unittest import TestCase
-from unittest.mock import call, patch
+from unittest.mock import call, patch, MagicMock
 
 from piccolo.apps.app.commands.show_all import show_all
 
 
 class TestShowAll(TestCase):
     @patch("piccolo.apps.app.commands.show_all.print")
-    def test_show_all(self, print_):
+    def test_show_all(self, print_: MagicMock):
         show_all()
 
         self.assertEqual(
diff --git a/tests/apps/migrations/commands/test_forwards_backwards.py b/tests/apps/migrations/commands/test_forwards_backwards.py
