97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
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 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,
|
|
"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")
|
|
# 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 e: # pylint: disable=broad-except
|
|
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")
|