90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
import logging
|
|
from multiprocessing import Process
|
|
|
|
import pika
|
|
|
|
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, make_storage_data_loader, make_analyzer
|
|
from pyinfra.exceptions import UnknownStorageBackend
|
|
from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
|
|
from pyinfra.storage.azure_blob_storage import AzureBlobStorageHandle
|
|
from pyinfra.storage.minio import MinioHandle
|
|
|
|
|
|
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():
|
|
|
|
load_data = make_storage_data_loader(get_storage())
|
|
analyze_file = make_analyzer(CONFIG.service.analysis_endpoint)
|
|
|
|
json_wrapped_body_processor = make_payload_processor(load_data, analyze_file)
|
|
|
|
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():
|
|
# TODO: implement meaningful checks
|
|
webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
|
|
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()
|