diff --git a/.gitignore b/.gitignore index 9b09122..230bbd6 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ __pycache__ /.vscode dist *.egg-info +.newt diff --git a/newt/commands/migrate/init.py b/newt/commands/migrate/init.py index 49c5e85..02d0639 100644 --- a/newt/commands/migrate/init.py +++ b/newt/commands/migrate/init.py @@ -1,5 +1,7 @@ +import re import shutil import subprocess +import tempfile from pathlib import Path from typing import Annotated @@ -7,34 +9,66 @@ app = typer.Typer(no_args_is_help=True) + @app.command() def init( - path: Annotated[str, typer.Argument(help="Relative path to the alembic directory")] = "alembic", - template: Annotated[str, typer.Option("--template", "-t", help="Custom template for the alembic directory")] = None, + path: Annotated[ + str, typer.Argument(help="Relative path to the alembic directory") + ] = "alembic", + template: Annotated[ + str, + typer.Option( + "--template", "-t", help="Custom template for the alembic directory" + ), + ] = None, ): """ - Creates a new alembic directory in the current or given path. - If alembic is already initialized in the specified directory, it will ask if the user wants to reinitialize it. - If the user chooses to reinitialize, it will remove the existing alembic directory and create a new one. - If the user chooses not to reinitialize, it will exit the program. + Creates a new alembic directory structure with alembic.ini at the top level + and migration files in an 'alembic' subdirectory. """ - # Check if alembic is already initialized in the current directory - project_dir = Path(f"{Path.cwd()}/{path}") + project_dir = Path.cwd() / path - #alembic_dir = project_dir / "alembic" + # Check if directory exists and handle reinitialize if project_dir.exists(): - reinitialize = typer.confirm("Alembic is already initialized. Do you want to reinitialize it?") - if not reinitialize: + if not typer.confirm( + "Alembic is already initialized. Do you want to reinitialize it?" + ): raise typer.Abort() - else: - shutil.rmtree(project_dir) + shutil.rmtree(project_dir) + project_dir.mkdir(parents=True, exist_ok=True) - # Create command - command = ["alembic", "init", str(project_dir)] - if template: - command.extend(["--template", template]) + # Create alembic structure in temporary directory + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Build and run alembic init command + command = ["alembic", "init", str(temp_path / "alembic")] + if template: + command.extend(["--template", template]) + + try: + subprocess.run(command, cwd=temp_dir, check=True) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + typer.echo(f"Error: {e}", err=True) + typer.echo( + "Please ensure alembic is installed: pip install alembic", err=True + ) + raise typer.Abort() + + # Move alembic directory to target location + shutil.move(str(temp_path / "alembic"), str(project_dir / "alembic")) + + # Move and update alembic.ini + with open(temp_path / "alembic.ini", "r") as f: + ini_content = f.read() - # Run alembic init command from a subfolder - subprocess.run(command) + # Update script_location to point to alembic subdirectory + ini_content = re.sub( + r"script_location = .*", + "script_location = %(here)s/alembic", + ini_content, + ) + with open(project_dir / "alembic.ini", "w") as f: + f.write(ini_content)