Add retrying to upload functionality (#36)

(DIS-1804)

Author: Miauwkeru

acquire/acquire.py

@@ -36,7 +36,7 @@
 from acquire.log import get_file_handler, reconfigure_log_file, setup_logging
 from acquire.outputs import OUTPUTS
 from acquire.uploaders.minio import MinIO
-from acquire.uploaders.plugin import UploaderPlugin
+from acquire.uploaders.plugin import UploaderPlugin, upload_files_using_uploader
 from acquire.uploaders.plugin_registry import UploaderRegistry
 from acquire.utils import (
     check_and_set_acquire_args,
@@ -1689,7 +1689,7 @@ def upload_files(
     log.debug("Proxies: %s (no_proxy = %s)", proxies, no_proxy)
 
     try:
-        upload_plugin.upload_files(paths, proxies)
+        upload_files_using_uploader(upload_plugin, paths, proxies)
     except Exception:
         log.error("Upload %s FAILED. See log file for details.", paths)
         log.exception("")

acquire/uploaders/minio.py

@@ -1,12 +1,9 @@
-import logging
 import os
 from pathlib import Path
-from typing import Any, List, Optional
+from typing import Any, Optional
 
 from acquire.uploaders.plugin import UploaderPlugin
 
-log = logging.getLogger(__name__)
-
 
 class MinIO(UploaderPlugin):
     def __init__(self, upload: dict[str, str], **kwargs: dict[str, Any]) -> None:
@@ -27,8 +24,8 @@ def __init__(self, upload: dict[str, str], **kwargs: dict[str, Any]) -> None:
         if not all((self.endpoint, self.access_id, self.access_key, self.bucket_name)):
             raise ValueError("Invalid cloud upload configuration")
 
-    def upload_files(self, paths: List[Path], proxies: Optional[dict[str, str]] = None) -> None:
-        """Uploads files using MinIO
+    def prepare_client(self, paths: list[Path], proxies: Optional[dict[str, str]] = None) -> Any:
+        """Prepares a Minio client used to upload files.
 
         Args:
             paths: The files to upload.
@@ -45,12 +42,10 @@ def upload_files(self, paths: List[Path], proxies: Optional[dict[str, str]] = No
 
         http_client = urllib3.proxy_from_url(proxies["http"]) if proxies else None
 
-        client = Minio(self.endpoint, self.access_id, self.access_key, http_client=http_client)
-        for path in paths:
-            try:
-                log.info("Uploading %s", path)
-                client.fput_object(self.bucket_name, os.path.basename(path), path)
-                log.info("Uploaded %s", path)
-            except Exception:
-                log.error("Upload %s FAILED. See log file for details.", path)
-                log.exception("")
+        return Minio(self.endpoint, self.access_id, self.access_key, http_client=http_client)
+
+    def upload_file(self, client: Any, path: Path) -> None:
+        client.fput_object(self.bucket_name, os.path.basename(path), path)
+
+    def finish(self, client: Any) -> None:
+        pass

acquire/uploaders/plugin.py

@@ -1,16 +1,60 @@
+import logging
 from pathlib import Path
-from typing import List, Optional, Protocol, runtime_checkable
+from typing import Any, Optional
 
+log = logging.getLogger(__name__)
 
-@runtime_checkable
-class UploaderPlugin(Protocol):
+__all__ = [
+    "UploaderPlugin",
+    "upload_files_using_uploader",
+]
+
+MAX_RETRIES = 4
+
+
+class UploaderPlugin:
     """Creates a typing definition to which an UploaderPlugin should adhere."""
 
-    def upload_files(self, paths: List[Path], proxies: Optional[dict[str, str]] = None) -> None:
-        """Uploads the files in ``paths`` to a destination.
+    def prepare_client(self, paths: list[Path], proxies: Optional[dict[str, str]] = None) -> Any:
+        """Prepares a client for the upload."""
+        raise NotImplementedError()
+
+    def upload_file(self, client: Any, path: Path) -> None:
+        """Uploads a file/path using the ``client``."""
+        raise NotImplementedError()
+
+    def finish(self, client: Any) -> None:
+        """A cleanup step or anything required to finish the upload."""
+        raise NotImplementedError()
+
+
+def upload_files_using_uploader(
+    uploader: UploaderPlugin, paths: list[Path], proxies: Optional[dict[str, str]] = None
+) -> None:
+    """Uploads the files in ``paths`` to a destination.
+
+    Args:
+        uploader: The plugin used to upload files.
+        paths: A list of files to upload.
+        proxies: Proxies used as an intermediate during an upload.
+    """
+    paths = [Path(path) if isinstance(path, str) else path for path in paths]
+    client = uploader.prepare_client(paths, proxies)
+
+    for path in paths:
+        for retry in range(MAX_RETRIES):
+            if retry == MAX_RETRIES - 1:
+                error_log = ("Upload %s FAILED after too many attempts. Stopping.", path)
+            else:
+                error_log = ("Upload %s FAILED. See log file for details. Retrying", path)
+
+            try:
+                log.info("Uploading %s", path)
+                uploader.upload_file(client, path)
+                log.info("Uploaded %s", path)
+                break
+            except Exception:
+                log.error(*error_log)
+                log.exception("")
 
-        Args:
-            paths: A list of files to upload.
-            proxies: Proxies used as an intermediate during an upload.
-        """
-        ...
+    uploader.finish(client)

tests/test_minio_uploader.py

@@ -1,8 +1,10 @@
 from pathlib import Path
+from unittest.mock import Mock, patch
 
 import pytest
 
 from acquire.uploaders.minio import MinIO
+from acquire.uploaders.plugin import upload_files_using_uploader
 from acquire.uploaders.plugin_registry import UploaderRegistry
 
 
@@ -59,4 +61,21 @@ def test_minio_valueerror(minio_plugin):
 
 def test_upload_files(minio_instance):
     # Generates an internal error, which is caught
-    minio_instance.upload_files([Path("hello")])
+    upload_files_using_uploader(minio_instance, [Path("hello")])
+
+
+def test_upload_file_multiple_failures(minio_instance: MinIO):
+    mocked_upload = Mock()
+    mocked_upload.side_effect = [Exception, Exception, Exception, "Hello"]
+    minio_instance.upload_file = mocked_upload
+
+    with patch("acquire.uploaders.plugin.log") as mocked_logger:
+        # Only three retry attempts get made
+        upload_files_using_uploader(minio_instance, [Path("hello")])
+        mocked_logger.error.assert_called_with("Upload %s FAILED. See log file for details. Retrying", Path("hello"))
+
+    # One extra gives a value error.
+    mocked_upload.side_effect = [Exception, Exception, Exception, Exception]
+    with patch("acquire.uploaders.plugin.log") as mocked_logger:
+        upload_files_using_uploader(minio_instance, [Path("hello")])
+        mocked_logger.error.assert_called_with("Upload %s FAILED after too many attempts. Stopping.", Path("hello"))