majuna/app/cli/db.py

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")