pyinfra/mini_queue/utils/rabbitmq.py

import json
import logging
import tempfile
import time
from operator import itemgetter

import pika

from mini_queue.utils.config import CONFIG
from mini_queue.utils.file import dossier_id_and_file_id_to_compressed_storage_pdf_object_name, download, unzip
from mini_queue.utils.minio import MinioHandle


def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChannel:
    channel = connection.channel()
    channel.basic_qos(prefetch_count=CONFIG.rabbitmq.prefetch_count)
    return channel


def declare_queue(channel, queue: str):
    args = {
        # "x-message-ttl": CONFIG.rabbitmq.message_ttl,
        "x-dead-letter-exchange": "",
        "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter,
    }
    return channel.queue_declare(queue=queue, auto_delete=False, arguments=args, durable=True)


def make_callback():

    storage_client = MinioHandle()

    def callback(channel, method, properties, body):
        logging.info(f" [R] Received {body}")
        response = json.dumps(process(body))
        channel.basic_publish(exchange="", routing_key=CONFIG.rabbitmq.queues.output, body=response)
        channel.basic_ack(delivery_tag=method.delivery_tag)

    def process(payload):
        payload = json.loads(payload)
        with tempfile.TemporaryDirectory() as pdf_compressed_dir:
            with tempfile.TemporaryDirectory() as pdf_dir:

                dossier_id, file_id = itemgetter("dossierId", "fileId")(payload)
                time.sleep(0.4)
                object_name = dossier_id_and_file_id_to_compressed_storage_pdf_object_name(dossier_id, file_id)
                time.sleep(0.6)
                downloaded_file_path = download(storage_client, object_name, pdf_compressed_dir)
                time.sleep(0.3)
                unzipped_file_path = unzip(downloaded_file_path, pdf_dir)
                time.sleep(1)
                payload["imageMetadata"] = []
                return json.dumps(payload)

    return callback


def read_connection_params():
    credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
    parameters = pika.ConnectionParameters(
        host=CONFIG.rabbitmq.host,
        port=CONFIG.rabbitmq.port,
        heartbeat=CONFIG.rabbitmq.heartbeat,
        credentials=credentials,
    )
    return parameters