lint: tidy up code in cli module

This commit is contained in:
Iain Learmonth 2022-06-17 13:21:35 +01:00
parent 5f7733d064
commit 98895e47de
6 changed files with 61 additions and 58 deletions

View file

@ -1,22 +1,16 @@
import argparse
import csv
import datetime
import logging
import sys
from typing import TYPE_CHECKING, Any
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
if TYPE_CHECKING:
_SubparserType = argparse._SubParsersAction[argparse.ArgumentParser]
else:
_SubparserType = Any
models = {
"bridge": Bridge,
"bridgeconf": BridgeConf,
@ -31,8 +25,8 @@ models = {
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())
for row in model.query.all():
out.writerow(row.csv_row())
def impot(model: db.Model) -> None:
@ -46,35 +40,35 @@ def impot(model: db.Model) -> None:
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"]:
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[i] == "":
line[i] = None # type: ignore
if line[idx] == "":
line[idx] = 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", "auto_rotation", "smart"]:
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[i] = line[i] == "True" # type: ignore
elif header[i].endswith("_id") and line[i] == "":
line[idx] = line[idx] == "True" # type: ignore
elif field_name.endswith("_id") and line[idx] == "":
# integer foreign keys
line[i] = None # type: ignore
elif header[i] in ["alarm_state"]:
line[idx] = None # type: ignore
elif field_name in ["alarm_state"]:
# alarm states
line[i] = getattr(AlarmState, line[i][len("AlarmState."):])
setattr(x, header[i], line[i])
db.session.add(x)
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 e: # pylint: disable=broad-except
logging.exception(e)
except Exception as exc: # pylint: disable=broad-except
logging.exception(exc)
db.session.rollback()
class DbCliHandler:
class DbCliHandler(BaseCliHandler):
@classmethod
def add_subparser_to(cls, subparsers: _SubparserType) -> None:
parser = subparsers.add_parser("db", help="database operations")
@ -84,9 +78,6 @@ class DbCliHandler:
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: