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.models import Person, Settlement, ConsumptionType, Consumption
from piket_server.flask import db


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
    parser_load_persons = subparsers.add_parser(
        "load_persons", help="Load Person data."
    )
    parser_load_persons.set_defaults(func=cmd_load_persons)
    parser_load_persons.add_argument("datafile")

    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  # type: ignore
        from alembic import command  # type: ignore

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

        print("Done.")
        return

    print("Aborting.")


def cmd_load_persons(args) -> None:
    """ Entrypoint for 'load_persons" subcommand. """

    print("Loading persons.")

    with open(args.datafile) as csvfile:
        reader = csv.DictReader(csvfile)

        for row in reader:
            p = Person(name=row["name"], active=bool(int(row["active"])))
            db.session.add(p)

    db.session.commit()

    print("Done.")


if __name__ == "__main__":
    main()