import argparse import csv import datetime import logging import sys from app import app from app.extensions import db from app.models.base import Group, MirrorList from app.models.bridges import Bridge, BridgeConf from app.models.mirrors import Mirror, Origin, Proxy from app.models.alarms import Alarm, AlarmState models = { "bridge": Bridge, "bridgeconf": BridgeConf, "alarm": Alarm, "group": Group, "list": MirrorList, "mirror": Mirror, "origin": Origin, "proxy": Proxy } def export(model: db.Model): out = csv.writer(sys.stdout) out.writerow(model.csv_header()) for r in model.query.all(): out.writerow(r.csv_row()) def impot(model: db.Model): first = True header = model.csv_header() try: for line in csv.reader(sys.stdin): if first: if line != header: logging.error("CSV header mismatch") sys.exit(1) first = False continue x = model() for i in range(len(header)): if header[i] in ["added", "updated", "destroyed", "deprecated", "last_updated", "terraform_updated"]: # datetime fields if line[i] == "": line[i] = None else: line[i] = datetime.datetime.strptime(line[i], "%Y-%m-%d %H:%M:%S.%f") elif header[i] in ["eotk"]: # boolean fields line[i] = line[i] == "True" elif header[i].endswith("_id") and line[i] == "": # integer foreign keys line[i] = None elif header[i] in ["alarm_state"]: # alarm states line[i] = getattr(AlarmState, line[i][len("AlarmState."):]) setattr(x, header[i], line[i]) db.session.add(x) db.session.commit() logging.info("Import completed successfully") except Exception as e: logging.exception(e) db.session.rollback() class DbCliHandler: @classmethod def add_subparser_to(cls, subparsers: argparse._SubParsersAction) -> None: parser = subparsers.add_parser("db", help="database operations") parser.add_argument("--export", choices=sorted(models.keys()), help="export data to CSV format") parser.add_argument("--import", choices=sorted(models.keys()), help="import data from CSV format", dest="impot") parser.set_defaults(cls=cls) def __init__(self, args): self.args = args def run(self): with app.app_context(): if self.args.export: export(models[self.args.export]) elif self.args.impot: impot(models[self.args.impot]) else: logging.error("No action requested")