pyinfra/scripts/mock_knock.py

import json

import pika

from mini_queue.utils.config import CONFIG
from mini_queue.utils.minio import MinioHandle
from mini_queue.utils.rabbitmq import make_channel, declare_queue


def build_message_bodies():
    minio_client = MinioHandle()
    for dossier_id, pdf_name in minio_client.list_folders_and_files():
        file_id = pdf_name.split(".")[0]
        yield json.dumps({"dossierId": dossier_id, "fileId": file_id})


if __name__ == "__main__":
    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,
    )

    connection = pika.BlockingConnection(parameters)
    channel = make_channel(connection)

    declare_queue(channel, CONFIG.rabbitmq.queues.input)
    declare_queue(channel, CONFIG.rabbitmq.queues.output)

    for body in build_message_bodies():
        channel.basic_publish("", CONFIG.rabbitmq.queues.input, body)
        print(f" [x] Put {body} on {CONFIG.rabbitmq.queues.input}")

    for method_frame, _, body in channel.consume(queue=CONFIG.rabbitmq.queues.output):
        print(json.loads(body))
        channel.basic_ack(method_frame.delivery_tag)