Built-in Database Migration Tool For Robyn

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
 
 [project.optional-dependencies]
 "templating" = ["jinja2 == 3.0.1"]
+"db" = ["alembic >= 1.16.3"]
 
 [project.urls]
 Documentation = "https://robyn.tech/"
@@ -79,11 +80,13 @@ watchdog = "4.0.1"
 multiprocess = "0.70.14"
 uvloop = { version = "0.21.0", markers = "sys_platform != 'win32' and (sys_platform != 'cygwin' and platform_python_implementation != 'PyPy')" }
 jinja2 = { version = "3.0.1", optional = true }
+alembic = { version = "^1.16.3", optional = true }
 rustimport = "^1.3.4"
 orjson = "^3.9.15"
 
 [tool.poetry.extras]
 templating = ["jinja2"]
+db = ["alembic"]
 
 [tool.poetry.group.dev]
 optional = true

robyn/argument_parser.py

@@ -84,6 +84,12 @@ def __init__(self) -> None:
             default=False,
             help="Fast mode. It sets the optimal values for processes, workers and log level. However, you can override them.",
         )
+        parser.add_argument(
+            "db",
+            nargs="?",
+            default=None,
+            help="Database migration commands. Use 'robyn db' to see more information.",
+        )
 
         args, unknown_args = parser.parse_known_args()
         self.fast = args.fast
@@ -99,6 +105,7 @@ def __init__(self) -> None:
         self.file_path = None
         self.disable_openapi = args.disable_openapi
         self.log_level = args.log_level
+        self.db = args.db
 
         if self.fast:
             # doing this here before every other check

robyn/cli.py

@@ -3,14 +3,18 @@
 import subprocess
 import sys
 import webbrowser
+import argparse
+import importlib.util
 from pathlib import Path
 from typing import Optional
+from pip._internal.cli.main import main as pip_main
 
 from InquirerPy.base.control import Choice
 from InquirerPy.resolver import prompt
 
 from robyn.env_populator import load_vars
 from robyn.robyn import get_version
+from robyn.migrate import configure_parser, execute_command
 
 from .argument_parser import Config
 from .reloader import create_rust_file, setup_reloader
@@ -101,6 +105,59 @@ def start_app_normally(config: Config):
     subprocess.run(command, start_new_session=False)
 
 
+def handle_db_command():
+    """Handle database migration commands."""
+    alembic_spec = importlib.util.find_spec("alembic")
+
+    if alembic_spec is None:
+        print("ERROR: Alembic has not been installed.")
+        install_choice = input("Would you like to install alembic now? (y/n): ").strip().lower()
+
+        if install_choice == "y":
+            try:
+                try:
+                    print("Installing alembic...")
+                    pip_main(["install", "alembic", "--quiet"])
+                    print("Successfully installed alembic.")
+                except ImportError:
+                    print("Installing alembic using subprocess...")
+                    subprocess.run([sys.executable, "-m", "pip", "install", "alembic", "-q"], check=True)
+                    print("Successfully installed alembic.")
+
+                importlib.invalidate_caches()
+                alembic_spec = importlib.util.find_spec("alembic")
+                if alembic_spec is None:
+                    print("ERROR: Failed to install alembic. Please install it manually using 'pip install alembic'.")
+                    sys.exit(1)
+            except Exception as e:
+                print(f"ERROR: Failed to install alembic: {str(e)}")
+                print("Please install it manually using 'pip install alembic'.")
+                sys.exit(1)
+        else:
+            print("Please install alembic manually using 'pip install alembic' before using database commands.")
+            sys.exit(1)
+
+    parser = argparse.ArgumentParser(
+        usage=argparse.SUPPRESS,  # omit usage hint
+        description="Robyn database migration commands.",
+    )
+    parser = configure_parser(parser)
+
+    if len(sys.argv) == 2 and sys.argv[1] == "db":
+        parser.print_help()
+        sys.exit(1)
+    # Remove the first two arguments (robyn and db)
+    if len(sys.argv) > 2 and sys.argv[1] == "db":
+        if sys.argv[2] == "--help" or sys.argv[2] == "-h" or sys.argv[2] == "-H":
+            parser.print_help()
+            sys.exit(1)
+        db_args = parser.parse_args(sys.argv[2:])
+        execute_command(db_args)
+    else:
+        print("ERROR: Invalid command. Please run 'robyn db' to see more information.")
+        sys.exit(1)
+
+
 def run():
     config = Config()
 
@@ -113,6 +170,11 @@ def run():
     if config.dev is None:
         config.dev = os.getenv("ROBYN_DEV_MODE", False) == "True"
 
+    # Handle db command
+    if config.db == "db" and len(sys.argv) > 1 and sys.argv[1] == "db":
+        handle_db_command()
+        return
+
     if config.create:
         create_robyn_app()