majuna/app/cli/db.py

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"]:
                    if line[i] == "":
                        line[i] = None
                    else:
                        line[i] = datetime.datetime.strptime(line[i], "%Y-%m-%d %H:%M:%S.%f")
                if header[i] in ["eotk"]:
                    line[i] = line[i] == "True"
                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")