amq-adapter-python/amqp/adapter/file_downloader.py

"""
implements the file downloader for the amqp adapter, where file is sent in chunks over RabbitMQ
"""

import asyncio
import json
import os
import sys
import traceback
import logging
from collections import defaultdict
from aio_pika import connect_robust
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel, AbstractRobustQueue, AbstractQueue

# 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

    eof = IN_PROGRESS
    fsize = 0
    logging.info(f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}")
    if queue_name:
        connection = await connect_robust(rabbitmq_url, loop=loop)
        try:
            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)
                        fsize = fsize + message.body_size
                        eof = message.headers.get('eof', END_OF_FILE)
                        if eof != IN_PROGRESS:
                            break
        except Exception as e:
            logging.error(f"Error downloading file: {e}")
            tb = traceback.extract_tb(sys.exc_info()[2])[-1]
            logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
            await connection.close()
    return (fsize, eof)


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.
    """
    res = bytearray()
    eof = IN_PROGRESS
    logging.info(f"Downloading buffer from stream: {queue_name}")
    if queue_name:
        connection = await connect_robust(rabbitmq_url, loop=loop)
        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"Error downloading buffer: {e}")
            tb = traceback.extract_tb(sys.exc_info()[2])[-1]
            logging.error(f" [E] Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")

        await connection.close()

    return (res, eof)


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("No files to download in this message.")
            await message.ack()
            return False

        # Create a task for each file download
        tasks = []
        for file_info in files_to_download:
            logging.info(f"about ti 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("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"[E] [{message.delivery_tag}] Error processing message: {e}")
        tb = traceback.extract_tb(sys.exc_info()[2])[-1]
        logging.error(f"   Error in {tb.filename}, line {tb.lineno}, function '{tb.name}': {e}")
        await message.ack() # ignore the message

    return download_locations