pyinfra/pyinfra/server/server.py

94 lines
2.5 KiB
Python

import logging
from flask import Flask, jsonify, request
from funcy import compose, identity
from pyinfra.server.bufferizer.lazy_bufferizer import Queue, StreamBuffer, stream_queue, make_consumer
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()
class ServerPipeline:
def __init__(self):
"""
- dequeue
-
"""
pass
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, batch_size):
fn = make_streamable(fn, batched)
fn = unpack_fn_pack(fn)
return fn
class RestStreamProcessor:
"""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"):
self.submit_suffix = submit_suffix
self.pickup_suffix = pickup_suffix
self.queue = Queue()
# self.processor = QueueBufferCoupler(self.queue, StreamBuffer(fn))
# self.output_stream = chain.from_iterable(map(StreamBuffer(fn), stream_queue(self.queue)))
self.fn = make_consumer(StreamBuffer(fn))
def submit(self, request):
self.queue.append(request.json)
return jsonify(request.base_url.replace(self.submit_suffix, self.pickup_suffix))
def pickup(self):
result = self.fn(stream_queue(self.queue))
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