Stanislav Hejny
408 lines
16 KiB
Python
408 lines
16 KiB
Python
"""
|
|
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
|
|
|
|
|
|
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.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
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
|
|
|
|
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.")
|
|
|
|
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.")
|
|
|
|
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.")
|
|
|
|
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
|
|
class StreamingFileHandler(FileHandler):
|
|
"""
|
|
Handles file downloads from RabbitMQ streams.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
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}"
|
|
)
|
|
|
|
async def wrap_rabbit_mq_calls() -> tuple[int, int] | None:
|
|
connection = await connect_robust(rabbitmq_url, loop=loop)
|
|
try:
|
|
_file_size = 0
|
|
_eof = IN_PROGRESS
|
|
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 {filename}"
|
|
)
|
|
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
|
|
|
|
if queue_name:
|
|
_local_file_size, _local_eof = await await_result(
|
|
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
|
|
)
|
|
return _local_file_size, _local_eof
|
|
|
|
return 0, CANCELLED
|
|
|
|
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (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:
|
|
|
|
async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]:
|
|
_connection = await connect_robust(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
|
|
|
|
_local_res, _local_eof = await await_result(
|
|
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
|
|
)
|
|
return _local_res, _local_eof
|
|
return 0, CANCELLED
|
|
|
|
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
|
|
|
|
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
|