""" 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 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 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.") 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()