93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
import logging
|
|
from asyncio import AbstractEventLoop
|
|
|
|
import aio_pika
|
|
import asyncio
|
|
import os
|
|
|
|
from pathlib import Path
|
|
from aio_pika.abc import AbstractRobustChannel, AbstractRobustExchange, AbstractExchange
|
|
from aio_pika import Message, DeliveryMode
|
|
|
|
from amqp.adapter.file_downloader import IN_PROGRESS, END_OF_FILE
|
|
from amqp.adapter.logging_utils import logging_debug
|
|
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
|
|
|
|
|
|
async def publish_file_to_rabbitmq(
|
|
exchange: AbstractExchange,
|
|
file_path: Path,
|
|
routing_key: str,
|
|
loop: AbstractEventLoop,
|
|
max_chunk_size: int = 2**20, # 1MB default chunk size
|
|
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:
|
|
exchange: 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).
|
|
routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue.
|
|
loop: correct event loop
|
|
max_chunk_size: chunk size to send
|
|
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)
|
|
|
|
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,
|
|
"pos": _offset,
|
|
"total_size": _file_size,
|
|
"chunk_size": len(_chunk),
|
|
"eof": (
|
|
END_OF_FILE
|
|
if max_chunk_size + _offset >= _file_size
|
|
else IN_PROGRESS
|
|
),
|
|
},
|
|
)
|
|
logging.debug(
|
|
f"Going to Publish chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
|
|
)
|
|
_future = asyncio.run_coroutine_threadsafe(
|
|
exchange.publish(message, routing_key=routing_key), loop
|
|
)
|
|
while not _future.done():
|
|
# sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
|
|
await asyncio.sleep(0.02)
|
|
logging.debug(
|
|
f"Waiting for chunk {_offset}/{_file_size} future= {_future}"
|
|
)
|
|
|
|
_offset += len(_chunk)
|
|
logging.debug(
|
|
f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
|
|
)
|
|
|
|
logging_debug(
|
|
f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}"
|
|
)
|
|
return routing_key, _file_size
|