refactoring migration commands to work with auto

Author: dantownsend

piccolo/commands/app/new.py

@@ -3,6 +3,7 @@
 import sys
 import typing as t
 
+import black
 import click
 import jinja2
 
@@ -34,8 +35,11 @@ def new_app(app_name: str):
     for filename, context in templates.items():
         with open(os.path.join(app_name, filename), "w") as f:
             template = JINJA_ENV.get_template(filename + ".jinja")
-            rendered = template.render(**context)
-            f.write(rendered)
+            file_contents = template.render(**context)
+            file_contents = black.format_str(
+                file_contents, mode=black.FileMode(line_length=82)
+            )
+            f.write(file_contents)
 
     migrations_folder_path = os.path.join(app_name, "piccolo_migrations")
     os.mkdir(migrations_folder_path)

piccolo/commands/app/templates/piccolo_app.py.jinja

@@ -1,11 +1,19 @@
+"""
+Import all of the Tables subclasses in your app here, and register them with
+the APP_CONFIG.
+"""
+
 import os
 
-APP_NAME = '{{ app_name }}'
+from piccolo.conf.apps import AppConfig
 
-CURRENT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
 
-TABLES_MODULE = 'tables'
+CURRENT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
 
-MIGRATIONS_FOLDER = os.path.join(CURRENT_DIRECTORY, 'piccolo_migrations')
 
-MIGRATION_DEPENDENCIES = []
+APP_CONFIG = AppConfig(
+    app_name='{{ app_name }}',
+    migrations_folder_path=os.path.join(CURRENT_DIRECTORY, 'piccolo_migrations'),
+    table_classes=[],
+    migration_dependencies=[],
+)

piccolo/commands/conf/new.py

@@ -2,6 +2,7 @@
 import os
 import sys
 
+import black
 import click
 import jinja2
 
@@ -26,8 +27,12 @@ def new_piccolo_conf(engine_name: str, force: bool = False):
 
     with open("piccolo_conf.py", "w") as f:
         template = JINJA_ENV.get_template("piccolo_conf.py.jinja")
