"""
implements the file downloader for the amqp adapter, where file is sent in chunks over RabbitMQ
"""
import asyncio
import json
import logging
import os
from asyncio import AbstractEventLoop
from collections import defaultdict
from pathlib import Path
from aio_pika import Message, connect_robust
from aio_pika.abc import (
AbstractExchange,
AbstractIncomingMessage,
AbstractQueue,
AbstractRobustChannel,
DeliveryMode,
)
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.router.utils import await_result
# Global dictionary to track completed downloads
download_locations = defaultdict(str)
IN_PROGRESS = 0
END_OF_FILE = 1
CANCELLED = 2
[docs]
class FileHandler:
"""
The abstract base class for file handling. This is mainly to make the code testable, as there are only 2
options for pushing (uploading) a file(s) to a service: as HTTP multipart message delivered via RabbitMQ
or via stream in RabbitMQ. Downloading a file is always done via RabbitMQ stream. Thus the method signiture
is tighly coupled to the RabbitMQ API in v1. Only alternative implementation is a mock object in unittests.
"""
[docs]
def __init__(self):
pass
[docs]
def ensure_directory_exists(self, filepath) -> str:
"""
Ensures that the directory for the given filepath exists, and handles file versioning.
Args:
filepath: The path to the file.
Returns:
The original filepath if the file does not exist, or a modified filepath
pointing to the next available version of the file if it does.
"""
# 1. Ensure directory exists
directory = os.path.dirname(filepath)
if directory and not os.path.exists(directory):
os.makedirs(directory)
# 2. Handle file versioning
if not os.path.exists(filepath):
return filepath # File doesn't exist, return original path
# File exists, find next available version
base, ext = os.path.splitext(filepath)
version = 1
while True:
new_filepath = f"{base}.{version}{ext}"
if not os.path.exists(new_filepath):
return new_filepath # Found a free version, return this path
version += 1
[docs]
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int):
"""
Asynchronously downloads a bytearray from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
queue_name (str): The directory where the file should be saved.
"""
raise NotImplementedError("Subclasses must implement this method.")
[docs]
async def download_file(
self, loop, rabbitmq_url, file_info, message_data, output_dir
) -> (int, int):
"""
Asynchronously downloads a file from RabbitMQ based on the provided file information.
takes the file_info and message_data and output_dir, creates queue consumer and waits for the file
to be downloaded. File may consist of multiple chunks, so the function will wait until all chunks are received.
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
file_info (dict): A dictionary containing file metadata
message_data (dict): The entire message
output_dir (str): The directory where the file should be saved.
"""
raise NotImplementedError("Subclasses must implement this method.")
[docs]
async def on_message(
self,
message: AbstractIncomingMessage,
json_data: bytes,
rabbitmq_url: str,
loop,
output_dir="downloaded_files",
) -> defaultdict:
"""
Asynchronous callback function for handling messages from the 'amq-cleverswarm' queue.
Args:
message: The incoming RabbitMQ message.
rabbitmq_url: The RabbitMQ connection URL.
"""
raise NotImplementedError("Subclasses must implement this method.")
[docs]
async def publish_file(
self,
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.
"""
raise NotImplementedError("Subclasses must implement this method.")
# The StreamingFileHandler class is a subclass of FileHandler that implements the download_file method
[docs]
class StreamingFileHandler(FileHandler):
"""
Handles file downloads from RabbitMQ streams.
"""
[docs]
def __init__(self):
super().__init__()
# This method MUST NOT be static.
# For testing purpose, we will mock this method and inject a mock connection.
async def _get_connection(self, rabbitmq_url, loop):
return await connect_robust(rabbitmq_url, loop=loop)
async def _download_from_queue_to_file(
self, rabbitmq_url, queue_name, filepath, loop
) -> tuple[int, int] | None:
"""
Internal method for downloading a file from a queue to disk.
Return (file size, eof status).
"""
# Here the normal mock doesn't work since this method
# is async and will be run in another loop,
# out of the scope of method mock.
connection = await self._get_connection(rabbitmq_url, loop=loop)
_file_size = 0
_eof = IN_PROGRESS
try:
channel: AbstractRobustChannel = await connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all file chunks being downloaded for {filepath}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
with open(filepath, "ab") as f:
f.write(message.body)
_file_size = _file_size + message.body_size
_eof = int(message.headers.get("eof", END_OF_FILE))
if _eof != IN_PROGRESS:
break
except Exception as e:
logging_error(f"File download failed: {e}")
finally:
await connection.close()
return _file_size, _eof
[docs]
async def download_file(
self, loop, rabbitmq_url, file_info, message_data, output_dir
) -> (int, int):
"""
Asynchronously downloads a file from RabbitMQ based on the provided file information.
takes the file_info and message_data and output_dir, creates queue consumer and waits for the file to be downloaded
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
file_info (dict): A dictionary containing file metadata
message_data (dict): The entire message
output_dir (str): The directory where the file should be saved.
"""
filename = file_info["filename"]
size = file_info["size"]
stream_name = file_info["streamName"]
queue_name = file_info["queue_name"]
filepath = os.path.join(output_dir, filename)
filepath = self.ensure_directory_exists(
filepath
) # if filename exists, create a new (next) version
download_locations[file_info["key"]] = filepath
filename = os.path.basename(
filepath
) # if a new version was created, ensure to use that new version
file_info["filename"] = filename
logging_debug(
f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}"
)
if queue_name:
_local_file_size, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_file(rabbitmq_url, queue_name, filepath, loop),
loop=loop,
)
)
return _local_file_size, _local_eof
return 0, CANCELLED
async def _download_from_queue_to_buffer(
self, rabbitmq_url, queue_name, loop
) -> tuple[bytearray, int]:
"""
Internal method for downloading a file from a queue to RAM.
Return (bytearray, eof status).
"""
_connection = await self._get_connection(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
[docs]
async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int):
"""
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
queue_name (str): The directory where the file should be saved.
"""
logging_debug(f"Downloading buffer from stream: {queue_name}")
if queue_name:
_local_res, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_buffer(rabbitmq_url, queue_name, loop), loop=loop
)
)
return _local_res, _local_eof
return 0, CANCELLED
[docs]
async def on_message(
self,
message: AbstractIncomingMessage,
json_data: bytes,
rabbitmq_url: str,
loop,
output_dir="downloaded_files",
) -> defaultdict:
"""
Asynchronous callback function for handling messages from the 'amq-cleverswarm' queue.
Args:
message: The incoming RabbitMQ message.
json_data: The JSON data from the message.
rabbitmq_url: The RabbitMQ connection URL.
loop: The asyncio event loop of the RabbitMQ API
output_dir
"""
try:
_message_data = json.loads(message.body.decode())
_amq_message_data = json.loads(json_data.decode())
files_to_download = []
for _key, _file_type in _amq_message_data["files"].items():
if len(_file_type) > 0:
files_info = _file_type[0]
stream_name = files_info["streamName"]
files_info["queue_name"] = (
_message_data.get(stream_name, "") if stream_name else ""
)
files_info["key"] = _key
files_to_download.append(files_info)
if not files_to_download:
logging_info(" [%s] No files to download in this message.", message.delivery_tag)
# await message.ack()
return defaultdict()
# Create a task for each file download
_tasks = []
for file_info in files_to_download:
logging_info(f" [{message.delivery_tag}] about to create task for {file_info}")
_task = asyncio.create_task(
self.download_file(loop, rabbitmq_url, file_info, _message_data, output_dir)
)
_tasks.append(_task)
logging_info(
f" [{message.delivery_tag}] Waiting for all files to download...gather(*tasks)"
)
await asyncio.gather(*_tasks) # Wait for all downloads to complete
logging_info(f"All files downloaded. d-tag: {message.delivery_tag}")
except Exception as e:
logging_error(f"Building DataMessage: {e}")
asyncio.run_coroutine_threadsafe(
message.ack(), loop=loop
) # ACK -> ignore the message (do not redeliver)
return download_locations
[docs]
async def publish_file(
self,
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