93 lines
2.5 KiB
Python
93 lines
2.5 KiB
Python
import logging
|
|
|
|
from flask import Flask, jsonify, request
|
|
from funcy import compose, identity, first
|
|
|
|
from pyinfra.server.bufferizer.lazy_bufferizer import Queue, stream_queue, FlatStreamBuffer
|
|
from pyinfra.server.dispatcher.dispatcher import Nothing
|
|
from pyinfra.server.utils import unpack, normalize, pack
|
|
from pyinfra.utils.func import starlift, lift
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
def make_streamable(fn, batched):
|
|
return compose(normalize, (identity if batched else starlift)(fn))
|
|
|
|
|
|
def unpack_fn_pack(fn):
|
|
return compose(starlift(pack), fn, lift(unpack))
|
|
|
|
|
|
def make_streamable_and_wrap_in_packing_logic(fn, batched):
|
|
fn = make_streamable(fn, batched)
|
|
fn = unpack_fn_pack(fn)
|
|
return fn
|
|
|
|
|
|
class LazyProcessor:
|
|
"""Accepts computation requests (push) and lazily produces results (pop)."""
|
|
def __init__(self, flat_stream_buffer: FlatStreamBuffer):
|
|
self.queue = Queue()
|
|
self.flat_stream_buffer = flat_stream_buffer
|
|
|
|
def push(self, item):
|
|
self.queue.append(item)
|
|
|
|
def pop(self):
|
|
items = stream_queue(self.queue)
|
|
return first(self.flat_stream_buffer(items))
|
|
|
|
|
|
class LazyRestProcessor:
|
|
def __init__(self, lazy_processor: LazyProcessor, submit_suffix="submit", pickup_suffix="pickup"):
|
|
self.submit_suffix = submit_suffix
|
|
self.pickup_suffix = pickup_suffix
|
|
self.lazy_processor = lazy_processor
|
|
|
|
def push(self, request):
|
|
self.lazy_processor.push(request.json)
|
|
return jsonify(request.base_url.replace(self.submit_suffix, self.pickup_suffix))
|
|
|
|
def pop(self):
|
|
result = self.lazy_processor.pop() or Nothing
|
|
|
|
if not valid(result):
|
|
logger.error(f"Received invalid result: {result}")
|
|
result = Nothing
|
|
|
|
if result is Nothing:
|
|
resp = jsonify("No more items left")
|
|
resp.status_code = 204
|
|
|
|
else:
|
|
resp = jsonify(result)
|
|
resp.status_code = 206
|
|
|
|
return resp
|
|
|
|
|
|
def valid(result):
|
|
return isinstance(result, dict) or result is Nothing
|
|
|
|
|
|
def set_up_processing_server(package_processor):
|
|
app = Flask(__name__)
|
|
processor = LazyRestProcessor(package_processor, submit_suffix="submit", pickup_suffix="pickup")
|
|
|
|
@app.route("/ready", methods=["GET"])
|
|
def ready():
|
|
resp = jsonify("OK")
|
|
resp.status_code = 200
|
|
return resp
|
|
|
|
@app.route("/submit", methods=["POST", "PATCH"])
|
|
def submit():
|
|
return processor.push(request)
|
|
|
|
@app.route("/pickup", methods=["GET"])
|
|
def pickup():
|
|
return processor.pop()
|
|
|
|
return app
|