1import os
  2import sys
  3
  4import click
  5
  6from ..services import Services
  7from .container import DBContainer
  8
  9
 10@click.group("db")
 11def cli():
 12    """Start, stop, and manage the local Postgres database"""
 13    pass
 14
 15
 16# @cli.command()
 17# def reset():
 18#     DBContainer().reset(create=True)
 19#     click.secho("Local development database reset", fg="green")
 20
 21
 22@cli.command()
 23@click.argument("export_path", default="")
 24def export(export_path):
 25    """Export the local database to a file"""
 26    if not export_path:
 27        current_dir_name = os.path.basename(os.getcwd())
 28        export_path = f"{current_dir_name}-dev-db.sql"
 29    with Services():
 30        export_successful = DBContainer().export(export_path)
 31
 32    if export_successful:
 33        click.secho(f"Local development database exported to {export_path}", fg="green")
 34    else:
 35        click.secho("Export failed", fg="red")
 36        sys.exit(1)
 37
 38
 39@cli.command("import")
 40@click.argument("sql_file")
 41def import_db(sql_file):
 42    """Import a database file into the local database"""
 43
 44    print(f"Importing {sql_file} ({os.path.getsize(sql_file) / 1024 / 1024:.2f} MB)")
 45
 46    with Services():
 47        successful = DBContainer().import_sql(sql_file)
 48
 49    if successful:
 50        click.secho(f"Local development database imported from {sql_file}", fg="green")
 51    else:
 52        click.secho("Import failed", fg="red")
 53        sys.exit(1)
 54
 55
 56@cli.group()
 57def snapshot():
 58    """Manage local database snapshots"""
 59    pass
 60
 61
 62@snapshot.command("create")
 63@click.argument("name")
 64@click.pass_context
 65def snapshot_create(ctx, name):
 66    """Create a snapshot of the main database"""
 67    created = DBContainer().create_snapshot(name)
 68    if not created:
 69        click.secho(f'Snapshot "{name}" already exists', fg="red")
 70        sys.exit(1)
 71
 72    click.secho(f'Snapshot "{name}" created', fg="green")
 73    print()
 74    ctx.invoke(snapshot_list)
 75
 76
 77@snapshot.command("list")
 78def snapshot_list():
 79    """List all snapshots"""
 80    DBContainer().list_snapshots()
 81
 82
 83@snapshot.command("restore")
 84@click.argument("name")
 85@click.option("--yes", "-y", is_flag=True)
 86def snapshot_restore(name, yes):
 87    """Restore a snapshot to the main database"""
 88    if not yes:
 89        click.confirm(
 90            f'Are you sure you want to restore snapshot "{name}" to the main database?',
 91            abort=True,
 92        )
 93
 94    DBContainer().restore_snapshot(name)
 95    click.secho(f'Snapshot "{name}" restored', fg="green")
 96
 97
 98@snapshot.command("delete")
 99@click.argument("name")
100@click.pass_context
101def snapshot_delete(ctx, name):
102    """Delete a snapshot"""
103    deleted = DBContainer().delete_snapshot(name)
104    if not deleted:
105        click.secho(f'Snapshot "{name}" does not exist', fg="red")
106        sys.exit(1)
107    click.secho(f'Snapshot "{name}" deleted', fg="green")
108    print()
109    ctx.invoke(snapshot_list)
110
111
112if __name__ == "__main__":
113    cli()