111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
import logging
|
|
from multiprocessing import Process
|
|
|
|
import pika
|
|
from flask import Flask, jsonify
|
|
from waitress import serve
|
|
|
|
from pyinfra.callback import (
|
|
make_retry_callback_for_output_queue,
|
|
make_retry_callback,
|
|
make_callback_for_output_queue,
|
|
)
|
|
from pyinfra.config import CONFIG
|
|
from pyinfra.consume import consume, ConsumerError
|
|
from pyinfra.core import make_payload_processor
|
|
from pyinfra.exceptions import UnknownStorageBackend
|
|
from pyinfra.storage.azure_blob_storage import AzureBlobStorageHandle
|
|
from pyinfra.storage.minio import MinioHandle
|
|
|
|
|
|
# TODO: implement meaningful checks
|
|
def start_integrity_checks_webserver(mode="debug"):
|
|
|
|
app = Flask(__name__)
|
|
|
|
@app.route("/ready", methods=["GET"])
|
|
def ready():
|
|
resp = jsonify("OK")
|
|
resp.status_code = 200
|
|
return resp
|
|
|
|
@app.route("/health", methods=["GET"])
|
|
def healthy():
|
|
resp = jsonify("OK")
|
|
resp.status_code = 200
|
|
return resp
|
|
|
|
if mode == "debug":
|
|
app.run(host=CONFIG.webserver.host, port=CONFIG.webserver.port, debug=True)
|
|
elif mode == "production":
|
|
serve(app, host=CONFIG.webserver.host, port=CONFIG.webserver.port)
|
|
|
|
|
|
def get_storage():
|
|
|
|
storage_backend = CONFIG.service.storage_backend
|
|
if storage_backend == "s3":
|
|
storage = MinioHandle()
|
|
elif storage_backend == "azure":
|
|
storage = AzureBlobStorageHandle()
|
|
else:
|
|
raise UnknownStorageBackend(f"Unknown storage backend '{storage_backend}'.")
|
|
|
|
return storage
|
|
|
|
|
|
def republish(channel, body, n_current_attempts):
|
|
channel.basic_publish(
|
|
exchange="",
|
|
routing_key=CONFIG.rabbitmq.queues.input,
|
|
body=body,
|
|
properties=pika.BasicProperties(headers={"x-retry-count": n_current_attempts}),
|
|
)
|
|
|
|
|
|
def make_callback():
|
|
|
|
json_wrapped_body_processor = make_payload_processor(CONFIG.service.analysis_endpoint)
|
|
|
|
if CONFIG.rabbitmq.retry.enabled:
|
|
retry_callback = make_retry_callback(republish, max_attempts=CONFIG.rabbitmq.retry.max_attempts)
|
|
|
|
callback = make_retry_callback_for_output_queue(
|
|
json_wrapped_body_processor=json_wrapped_body_processor,
|
|
output_queue_name=CONFIG.rabbitmq.queues.output,
|
|
retry_callback=retry_callback,
|
|
)
|
|
else:
|
|
callback = make_callback_for_output_queue(
|
|
json_wrapped_body_processor=json_wrapped_body_processor, output_queue_name=CONFIG.rabbitmq.queues.output
|
|
)
|
|
|
|
return callback
|
|
|
|
|
|
def main():
|
|
|
|
webserver = Process(target=start_integrity_checks_webserver, args=("production",))
|
|
logging.info("Starting webserver...")
|
|
webserver.start()
|
|
|
|
try:
|
|
consume(CONFIG.rabbitmq.queues.input, make_callback())
|
|
except KeyboardInterrupt:
|
|
pass
|
|
except ConsumerError:
|
|
webserver.terminate()
|
|
raise
|
|
|
|
webserver.join()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging_level = CONFIG.service.logging_level
|
|
logging.basicConfig(level=logging_level)
|
|
logging.getLogger("pika").setLevel(logging.ERROR)
|
|
logging.getLogger("flask").setLevel(logging.ERROR)
|
|
logging.getLogger("urllib3").setLevel(logging.ERROR)
|
|
|
|
main()
|