diff --git a/app/collect/database.py b/app/collect/database.py index 63ccb58..f66d1f9 100644 --- a/app/collect/database.py +++ b/app/collect/database.py @@ -25,7 +25,7 @@ def upsert(session, model, rows, update_cols): def update_device_infos(session, address_origin, path=None): - if address_origin == DeviceInfoOrigin.flarmnet: + if address_origin == DeviceInfoOrigin.FLARMNET: device_infos = get_flarmnet(fln_file=path) else: device_infos = get_ddb(csv_file=path) @@ -49,7 +49,7 @@ def import_ddb(session, logger=None): logger = app.logger logger.info("Import registered devices fom the DDB...") - counter = update_device_infos(session, DeviceInfoOrigin.ogn_ddb) + counter = update_device_infos(session, DeviceInfoOrigin.OGN_DDB) finish_message = "DeviceInfo: {} inserted.".format(counter) logger.info(finish_message) diff --git a/app/commands/database.py b/app/commands/database.py index 0410f4e..8025b5d 100644 --- a/app/commands/database.py +++ b/app/commands/database.py @@ -94,7 +94,7 @@ def import_ddb(): """Import registered devices from the DDB.""" print("Import registered devices fom the DDB...") - counter = update_device_infos(db.session, DeviceInfoOrigin.ogn_ddb) + counter = update_device_infos(db.session, DeviceInfoOrigin.OGN_DDB) print("Imported %i devices." % counter) @@ -104,7 +104,7 @@ def import_file(path="tests/custom_ddb.txt"): """Import registered devices from a local file.""" print("Import registered devices from '{}'...".format(path)) - counter = update_device_infos(db.session, DeviceInfoOrigin.user_defined, path=path) + counter = update_device_infos(db.session, DeviceInfoOrigin.USER_DEFINED, path=path) print("Imported %i devices." % counter) @@ -114,7 +114,7 @@ def import_flarmnet(path=None): """Import registered devices from a local file.""" print("Import registered devices from '{}'...".format("internet" if path is None else path)) - counter = update_device_infos(db.session, DeviceInfoOrigin.flarmnet, path=path) + counter = update_device_infos(db.session, DeviceInfoOrigin.FLARMNET, path=path) print("Imported %i devices." % counter) diff --git a/app/model/device_info.py b/app/model/device_info.py index e0d134a..d721a20 100644 --- a/app/model/device_info.py +++ b/app/model/device_info.py @@ -1,4 +1,5 @@ from app import db +from .device_info_origin import DeviceInfoOrigin class DeviceInfo(db.Model): @@ -15,7 +16,7 @@ class DeviceInfo(db.Model): identified = db.Column(db.Boolean) aircraft_type = db.Column(db.SmallInteger) - address_origin = db.Column(db.SmallInteger) + address_origin = db.Column(db.Enum(DeviceInfoOrigin), nullable=False, default=DeviceInfoOrigin.UNKNOWN) def __repr__(self): return "" % ( diff --git a/app/model/device_info_origin.py b/app/model/device_info_origin.py index db31572..b660bb2 100644 --- a/app/model/device_info_origin.py +++ b/app/model/device_info_origin.py @@ -1,22 +1,7 @@ -class DeviceInfoOrigin: - unknown = 0 - ogn_ddb = 1 - flarmnet = 2 - user_defined = 3 +import enum - def __init__(self, origin=0): - if origin in [0, 1, 2, 3]: - self.origin = origin - else: - raise ValueError("no address origin with id {} known".format(origin)) - - def name(self): - if self.origin == self.unknown: - return "unknown" - elif self.origin == self.ogn_ddb: - return "OGN-DDB" - elif self.origin == self.flarmnet: - return "FlarmNet" - elif self.origin == self.user_defined: - return "user-defined" - return "" +class DeviceInfoOrigin(enum.Enum): + UNKNOWN = 0 + OGN_DDB = 1 + FLARMNET = 2 + USER_DEFINED = 3 \ No newline at end of file diff --git a/app/utils.py b/app/utils.py index a96d33d..69bc235 100644 --- a/app/utils.py +++ b/app/utils.py @@ -31,7 +31,7 @@ def date_to_timestamps(date): return (start, end) -def get_ddb(csv_file=None, address_origin=DeviceInfoOrigin.unknown): +def get_ddb(csv_file=None, address_origin=DeviceInfoOrigin.UNKNOWN): if csv_file is None: r = requests.get(DDB_URL) rows = "\n".join(i for i in r.text.splitlines() if i[0] != "#") @@ -59,7 +59,7 @@ def get_ddb(csv_file=None, address_origin=DeviceInfoOrigin.unknown): return device_infos -def get_flarmnet(fln_file=None, address_origin=DeviceInfoOrigin.flarmnet): +def get_flarmnet(fln_file=None, address_origin=DeviceInfoOrigin.FLARMNET): if fln_file is None: r = requests.get(FLARMNET_URL) rows = [bytes.fromhex(line).decode("latin1") for line in r.text.split("\n") if len(line) == 172] diff --git a/migrations/versions/be9a6dad551e_use_enums_for_deviceinfoorigin.py b/migrations/versions/be9a6dad551e_use_enums_for_deviceinfoorigin.py new file mode 100644 index 0000000..526b408 --- /dev/null +++ b/migrations/versions/be9a6dad551e_use_enums_for_deviceinfoorigin.py @@ -0,0 +1,41 @@ +"""Use Enums for DeviceInfoOrigin + +Revision ID: be9a6dad551e +Revises: 885123e6a2d6 +Create Date: 2019-09-15 14:38:25.838089 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + + +# revision identifiers, used by Alembic. +revision = 'be9a6dad551e' +down_revision = '885123e6a2d6' +branch_labels = None +depends_on = None + +deviceinfoorigin = postgresql.ENUM('UNKNOWN', 'OGN_DDB', 'FLARMNET', 'USER_DEFINED', name='deviceinfoorigin') + +def upgrade(): + deviceinfoorigin.create(op.get_bind()) + + op.add_column('device_infos', sa.Column('address_origin_enum', sa.Enum('UNKNOWN', 'OGN_DDB', 'FLARMNET', 'USER_DEFINED', name='deviceinfoorigin'), nullable=False, server_default='UNKNOWN')) + op.execute("UPDATE device_infos SET address_origin_enum = 'UNKNOWN' WHERE address_origin = 0") + op.execute("UPDATE device_infos SET address_origin_enum = 'OGN_DDB' WHERE address_origin = 1") + op.execute("UPDATE device_infos SET address_origin_enum = 'FLARMNET' WHERE address_origin = 2") + op.execute("UPDATE device_infos SET address_origin_enum = 'USER_DEFINED' WHERE address_origin = 3") + op.drop_column('device_infos', 'address_origin') + op.alter_column('device_infos', 'address_origin_enum', new_column_name='address_origin') + +def downgrade(): + op.add_column('device_infos', sa.Column('address_origin_int', sa.SmallInteger)) + op.execute("UPDATE device_infos SET address_origin_int = 0 WHERE address_origin = 'UNKNOWN'") + op.execute("UPDATE device_infos SET address_origin_int = 1 WHERE address_origin = 'OGN_DDB'") + op.execute("UPDATE device_infos SET address_origin_int = 2 WHERE address_origin = 'FLARMNET'") + op.execute("UPDATE device_infos SET address_origin_int = 3 WHERE address_origin = 'USER_DEFINED'") + op.drop_column('device_infos', 'address_origin') + op.alter_column('device_infos', 'address_origin_int', new_column_name='address_origin') + + deviceinfoorigin.drop(op.get_bind())