233 lines
8.5 KiB
Python
233 lines
8.5 KiB
Python
"""
|
|
implements the file downloader for the amqp adapter, where file is sent in chunks over RabbitMQ
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import logging
|
|
from asyncio import Future
|
|
from collections import defaultdict
|
|
from typing import Any
|
|
|
|
from aio_pika import connect_robust
|
|
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel, AbstractQueue
|
|
|
|
from amqp.adapter.logging_utils import logging_error
|
|
|
|
# Global dictionary to track completed downloads
|
|
download_locations = defaultdict(str)
|
|
IN_PROGRESS = 0
|
|
END_OF_FILE = 1
|
|
CANCELLED = 2
|
|
|
|
|
|
def ensure_directory_exists(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_file(
|
|
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.
|
|
"""
|
|
global download_locations
|
|
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 = 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.info(
|
|
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.info(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:
|
|
_future: Future = asyncio.run_coroutine_threadsafe(
|
|
wrap_rabbit_mq_calls(), loop=loop
|
|
)
|
|
while not _future.done():
|
|
await asyncio.sleep(0.05)
|
|
_local_file_size, _local_eof = _future.result()
|
|
return _local_file_size, _local_eof
|
|
|
|
return 0, CANCELLED
|
|
|
|
|
|
async def download_buffer(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.info(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.info(
|
|
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
|
|
|
|
_future: Future = asyncio.run_coroutine_threadsafe(
|
|
wrap_rabbit_mq_calls(), loop=loop
|
|
)
|
|
while not _future.done():
|
|
await asyncio.sleep(0.05)
|
|
_local_res, _local_eof = _future.result()
|
|
return _local_res, _local_eof
|
|
return 0, CANCELLED
|
|
|
|
|
|
async def on_message(
|
|
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.
|
|
"""
|
|
global download_locations
|
|
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 False
|
|
|
|
# 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(
|
|
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
|
|
|
|
# After all files are downloaded, acknowledge the original message
|
|
logging.info(f"All files downloaded. d-tag: {message.delivery_tag}")
|
|
# logging.info(f"All files downloaded. Acknowledging the main message d-tag: {message.delivery_tag}")
|
|
# await message.ack()
|
|
# logging.info(f"Message d-tag: {message.delivery_tag} acknowledged.")
|
|
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
|