108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
import logging
|
|
|
|
from flask import Flask, jsonify, request
|
|
from funcy import compose, flatten, rcompose
|
|
|
|
from pyinfra.server.dispatcher.dispatcher import Nothing
|
|
from pyinfra.server.processor.processor import OnDemandProcessor
|
|
from pyinfra.server.utils import unpack_batchop_pack, unpack, normalize, pack
|
|
from pyinfra.utils.buffer import bufferize
|
|
from pyinfra.utils.func import starlift, star, lift
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
class ServerPipeline:
|
|
def __init__(self, fn):
|
|
"""Function `fn` has to return an iterable and ideally is a generator."""
|
|
self.pipe = rcompose(
|
|
lift(unpack),
|
|
fn,
|
|
normalize,
|
|
starlift(pack),
|
|
)
|
|
|
|
def __call__(self, packages):
|
|
return self.pipe(packages)
|
|
|
|
|
|
def make_streamable(operation, buffer_size, batched):
|
|
|
|
operation = operation if batched else starlift(operation)
|
|
operation = ServerPipeline(operation)
|
|
operation = BufferedProcessor(operation, buffer_size=buffer_size)
|
|
operation = compose(flatten, operation)
|
|
|
|
return operation
|
|
|
|
|
|
class BufferedProcessor:
|
|
def __init__(self, fn, buffer_size):
|
|
self.fn = bufferize(fn, buffer_size=buffer_size)
|
|
|
|
def __call__(self, item, final):
|
|
return self.fn(item, final=final)
|
|
|
|
|
|
class RestOndDemandProcessor(OnDemandProcessor):
|
|
def __init__(self, fn):
|
|
super(RestOndDemandProcessor, self).__init__(fn=fn)
|
|
|
|
def submit(self, request, **kwargs) -> None:
|
|
super(RestOndDemandProcessor, self).submit(request.json, final=request.method == "POST")
|
|
|
|
|
|
class RestStreamProcessor(RestOndDemandProcessor):
|
|
"""Wraps an on-demand-processor. Combine with a webserver that provides the endpoints 'submit' and 'pickup'."""
|
|
|
|
def __init__(self, fn, submit_suffix="submit", pickup_suffix="pickup"):
|
|
super(RestStreamProcessor, self).__init__(fn=fn)
|
|
self.submit_suffix = submit_suffix
|
|
self.pickup_suffix = pickup_suffix
|
|
|
|
def submit(self, request, **kwargs):
|
|
super(RestStreamProcessor, self).submit(request)
|
|
return jsonify(request.base_url.replace(self.submit_suffix, self.pickup_suffix))
|
|
|
|
def pickup(self):
|
|
result = self.compute_next()
|
|
|
|
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(process_fn):
|
|
app = Flask(__name__)
|
|
stream = RestStreamProcessor(process_fn, 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 stream.submit(request)
|
|
|
|
@app.route("/pickup", methods=["GET"])
|
|
def pickup():
|
|
return stream.pickup()
|
|
|
|
return app
|