76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
import logging
|
|
import aio_pika
|
|
import asyncio
|
|
import os
|
|
|
|
from pathlib import Path
|
|
from aio_pika.abc import AbstractRobustChannel
|
|
from aio_pika import Message, DeliveryMode
|
|
|
|
from amqp.adapter.file_downloader import IN_PROGRESS, END_OF_FILE
|
|
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
|
|
|
|
|
|
async def publish_file_to_rabbitmq(
|
|
channel: AbstractRobustChannel,
|
|
file_path: Path,
|
|
max_chunk_size: int = 2**20, # 1MB default chunk size
|
|
exchange_name: str = AMQ_REPLY_TO_EXCHANGE,
|
|
content_type: str = 'application/octet-stream', # Default content type
|
|
delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
|
|
) -> tuple[str, int]:
|
|
"""
|
|
Reads a file from the file system and publishes its content to a RabbitMQ exchange
|
|
in chunks. It declares a queue with a random name, binds it to the specified
|
|
exchange, and uses the queue name as the routing key.
|
|
|
|
Args:
|
|
channel: The aio_pika RobustChannel to use for communication with RabbitMQ.
|
|
file_path: The path to the file to publish.
|
|
max_chunk_size: The maximum size of each chunk (in bytes).
|
|
exchange_name: The name of the RabbitMQ exchange to publish to.
|
|
content_type: The content type of the file data.
|
|
delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT).
|
|
|
|
Returns:
|
|
The name of the randomly generated queue where the file content was published.
|
|
"""
|
|
if not file_path.is_file():
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
_file_size = os.path.getsize(file_path)
|
|
_exchange = await channel.get_exchange(name=exchange_name, ensure=True)
|
|
_queue = await channel.declare_queue(exclusive=True) # Declare a unique, exclusive queue
|
|
_queue_name = _queue.name
|
|
await _queue.bind(
|
|
exchange=exchange_name,
|
|
routing_key=_queue_name, # Use queue name as routing key
|
|
)
|
|
|
|
logging.info(f"Publishing file: {file_path} (size: {_file_size}) to exchg: {exchange_name}, q: {_queue_name}")
|
|
|
|
with open(file_path, 'rb') as f:
|
|
_offset = 0
|
|
while True:
|
|
_chunk = f.read(max_chunk_size)
|
|
if not _chunk:
|
|
break # End of file
|
|
|
|
message = Message(
|
|
body=_chunk,
|
|
content_type=content_type,
|
|
delivery_mode=delivery_mode,
|
|
headers={
|
|
'file_name': file_path.name,
|
|
'offset': _offset,
|
|
'total_size': _file_size,
|
|
'chunk_size': len(_chunk),
|
|
'eof': END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS
|
|
}
|
|
)
|
|
await _exchange.publish(message, routing_key=_queue_name)
|
|
_offset += len(_chunk)
|
|
logging.debug(f"Published chunk {_offset}/{_file_size} to {exchange_name}:{_queue_name}")
|
|
|
|
print(f"File {file_path} published to RabbitMQ exchange {exchange_name}, queue {_queue_name}")
|
|
return _queue_name, _file_size
|