Improve migration script with prompts, warnings, finish TODOs

backend/linkpulse/migrate.py

@@ -1,6 +1,9 @@
+import os
 import pkgutil
+import re
 import sys
-from typing import Any, List, Optional
+import questionary
+from typing import Any, List, Optional, Tuple
 from dotenv import load_dotenv
 from peewee_migrate import Router, router
 from peewee import PostgresqlDatabase
@@ -10,7 +13,7 @@ from linkpulse.formatting import pluralize
 load_dotenv(dotenv_path=".env")
 
 class ExtendedRouter(Router):
-    def show(self, module: str) -> Optional[str]:
+    def show(self, module: str) -> Optional[Tuple[str, str]]:
         """
         Show the suggested migration that will be created, without actually creating it.
 
@@ -65,10 +68,6 @@ def main(*args: str) -> None:
     router = ExtendedRouter(database=db, migrate_dir='linkpulse/migrations', ignore=[models.BaseModel._meta.table_name])
     auto = 'linkpulse.models'
 
-    # TODO: Show unapplied migrations before applying all
-    # TODO: Suggest merging migrations if many are present + all applied
-    # TODO: Show prepared migration before naming (+ confirmation option for pre-provided name)
-
     current = router.all_migrations()
     if len(current) == 0:
         diff = router.diff
@@ -82,15 +81,27 @@ def main(*args: str) -> None:
             else:
                 print(f"Migration created: {migration}")
                 router.run(migration)
+    
+    diff = router.diff
+    if len(diff) > 0:
+        print('Note: Selecting a migration will apply all migrations up to and including the selected migration.')
+        print('e.g. Applying 004 while only 001 is applied would apply 002, 003, and 004.')
+        
+        choice = questionary.select("Select highest migration to apply:", choices=diff).ask()
+        if choice is None:
+            print("For safety reasons, you won't be able to create migrations without applying the pending ones.")
+            if len(current) == 0:
+                print("Warn: No migrations have been applied globally, which is dangerous. Something may be wrong.")
+            return
+        
+        result = router.run(choice)
+        print(f"Done. Applied migrations: {result}")
+        print("Warning: You should commit and push any new migrations immediately!")
     else:
-            print("{} migration{} found, applying all ({}).".format(len(diff), pluralize(len(diff)), ', '.join(diff)))
+        print("No pending migrations to apply.")
-            applied = router.run()
-            print('Done ({}).'.format(', '.join(applied)))
-    else:
-        print('No migrations found, all migrations applied.')
 
     migration_available = router.show(auto)
-    if migration_available:
+    if migration_available is not None:
         print("A migration is available to be applied:")
         migrate_text, rollback_text = migration_available
         
@@ -105,6 +116,35 @@ def main(*args: str) -> None:
                 continue
             print('\t' + line)
         
+        if questionary.confirm("Do you want to create this migration?").ask():
+            print('Lowercase letters and underscores only (e.g. "create_table", "remove_ipaddress_count").')
+            migration_name: Optional[str] = questionary.text("Enter migration name", validate=lambda text: re.match("^[a-z_]+$", text) is not None).ask()
+
+            if migration_name is None:
+                return
+            
+            migration = router.create(migration_name, auto=auto)
+            if migration:
+                print(f"Migration created: {migration}")
+                if len(router.diff) == 1:
+                    if questionary.confirm("Do you want to apply this migration immediately?").ask():
+                        router.run(migration)
+                        print('Done.')
+                        print("!!! Commit and push this migration file immediately!")
+            else:
+                print("No changes detected. Something went wrong.")
+                return
+    else:
+        print("No database changes detected.")
+
+    if len(current) > 5:
+        if questionary.confirm("There are more than 5 migrations applied. Do you want to merge them?", default=False).ask():
+            print("Merging migrations...")
+            router.merge(name="initial")
+            print("Done.")
+
+            print("!!! Commit and push this merged migration file immediately!")
+
     # Testing Code:
 """
     print(router.print('linkpulse.models'))

backend/poetry.lock

@@ -332,6 +332,20 @@ files = [
     {file = "priority-2.0.0.tar.gz", hash = "sha256:c965d54f1b8d0d0b19479db3924c7c36cf672dbf2aec92d43fbdaf4492ba18c0"},
 ]
 
+[[package]]
+name = "prompt-toolkit"
+version = "3.0.36"
+description = "Library for building powerful interactive command lines in Python"
+optional = false
+python-versions = ">=3.6.2"
+files = [
+    {file = "prompt_toolkit-3.0.36-py3-none-any.whl", hash = "sha256:aa64ad242a462c5ff0363a7b9cfe696c20d55d9fc60c11fd8e632d064804d305"},
+    {file = "prompt_toolkit-3.0.36.tar.gz", hash = "sha256:3e163f254bef5a03b146397d7c1963bd3e2812f0964bb9a24e6ec761fd28db63"},
+]
+
+[package.dependencies]
+wcwidth = "*"
+
 [[package]]
 name = "psutil"
 version = "6.0.0"
@@ -513,6 +527,20 @@ files = [
 [package.extras]
 cli = ["click (>=5.0)"]
 
+[[package]]
+name = "questionary"
+version = "2.0.1"
+description = "Python library to build pretty command line user prompts ⭐️"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "questionary-2.0.1-py3-none-any.whl", hash = "sha256:8ab9a01d0b91b68444dff7f6652c1e754105533f083cbe27597c8110ecc230a2"},
+    {file = "questionary-2.0.1.tar.gz", hash = "sha256:bcce898bf3dbb446ff62830c86c5c6fb9a22a54146f0f5597d3da43b10d8fc8b"},
+]
+
+[package.dependencies]
+prompt_toolkit = ">=2.0,<=3.0.36"
+
 [[package]]
 name = "six"
 version = "1.16.0"
@@ -614,6 +642,17 @@ h11 = ">=0.8"
 [package.extras]
 standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1)", "watchfiles (>=0.13)", "websockets (>=10.4)"]
 
+[[package]]
+name = "wcwidth"
+version = "0.2.13"
+description = "Measures the displayed width of unicode strings in a terminal"
+optional = false
+python-versions = "*"
+files = [
+    {file = "wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859"},
+    {file = "wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5"},
+]
+
 [[package]]
 name = "wsproto"
 version = "1.2.0"
@@ -631,4 +670,4 @@ h11 = ">=0.9.0,<1"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.12"
-content-hash = "297e93d8b8d987b9e98575c94874482c931fad147925d75f784fbc5138b1dad0"
+content-hash = "5f79b30ae45d71568bbce44ede4d93ac6e0c0f580a818cf01e919745d23fa091"

backend/pyproject.toml

@@ -20,6 +20,7 @@ peewee-migrate = "^1.13.0"
 types-peewee = "^3.17.7.20241017"
 types-psycopg2 = "^2.9.21.20241019"
 fastapi-cache2 = "^0.2.2"
+questionary = "^2.0.1"
 
 
 [tool.poetry.group.dev.dependencies]