Stanislav Hejny
141 lines
5.6 KiB
Python
141 lines
5.6 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import threading
|
|
from collections import defaultdict
|
|
|
|
from aio_pika.abc import HeadersType, AbstractIncomingMessage, AbstractQueue, AbstractRobustChannel
|
|
from aiormq import Channel
|
|
|
|
|
|
class MultipartDownload:
|
|
|
|
def __init__(self, headers: HeadersType, channel: AbstractRobustChannel):
|
|
self.headers = headers
|
|
self.channel = channel
|
|
self.addressing: int = headers.get('clevermicro.addressing', 0)
|
|
self.files_to_download = []
|
|
self.threads = []
|
|
# Global dictionary to track completed downloads
|
|
self.download_status = defaultdict(lambda: defaultdict(bool))
|
|
self.future = None
|
|
|
|
def ensure_directory_exists(self, filepath):
|
|
"""Ensures that the directory for the given filepath exists."""
|
|
dir_name = os.path.dirname(filepath)
|
|
if dir_name and not os.path.exists(dir_name):
|
|
os.makedirs(dir_name)
|
|
|
|
async def start_queue2(self, channel: Channel, queue_name: str, callback):
|
|
dok = await channel.queue_declare(queue=queue_name, durable=False, auto_delete=True)
|
|
await channel.basic_consume(queue=queue_name, consumer_callback=callback, no_ack=True, exclusive=True)
|
|
|
|
async def start_queue(self, queue_name: str, callback):
|
|
"""
|
|
Starts consuming from a queue.
|
|
|
|
Args:
|
|
:param queue_name: The name of the queue to consume from.
|
|
:param callback: The callback function to call when a message is received.
|
|
"""
|
|
queue: AbstractQueue = await self.channel.declare_queue(name=queue_name, durable=False, auto_delete=True)
|
|
await queue.consume(callback=callback, no_ack=True, exclusive=True)
|
|
|
|
def download_file(self, channel: Channel, file_info, message_data, output_dir):
|
|
"""
|
|
Downloads a file from RabbitMQ based on the provided file information.
|
|
|
|
Args:
|
|
channel: The RabbitMQ channel.
|
|
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']
|
|
filepath = os.path.join(output_dir, filename)
|
|
self.ensure_directory_exists(filepath)
|
|
|
|
print(f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name} to {filepath}")
|
|
|
|
bytes_received = 0
|
|
with open(filepath, 'wb') as f:
|
|
async def callback(in_message: AbstractIncomingMessage):
|
|
nonlocal bytes_received
|
|
f.write(in_message.body)
|
|
bytes_received += len(in_message.body)
|
|
# print(f"Received {len(body)} bytes for {filename}. Total: {bytes_received}/{size}") # uncomment for verbose
|
|
|
|
if bytes_received >= size:
|
|
print(f"Finished downloading {filename}")
|
|
f.close()
|
|
# Update global download status
|
|
message_req_info = json.loads(message_data.get('req-info', '{}'))
|
|
if message_req_info:
|
|
self.download_status[message_req_info['from']][filename] = True
|
|
await in_message.channel.close() # stop consuming after the file is downloaded
|
|
|
|
self.start_queue2(channel, stream_name, callback)
|
|
|
|
def on_message(self, in_message: AbstractIncomingMessage, body: bytes) -> asyncio.Future:
|
|
"""
|
|
Callback function for handling messages from the 'amq-cleverswarm' queue.
|
|
|
|
Args:
|
|
ch: The RabbitMQ channel.
|
|
method: The delivery method.
|
|
properties: The message properties.
|
|
body: The message body (a JSON string).
|
|
:param body:
|
|
"""
|
|
message_data = json.loads(body.decode())
|
|
req_info = json.loads(message_data.get('req-info', '{}')) # Handle missing 'req-info'
|
|
self.future = asyncio.Future()
|
|
|
|
if not req_info:
|
|
print("Received message without 'req-info', skipping.")
|
|
in_message.ack(False)
|
|
self.future.set_result(False)
|
|
return self.future
|
|
|
|
for file_key in req_info['files']:
|
|
files_info = req_info['files'].get(file_key, []) # Handle missing file types
|
|
self.files_to_download.extend(files_info)
|
|
|
|
if not self.files_to_download:
|
|
print("No files to download in this message.")
|
|
self.future.set_result(False)
|
|
return self.future
|
|
|
|
# Create a thread for each file download
|
|
output_dir = "downloaded_files" # You can change this
|
|
for file_info in self.files_to_download:
|
|
thread = threading.Thread(
|
|
target=self.download_file,
|
|
args=(in_message.channel, file_info, message_data, output_dir)
|
|
)
|
|
self.threads.append(thread)
|
|
thread.start()
|
|
|
|
# Wait for all download threads to complete
|
|
asyncio.run(self.wait_for_download_completion(self.threads))
|
|
|
|
# After all files are downloaded, acknowledge the original message
|
|
print("All files are being downloaded. Acknowledging the main message.")
|
|
in_message.ack(False)
|
|
print(f"Message from {in_message.routing_key} acknowledged.")
|
|
return self.future
|
|
|
|
async def wait_for_download_completion(self, threads):
|
|
"""
|
|
Waits for all download threads to complete.
|
|
|
|
Args:
|
|
threads (list): A list of threads to wait for.
|
|
"""
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
self.future.set_result(True)
|