-        rendered = template.render(engine_name=engine_name)
-        f.write(rendered)
+        file_contents = template.render(engine_name=engine_name)
+        file_contents = black.format_str(
+            file_contents, mode=black.FileMode(line_length=82)
+        )
+
+        f.write(file_contents)
 
 
 @click.option(

piccolo/commands/conf/templates/piccolo_conf.py.jinja

@@ -1,5 +1,4 @@
-from piccolo.conf.app import AppConfig
-
+from piccolo.conf.apps import AppRegistry
 {% if engine_name == 'postgres' %}
 from piccolo.engine.postgres import PostgresEngine
 
@@ -14,9 +13,6 @@ DB = SQLiteEngine(path='project.sqlite')
 {% endif %}
 
 
-MIGRATIONS = []
-
-
-# A list of paths to piccolo apps, which are basically just Python packages
-# containing a piccolo_app.py file.
-APPS = AppConfig()
+# A list of paths to piccolo apps/
+# e.g. ['blog.piccolo_app']
+APP_REGISTRY = AppRegistry(apps=[])

piccolo/commands/migration/backwards.py

@@ -1,5 +1,4 @@
 import asyncio
-import os
 import sys
 
 import click
@@ -14,19 +13,23 @@ def __init__(self, app_name: str, migration_name: str):
         self.app_name = app_name
 
     def run(self):
-        config_modules = self.get_config_modules()
+        app_modules = self.get_app_modules()
 
         migration_modules = []
 
-        for config_module in config_modules:
-            if self.get_app_name(config_module) == self.app_name:
+        for app_module in app_modules:
+            app_config = getattr(app_module, "APP_CONFIG")
+            if app_config.app_name == self.app_name:
                 migration_modules = self.get_migration_modules(
-                    os.path.dirname(config_module.__file__)
+                    app_config.migrations_folder_path
                 )
                 break
 
         migration_ids = self.get_migration_ids(migration_modules)
 
+        if self.migration_name == "all":
+            self.migration_name = migration_ids[0]
+
         if self.migration_name not in migration_ids:
             print(
                 f"Unrecognized migration name - must be one of {migration_ids}"
@@ -70,7 +73,8 @@ def run(self):
 @click.argument("migration_name")
 def backwards(migration_name: str, app_name: str):
     """
-    Undo migrations up to a specific migrations.
+    Undo migrations up to a specific migrations. Specify a migration_name of
+    'all' to undo all of the migrations.
 
     - make sure the migration name is valid
     - work out which to undo, and in which order

piccolo/commands/migration/base.py

@@ -5,6 +5,7 @@
 from types import ModuleType
 import typing as t
 
+from piccolo.conf.apps import AppConfig, AppRegistry
 from piccolo.migrations.tables import Migration
 
 
@@ -18,9 +19,8 @@ async def backwards() -> None:
         pass
 
 
-class ConfigModule(ModuleType):
-    NAME: str
-    DEPENDENCIES: t.List[str]
+class PiccoloAppModule(ModuleType):
+    APP_CONFIG: AppConfig
 
 
 class BaseMigrationManager:
@@ -35,52 +35,53 @@ def create_migration_table(self) -> bool:
         return False
 
     def _deduplicate(
-        self, config_modules: t.List[ConfigModule]
-    ) -> t.List[ConfigModule]:
+        self, config_modules: t.List[PiccoloAppModule]
+    ) -> t.List[PiccoloAppModule]:
         """
         Remove all duplicates - just leaving the first instance.
         """
         # Deduplicate, but preserve order - which is why set() isn't used.
         return list(dict([(c, None) for c in config_modules]).keys())
 
-    def import_config_modules(
-        self, migrations: t.List[str]
-    ) -> t.List[ConfigModule]:
+    def _import_app_modules(
+        self, config_module_paths: t.List[str]
+    ) -> t.List[PiccoloAppModule]:
         """
-        Import all config modules, and all dependencies.
+        Import all piccolo_app.py modules, and all dependencies.
         """
         config_modules = []
 
-        for config_module_path in migrations:
+        for config_module_path in config_module_paths:
             try:
                 config_module = t.cast(
-                    ConfigModule, importlib.import_module(config_module_path)
+                    PiccoloAppModule,
+                    importlib.import_module(config_module_path),
                 )
             except ImportError:
                 raise Exception(f"Unable to import {config_module_path}")
-            DEPENDENCIES = getattr(config_module, "DEPENDENCIES", [])
-            dependency_config_modules = self.import_config_modules(
-                DEPENDENCIES
+            app_config: AppConfig = getattr(config_module, "APP_CONFIG")
+            dependency_config_modules = self._import_app_modules(
+                app_config.migration_dependencies
             )
             config_modules.extend(dependency_config_modules + [config_module])
 
         return config_modules
 
-    def get_config_modules(self) -> t.List[ConfigModule]:
+    def get_app_modules(self) -> t.List[PiccoloAppModule]:
         try:
-            from piccolo_conf import MIGRATIONS
+            from piccolo_conf import APP_REGISTRY
         except ImportError:
             raise Exception(
-                "Unable to import MIGRATIONS from piccolo_conf - make sure "
+                "Unable to import APP_REGISTRY from piccolo_conf - make sure "
                 "it's in your path."
             )
 
-        config_modules = self.import_config_modules(MIGRATIONS)
+        app_modules = self._import_app_modules(APP_REGISTRY.apps)
 
         # Now deduplicate any dependencies
-        config_modules = self._deduplicate(config_modules)
+        app_modules = self._deduplicate(app_modules)
 
-        return config_modules
+        return app_modules
 
     def get_migration_modules(
         self, folder_path: str
@@ -93,7 +94,7 @@ def get_migration_modules(
         sys.path.insert(0, folder_path)
 
         folder_contents = os.listdir(folder_path)
-        excluded = ("__init__.py", "__pycache__", "config.py")
+        excluded = ("__init__.py", "__pycache__")
         migration_names = [
             i.split(".py")[0]
             for i in folder_contents
@@ -116,7 +117,7 @@ def get_migration_ids(
     ) -> t.List[str]:
         return sorted(list(migration_module_dict.keys()))
 
-    def get_app_name(self, config_module: ConfigModule) -> str:
+    def get_app_name(self, config_module: PiccoloAppModule) -> str:
         # TODO - default to package name instead.
         app_name = getattr(config_module, "NAME", "").strip()
         return app_name

piccolo/commands/migration/check.py

@@ -1,5 +1,3 @@
-import os
-
 import click
 
 from .base import BaseMigrationManager
@@ -18,14 +16,17 @@ def get_fixed_length_string(self, string: str, length=20) -> str:
     def run(self):
         print("Listing migrations ...")
 
-        config_modules = self.get_config_modules()
+        app_modules = self.get_app_modules()
+
+        for app_module in app_modules:
+            app_config = app_module.APP_CONFIG
 
-        for config_module in config_modules:
-            app_name = self.get_app_name(config_module)
+            app_name = app_config.app_name
             fixed_length_app_name = self.get_fixed_length_string(app_name)
 
-            path = os.path.dirname(config_module.__file__)
-            migration_modules = self.get_migration_modules(path)
+            migration_modules = self.get_migration_modules(
+                app_config.migrations_folder_path
+            )
             ids = self.get_migration_ids(migration_modules)
             for _id in ids:
                 has_ran = (

piccolo/commands/migration/forwards.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 import asyncio
-import os
 import pprint
 import typing as t
 
@@ -9,56 +8,63 @@
 from piccolo.commands.migration.base import (
     BaseMigrationManager,
     MigrationModule,
-    ConfigModule,
+    PiccoloAppModule,
 )
+from piccolo.conf.apps import AppConfig
 from piccolo.migrations.tables import Migration
+from piccolo.migrations.auto import MigrationManager
 
 
 class ForwardsMigrationManager(BaseMigrationManager):
-    def run_migrations(self, config_modules: t.List[ConfigModule]) -> None:
+    def __init__(self, app_name: str, *args, **kwargs):
+        self.app_name = app_name
+
+    def run_migrations(self, app_modules: t.List[PiccoloAppModule]) -> None:
         already_ran = Migration.get_migrations_which_ran()
         print(f"Already ran:\n{already_ran}\n")
 
-        for config_module in config_modules:
-            migrations_folder = os.path.dirname(config_module.__file__)
+        for app_module in app_modules:
+            app_config: AppConfig = getattr(app_module, "APP_CONFIG")
 
             migration_modules: t.Dict[
                 str, MigrationModule
-            ] = self.get_migration_modules(migrations_folder)
+            ] = self.get_migration_modules(app_config.migrations_folder_path)
 
             ids = self.get_migration_ids(migration_modules)
             print(f"Migration ids = {ids}")
 
-            app_name = self.get_app_name(config_module)
-
             havent_run = sorted(set(ids) - set(already_ran))
             for _id in havent_run:
                 migration_module = migration_modules[_id]
-                asyncio.run(migration_module.forwards())
+                response = asyncio.run(migration_module.forwards())
+
+                if isinstance(response, MigrationManager):
+                    asyncio.run(response.run())
 
                 print(f"Ran {_id}")
                 Migration.insert().add(
-                    Migration(name=_id, app_name=app_name)
+                    Migration(name=_id, app_name=self.app_name)
                 ).run_sync()
 
     def run(self):
         print("Running migrations ...")
         self.create_migration_table()
 
-        config_modules = self.get_config_modules()
+        app_modules = self.get_app_modules()
 
         print("Config Modules:")
-        pprint.pprint(config_modules)
+        pprint.pprint(app_modules)
         print("\n")
 
-        self.run_migrations(config_modules)
+        self.run_migrations(app_modules)
 
 
 @click.command()
-def forwards():
+@click.argument("app_name")
+def forwards(app_name: str):
     """
     Runs any migrations which haven't been run yet, or up to a specific
     migration.
     """
-    manager = ForwardsMigrationManager()
+    manager = ForwardsMigrationManager(app_name=app_name)
     manager.run()