import csv import datetime import logging import sys from app import app from app.cli import _SubparserType, BaseCliHandler from app.extensions import db from app.models.base import Group, MirrorList from app.models.bridges import Bridge, BridgeConf from app.models.mirrors import Origin, Proxy from app.models.alarms import Alarm, AlarmState models = { "bridge": Bridge, "bridgeconf": BridgeConf, "alarm": Alarm, "group": Group, "list": MirrorList, "origin": Origin, "proxy": Proxy } def export(model: db.Model) -> None: out = csv.writer(sys.stdout) out.writerow(model.csv_header()) for row in model.query.all(): out.writerow(row.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 new_entity = model() for idx, field_name in header: if field_name in ["added", "updated", "destroyed", "deprecated", "last_updated", "terraform_updated"]: # datetime fields if line[idx] == "": line[idx] = None # type: ignore else: line[idx] = datetime.datetime.strptime(line[idx], "%Y-%m-%d %H:%M:%S.%f") # type: ignore elif field_name in ["eotk", "auto_rotation", "smart"]: # boolean fields line[idx] = line[idx] == "True" # type: ignore elif field_name.endswith("_id") and line[idx] == "": # integer foreign keys line[idx] = None # type: ignore elif field_name in ["alarm_state"]: # alarm states line[idx] = getattr(AlarmState, line[idx][len("AlarmState."):]) setattr(new_entity, field_name, line[idx]) db.session.add(new_entity) db.session.commit() logging.info("Import completed successfully") # Many things can go wrong in the above, like IO, format or database errors. # We catch all the errors and ensure the database transaction is rolled back, and log it. except Exception as exc: # pylint: disable=broad-except logging.exception(exc) db.session.rollback() class DbCliHandler(BaseCliHandler): @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 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")