BlueskyDatabrokerPackServer

BlueskyDatabrokerPackServer.py

"""
This builds and serves databroker-pack files
"""
import asyncio
import os
import pathlib
import subprocess
import tempfile
import zipfile

from fastapi import FastAPI, status
from fastapi.responses import FileResponse
import databroker.tutorial_utils

HOST = '0.0.0.0'
PORT = 9090

catalog_name = os.environ.get("PACK_CATALOG",
            'bluesky-tutorial-BMM')
pack_directory = os.environ.get("PACK_DIRECTORY",
            os.path.join(pathlib.Path.home(), "packs"))

available_packs = set()
currently_packing = dict()

app = FastAPI()

def zipdir(path, ziph):
    # https://stackoverflow.com/a/1855118/6513183
    # ziph is zipfile handle
    for root, dirs, files in os.walk(path):
        for file in files:
            ziph.write(os.path.join(root, file),
                       os.path.relpath(os.path.join(root, file),
                                       os.path.join(path, '..')))

def initialize(pack_directory):
    global available_packs

    os.makedirs(pack_directory, exist_ok=True)
    available_packs = {filename.split('.')[0] for
                       filename in os.listdir(pack_directory)}
    databroker.tutorial_utils.fetch_BMM_example()


def authorized_run(token, run_uid):
    return True


def authorized_file(token, file_name):
    return True


def eta(process):
    return '?'


async def pack(run_uid):
    global available_packs
    global currently_packing

    if run_uid in currently_packing.keys():
        return "Already Packing"

    with open(pack_directory + '/' + run_uid + ".txt", "w+") as f:
        f.write(run_uid)

    process = asyncio.create_subprocess_exec(
          'databroker-pack',
          *f'{catalog_name} --uids {pack_directory + "/" + run_uid + ".txt"} {pack_directory + "/" + run_uid}'.split())

    currently_packing[run_uid] = lambda: '?' #eta(process)

    asyncio.create_task(finish_packing(process, run_uid))

    return process


async def finish_packing(process, run_uid):
    global currently_packing
    global available_packs

    print("Packing Started!", run_uid)
    try:
        await process
        zipf = zipfile.ZipFile(os.path.join(pack_directory, run_uid + '.tar.gz'),
                               'w', zipfile.ZIP_DEFLATED)
        zipdir(os.path.join(pack_directory, run_uid), zipf)
        zipf.close()
    except Exception as e:
        print("Packing failed :(", run_uid)
        currently_packing.pop(run_uid)
        raise e

    currently_packing.pop(run_uid, None)
    available_packs.add(run_uid)
    print("Packing Complete!", run_uid)


@app.get("/pack/{run_uid}")
async def get_pack(run_uid: str, token=None):
    global currently_packing

    pack_path = os.path.join(pack_directory, run_uid + '.tar.gz')

    if not authorized_run(token, run_uid):
        return "Authentication Failed"

    if run_uid in available_packs:
        return FileResponse(pack_path)

    await pack(run_uid)

    return {"status": "queued", "eta": currently_packing[run_uid]()}


@app.get("/file/")
async def get_file(filename: str, token=None):

    if not os.path.exists(filename):
        return "File not found"

    if not authorized_file(token, filename):
        return "Authentication Failed"

    return FileResponse(filename)


@app.on_event("startup")
def startup():
    print('Starting BlueskyPackServer')
    initialize(pack_directory)


@app.on_event("shutdown")
def shutdown():
    print('Stopping BlueskyPackServer')


if __name__ == "__main__":
    import uvicorn
    uvicorn.run('BlueskyDatabrokerPackServer:app', host=HOST, port=PORT,
                log_level='info', reload=True)