piket/piket_server/seed.py

"""
Provides functions to manage the database while the server is offline.
"""

import argparse
import csv
import os

from piket_server import db, Person, Settlement, ConsumptionType, Consumption


def main():
    """ Main entry point. """

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()

    # Clear command
    parser_clear = subparsers.add_parser("clear", help="Clear the database.")
    parser_clear.set_defaults(func=cmd_clear)
    parser_clear.add_argument("--removemydata", action="store_true")

    # Parser load_seeds
    # TODO

    args = parser.parse_args()
    args.func(args)


def cmd_clear(args) -> None:
    """ Entry point for 'clear' subcommand. """

    if not args.removemydata:
        print("WARNING! This command will delete all contents in your database!")
        print("Type 'removemydata' to continue, anything else or CTRL-C to abort.")

        confirmation = input("> ")

        do_wipe = confirmation == "removemydata"

    else:
        do_wipe = True

    if do_wipe:
        print("Dropping all tables...")
        db.drop_all()

        print("All data removed. Recreating database...")
        db.create_all()

        from alembic.config import Config
        from alembic import command

        alembic_cfg = Config(os.path.join(os.path.dirname(__file__), "alembic.ini"))
        command.stamp(alembic_cfg, "head")

        print("Done.")
        return

    print("Aborting.")


if __name__ == "__main__":
    main()