75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
|
import argparse
|
||
|
|
import csv
|
||
|
|
import datetime
|
||
|
|
import logging
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from app import app
|
||
|
|
from app.extensions import db
|
||
|
|
from app.models import Group, Origin, Proxy, BridgeConf, Alarm
|
||
|
|
|
||
|
|
models = {
|
||
|
|
"group": Group,
|
||
|
|
"origin": Origin,
|
||
|
|
"proxy": Proxy,
|
||
|
|
"bridge": BridgeConf,
|
||
|
|
"alarm": Alarm
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
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")
|
||
|
|
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=["group", "origin", "proxy", "bridge"],
|
||
|
|
help="export data to CSV format")
|
||
|
|
parser.add_argument("--import", choices=["group", "origin", "proxy", "bridge"],
|
||
|
|
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")
|