ogn-python/ogn/commands/database.py

from ogn.commands.dbutils import engine, session
from ogn.model import Base, AddressOrigin, AircraftBeacon, ReceiverBeacon, Device, Receiver
from ogn.utils import get_airports, open_file
from ogn.collect.database import update_device_infos

from sqlalchemy import insert, distinct
from sqlalchemy.sql import null

from manager import Manager
manager = Manager()

ALEMBIC_CONFIG_FILE = "alembic.ini"


@manager.command
def init():
    """Initialize the database."""

    from alembic.config import Config
    from alembic import command

    session.execute('CREATE EXTENSION IF NOT EXISTS postgis;')
    session.commit()
    Base.metadata.create_all(engine)
    alembic_cfg = Config(ALEMBIC_CONFIG_FILE)
    command.stamp(alembic_cfg, "head")
    print("Done.")


@manager.command
def upgrade():
    """Upgrade database to the latest version."""

    from alembic.config import Config
    from alembic import command

    alembic_cfg = Config(ALEMBIC_CONFIG_FILE)
    command.upgrade(alembic_cfg, 'head')


@manager.command
def drop(sure='n'):
    """Drop all tables."""
    if sure == 'y':
        Base.metadata.drop_all(engine)
        print('Dropped all tables.')
    else:
        print("Add argument '--sure y' to drop all tables.")


@manager.command
def import_ddb():
    """Import registered devices from the DDB."""

    print("Import registered devices fom the DDB...")
    address_origin = AddressOrigin.ogn_ddb
    counter = update_device_infos(session,
                                  address_origin)
    print("Imported %i devices." % counter)


@manager.command
def import_file(path='tests/custom_ddb.txt'):
    """Import registered devices from a local file."""
    # (flushes previously manually imported entries)

    print("Import registered devices from '{}'...".format(path))
    address_origin = AddressOrigin.user_defined
    counter = update_device_infos(session,
                                  address_origin,
                                  csvfile=path)
    print("Imported %i devices." % counter)


@manager.command
def import_airports(path='tests/SeeYou.cup'):
    """Import airports from a ".cup" file"""

    print("Import airports from '{}'...".format(path))
    airports = get_airports(path)
    session.bulk_save_objects(airports)
    session.commit()
    print("Imported {} airports.".format(len(airports)))


@manager.command
def update_relations():
    """Update AircraftBeacon and ReceiverBeacon relations"""

    # Create missing Receiver from ReceiverBeacon
    available_receivers = session.query(Receiver.name) \
        .subquery()

    missing_receiver_query = session.query(distinct(ReceiverBeacon.name)) \
        .filter(ReceiverBeacon.receiver_id == null()) \
        .filter(~ReceiverBeacon.name.in_(available_receivers))

    ins = insert(Receiver).from_select([Receiver.name], missing_receiver_query)
    session.execute(ins)

    # Create missing Device from AircraftBeacon
    available_addresses = session.query(Device.address) \
        .subquery()

    missing_addresses_query = session.query(distinct(AircraftBeacon.address)) \
        .filter(AircraftBeacon.device_id == null()) \
        .filter(~AircraftBeacon.address.in_(available_addresses))

    ins2 = insert(Device).from_select([Device.address], missing_addresses_query)
    session.execute(ins2)
    session.commit()
    print("Inserted {} Receivers and {} Devices".format(ins, ins2))
    return

    # Update AircraftBeacons
    upd = session.query(AircraftBeacon) \
        .filter(AircraftBeacon.device_id == null()) \
        .filter(AircraftBeacon.receiver_id == null()) \
        .filter(AircraftBeacon.address == Device.address) \
        .filter(AircraftBeacon.receiver_name == Receiver.name) \
        .update({AircraftBeacon.device_id: Device.id,
                 AircraftBeacon.receiver_id: Receiver.id},
                synchronize_session='fetch')

    upd2 = session.query(ReceiverBeacon) \
        .filter(ReceiverBeacon.receiver_id == null()) \
        .filter(ReceiverBeacon.receiver_name == Receiver.name) \
        .update({Receiver.name: ReceiverBeacon.receiver_name},
                synchronize_session='fetch')

    session.commit()
    print("Updated {} AircraftBeacons and {} ReceiverBeacons".
          format(upd, upd2))