import argparse import csv import datetime import logging import sys from typing import TYPE_CHECKING, Any 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 if TYPE_CHECKING: _SubparserType = argparse._SubParsersAction[argparse.ArgumentParser] else: _SubparserType = Any models = { "bridge": Bridge, "bridgeconf": BridgeConf, "alarm": Alarm, "group": Group, "list": MirrorList, "mirror": Mirror, "origin": Origin, "proxy": Proxy } def export(model: db.Model) -> None: 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) -> None: 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 # type: ignore else: line[i] = datetime.datetime.strptime(line[i], "%Y-%m-%d %H:%M:%S.%f") # type: ignore elif header[i] in ["eotk"]: # boolean fields line[i] = line[i] == "True" # type: ignore elif header[i].endswith("_id") and line[i] == "": # integer foreign keys line[i] = None # type: ignore 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: _SubparserType) -> 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: argparse.Namespace) -> None: self.args = args def run(self) -> None: 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")