Create new uploaded file version rather than append to it, if file of the same name already exists

amqp/adapter/file_downloader.py

@@ -15,11 +15,34 @@ END_OF_FILE = 1
 CANCELLED = 2
 
 
-def ensure_directory_exists(filepath):
+def ensure_directory_exists(filepath) -> str:
-    """Ensures that the directory for the given filepath exists."""
+    """
-    dir_name = os.path.dirname(filepath)
+    Ensures that the directory for the given filepath exists, and handles file versioning.
-    if dir_name and not os.path.exists(dir_name):
+
-        os.makedirs(dir_name)
+    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):
@@ -39,8 +62,11 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir)
     stream_name = file_info['streamName']
     queue_name = file_info['queue_name']
     filepath = os.path.join(output_dir, filename)
-    ensure_directory_exists(filepath)
+    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}")
@@ -67,7 +93,7 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir)
     return (fsize, eof)
 
 
-async def download_buffer(loop, rabbitmq_url, queue_name, ttl = 604800000) -> (bytearray, int):
+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
@@ -97,7 +123,9 @@ async def download_buffer(loop, rabbitmq_url, queue_name, ttl = 604800000) -> (b
             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()
+
+        await connection.close()
+
     return (res, eof)
 
 

amqp/adapter/test_file_downloader.py

@@ -0,0 +1,46 @@
+import os
+import tempfile
+from unittest import TestCase
+
+import pytest
+
+from amqp.adapter.file_downloader import ensure_directory_exists
+
+class TestFileDownloader(TestCase):
+    def test_ensure_directory_exists(self):
+        """
+        Test function for ensure_directory_exists.  Creates a temporary directory
+        and files to test the versioning logic.
+        """
+
+        # Create a temporary directory
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Test case 1: File does not exist
+            filepath1 = os.path.join(temp_dir, "test1.txt")
+            result1 = ensure_directory_exists(filepath1)
+            assert result1 == filepath1, f"Test Case 1 Failed: {result1} != {filepath1}"
+            # Create the file
+            with open(filepath1, "w") as f:
+                f.write("test")
+
+            # Test case 2: File exists, should create versioned file
+            filepath2 = os.path.join(temp_dir, "test2.txt")
+            result2 = ensure_directory_exists(filepath2)
+            assert result2 == filepath2, f"Test Case 2 Failed: {result2} != {filepath2}"
+            with open(filepath2, "w") as f:
+                f.write("test")
+            result2_v1 = ensure_directory_exists(filepath2)
+            assert result2_v1 == os.path.join(temp_dir, "test2.1.txt"), f"Test Case 2 Version 1 Failed: {result2_v1} != {os.path.join(temp_dir, 'test2.1.txt')}"
+
+            # Create a few versioned files
+            with open(os.path.join(temp_dir, "test3.txt"), "w") as f:
+                f.write("test")
+            with open(os.path.join(temp_dir, "test3.1.txt"), "w") as f:
+                f.write("test")
+            with open(os.path.join(temp_dir, "test3.2.txt"), "w") as f:
+                f.write("test")
+            filepath3 = os.path.join(temp_dir, "test3.txt")
+            result3_v3 = ensure_directory_exists(filepath3)
+            assert result3_v3 == os.path.join(temp_dir, "test3.3.txt"), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
+
+            print("All test cases passed!")

requirements.txt

@@ -7,3 +7,4 @@ opentelemetry-exporter-otlp
 opentelemetry-exporter-prometheus
 opentelemetry-instrumentation-pika
 python_multipart
+pytest