from orionclient.session import APISession
from orionclient.types import MolsearchDatabase
import time


def monitor_db_until_desired_state(db, desired_state):
    if desired_state not in ["LOADED", "UNLOADED"]:
        raise ValueError("Desired state must be either LOADED or UNLOADED")

    state = ''

    # Wait for completion
    while db.state != desired_state:
        APISession.refresh_resource(db)
        if state != db.state:
            print(f"DB state: {db.state}")
        time.sleep(1)

        state = db.state

    if state == '':
        print(f"DB state: {db.state}")


if __name__ == '__main__':
    db_id = < your_db_id >
    db = APISession.get_resource(MolsearchDatabase, db_id)

    db.unload()
    monitor_db_until_desired_state(db, 'UNLOADED')

    db.load()
    monitor_db_until_desired_state(db, 'LOADED')