pyinfra/pyinfra/server/server.py

from functools import singledispatch
from typing import Dict, Callable, Union

from flask import Flask, jsonify, request
from prometheus_client import generate_latest

from pyinfra.config import CONFIG
from pyinfra.server.buffering.stream import FlatStreamBuffer
from pyinfra.server.monitoring import OperationDispatcherMonitoringDecorator
from pyinfra.server.operation_dispatcher import OperationDispatcher
from pyinfra.server.stream.queued_stream_function import QueuedStreamFunction


@singledispatch
def set_up_processing_server(arg: Union[dict, Callable], buffer_size=1):
    """Produces a processing server given a streamable function or a mapping from operations to streamable functions.
    Streamable functions are constructed by calling pyinfra.server.utils.make_streamable_and_wrap_in_packing_logic on a
    function taking a tuple of data and metadata and also returning a tuple or yielding tuples of data and metadata.
    If the function doesn't produce data, data should be an empty byte string.
    If the function doesn't produce metadata, metadata should be an empty dictionary.

    Args:
        arg: streamable function or mapping of operations: str to streamable functions
        buffer_size: If your function operates on batches this parameter controls how many items are aggregated before
                     your function is applied.

    TODO: buffer_size has to be controllable on per function basis.

    Returns:
        Processing server: flask app
    """
    pass


@set_up_processing_server.register
def _(operation2stream_fn: dict, buffer_size=1):
    return __stream_fn_to_processing_server(operation2stream_fn, buffer_size)


@set_up_processing_server.register
def _(stream_fn: object, buffer_size=1):
    operation2stream_fn = {None: stream_fn}
    return __stream_fn_to_processing_server(operation2stream_fn, buffer_size)


def __stream_fn_to_processing_server(operation2stream_fn: dict, buffer_size):
    operation2stream_fn = {
        op: QueuedStreamFunction(FlatStreamBuffer(fn, buffer_size)) for op, fn in operation2stream_fn.items()
    }
    return __set_up_processing_server(operation2stream_fn)


def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunction]):
    app = Flask(__name__)

    dispatcher = OperationDispatcherMonitoringDecorator(
        OperationDispatcher(operation2function),
        naming_policy=naming_policy,
    )

    def ok():
        resp = jsonify("OK")
        resp.status_code = 200
        return resp

    @app.route("/ready", methods=["GET"])
    def ready():
        return ok()

    @app.route("/health", methods=["GET"])
    def healthy():
        return ok()

    @app.route("/prometheus", methods=["GET"])
    def prometheus():
        return generate_latest(registry=dispatcher.registry)

    @app.route("/<operation>", methods=["POST", "PATCH"])
    def submit(operation):
        return dispatcher.submit(operation, request)

    @app.route("/", methods=["POST", "PATCH"])
    def submit_default():
        return dispatcher.submit("", request)

    @app.route("/<operation>", methods=["GET"])
    def pickup(operation):
        return dispatcher.pickup(operation)

    return app


def naming_policy(op_name: str):
    pop_suffix = OperationDispatcher.pickup_suffix
    prefix = f"redactmanager_{CONFIG.service.name}"

    op_display_name = op_name.replace(f"_{pop_suffix}", "") if op_name != pop_suffix else "default"
    complete_display_name = f"{prefix}_{op_display_name}"

    return complete_display_name