88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
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: # type: ignore[name-defined]
|
|
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: # type: ignore[name-defined]
|
|
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"
|
|
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")
|