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 BridgeConf from app.models.mirrors import Mirror, Origin, Proxy from app.models.alarms import Alarm models = { "bridge": 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 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")