Merge pull request #212 from Aiven-Open/daniel.blasina-calculate-max-file-size

Transfer return max file size

Author: giacomo-alzetta-aiven

rohmu/object_storage/azure.py

@@ -30,6 +30,7 @@
 from rohmu.object_storage.config import (  # noqa: F401
     AZURE_ENDPOINT_SUFFIXES as ENDPOINT_SUFFIXES,
     AZURE_MAX_BLOCK_SIZE as MAX_BLOCK_SIZE,
+    AZURE_MAX_NUM_PARTS_PER_UPLOAD as MAX_NUM_PARTS_PER_UPLOAD,
     AzureObjectStorageConfig as Config,
     calculate_azure_max_block_size as calculate_max_block_size,
 )
@@ -139,6 +140,9 @@ def create_object_store_if_needed(self) -> None:
             else:
                 raise TransferObjectStoreInitializationError() from ex
 
+    def calculate_max_unknown_file_size(self) -> int:
+        return MAX_NUM_PARTS_PER_UPLOAD * MAX_BLOCK_SIZE
+
     def get_blob_service_client(self) -> BlobServiceClient:
         if self._blob_service_client is None:
             self._blob_service_client = BlobServiceClient.from_connection_string(
@@ -178,7 +182,7 @@ def conn_string(
     def copy_file(
         self, *, source_key: str, destination_key: str, metadata: Optional[Metadata] = None, **kwargs: Any
     ) -> None:
-        timeout = kwargs.get("timeout") or 15.0
+        timeout = kwargs.get("timeout") or 15
         self._copy_file_from_bucket(
             source_bucket=self, source_key=source_key, destination_key=destination_key, metadata=metadata, timeout=timeout
         )
@@ -190,7 +194,7 @@ def _copy_file_from_bucket(
         source_key: str,
         destination_key: str,
         metadata: Optional[Metadata] = None,
-        timeout: float = 15.0,
+        timeout: int = 15,
     ) -> None:
         source_path = source_bucket.format_key_for_backend(source_key, remove_slash_prefix=True, trailing_slash=False)
         source_client = source_bucket.get_blob_service_client().get_blob_client(source_bucket.container_name, source_path)

rohmu/object_storage/base.py

@@ -40,6 +40,7 @@
 
 import logging
 import os
+import sys
 
 KEY_TYPE_OBJECT = "object"
 KEY_TYPE_PREFIX = "prefix"
@@ -230,7 +231,7 @@ def _copy_file_from_bucket(
         source_key: str,
         destination_key: str,
         metadata: Optional[Metadata] = None,
-        timeout: float = 15.0,
+        timeout: int = 15,
     ) -> None:
         raise NotImplementedError
 
@@ -271,6 +272,9 @@ def delete_tree(self, key: str, preserve_trailing_slash: bool = False) -> None:
         names = [item["name"] for item in self.list_path(key, with_metadata=False, deep=True)]
         self.delete_keys(names, preserve_trailing_slash=preserve_trailing_slash)
 
+    def calculate_max_unknown_file_size(self) -> int:
+        return sys.maxsize
+
     def get_contents_to_file(
         self, key: str, filepath_to_store_to: AnyPath, *, progress_callback: ProgressProportionCallbackType = None
     ) -> Metadata:

rohmu/object_storage/config.py

@@ -50,14 +50,14 @@ def calculate_azure_max_block_size() -> int:
 # to that 5 TiB increase the block size based on host memory; we don't want to use the max 100 for all
 # hosts because the uploader will allocate (with default settings) 3 x block size of memory.
 AZURE_MAX_BLOCK_SIZE: Final[int] = calculate_azure_max_block_size()
-
+AZURE_MAX_NUM_PARTS_PER_UPLOAD: Final[int] = 10000
 
 # googleapiclient download performs some 3-4 times better with 50 MB chunk size than 5 MB chunk size;
 # but decrypting/decompressing big chunks needs a lot of memory so use smaller chunks on systems with less
 # than 2 GB RAM
 GOOGLE_DOWNLOAD_CHUNK_SIZE: Final[int] = 1024 * 1024 * 5 if (get_total_memory() or 0) < 2048 else 1024 * 1024 * 50
 GOOGLE_UPLOAD_CHUNK_SIZE: Final[int] = 1024 * 1024 * 5
-
+GOOGLE_MAX_NUM_PARTS_PER_UPLOAD: Final[int] = 10000
 
 LOCAL_CHUNK_SIZE: Final[int] = 1024 * 1024
 
@@ -85,6 +85,7 @@ def calculate_s3_chunk_size() -> int:
 
 SWIFT_CHUNK_SIZE: Final[int] = 1024 * 1024 * 5  # 5 Mi
 SWIFT_SEGMENT_SIZE: Final[int] = 1024 * 1024 * 1024 * 3  # 3 Gi
+SWIFT_MAX_NUM_PARTS_PER_UPLOAD = 10000
 
 
 class AzureObjectStorageConfig(StorageModel):

rohmu/object_storage/google.py

@@ -43,6 +43,7 @@
 )
 from rohmu.object_storage.config import (
     GOOGLE_DOWNLOAD_CHUNK_SIZE as DOWNLOAD_CHUNK_SIZE,
+    GOOGLE_MAX_NUM_PARTS_PER_UPLOAD as MAX_NUM_PARTS_PER_UPLOAD,
     GOOGLE_UPLOAD_CHUNK_SIZE as UPLOAD_CHUNK_SIZE,
     GoogleObjectStorageConfig as Config,
 )
@@ -401,6 +402,9 @@ def _copy_file_from_bucket(
                 self.notifier.object_copied(key=destination_key, size=size, metadata=metadata)
             reporter.report(self.stats)
 
+    def calculate_max_unknown_file_size(self) -> int:
+        return MAX_NUM_PARTS_PER_UPLOAD * UPLOAD_CHUNK_SIZE
+
     def get_metadata_for_key(self, key: str) -> Metadata:
         path = self.format_key_for_backend(key)
         with self._object_client(not_found=path) as clob:

rohmu/object_storage/s3.py

@@ -178,6 +178,9 @@ def __init__(
     def _verify_object_storage_unwrapped(self) -> None:
         self.check_or_create_bucket(create_if_needed=False)
 
+    def calculate_max_unknown_file_size(self) -> int:
+        return self.default_multipart_chunk_size * S3_MAX_NUM_PARTS_PER_UPLOAD
+
     def verify_object_storage(self) -> None:
         try:
             self._verify_object_storage_unwrapped()

rohmu/object_storage/swift.py

@@ -24,6 +24,7 @@
 )
 from rohmu.object_storage.config import (
     SWIFT_CHUNK_SIZE as CHUNK_SIZE,
+    SWIFT_MAX_NUM_PARTS_PER_UPLOAD as MAX_NUM_PARTS_PER_UPLOAD,
     SWIFT_SEGMENT_SIZE as SEGMENT_SIZE,
     SwiftObjectStorageConfig as Config,
 )
@@ -153,6 +154,9 @@ def get_metadata_for_key(self, key: str) -> Metadata:
         path = self.format_key_for_backend(key)
         return self._metadata_for_key(path)
 
+    def calculate_max_unknown_file_size(self) -> int:
+        return MAX_NUM_PARTS_PER_UPLOAD * CHUNK_SIZE
+
     def _metadata_for_key(self, key: str, *, resolve_manifest: bool = False) -> Metadata:
         try:
             headers = self.conn.head_object(self.container_name, key)

test/object_storage/test_local.py

@@ -15,6 +15,7 @@
 import json
 import os
 import pytest
+import sys
 
 
 def test_store_file_from_disk() -> None:
@@ -37,6 +38,15 @@ def test_store_file_from_disk() -> None:
         )
 
 
+def test_calculate_max_unknown_file_size() -> None:
+    with TemporaryDirectory() as destdir:
+        transfer = LocalTransfer(
+            directory=destdir,
+        )
+
+        assert transfer.calculate_max_unknown_file_size() == sys.maxsize
+
+
 def test_store_file_object() -> None:
     with TemporaryDirectory() as destdir:
         notifier = MagicMock()

test/object_storage/test_s3.py

@@ -49,6 +49,21 @@ def _get_session(cls: S3Transfer) -> Iterator[MagicMock]:
     return transfer
 
 
+def test_calculate_max_unknown_file_size(mocker: Any) -> None:
+    segment_size = 100
+    transfer = make_mock_transfer(
+        mocker,
+        {
+            "region": "test-region",
+            "bucket_name": "test-bucket",
+            "prefix": "test-prefix",
+            "segment_size": segment_size,
+        },
+    )
+
+    assert transfer.calculate_max_unknown_file_size() == segment_size * S3_MAX_NUM_PARTS_PER_UPLOAD
+
+
 @pytest.fixture(name="infra")
 def fixture_infra(mocker: Any) -> Iterator[S3Infra]:
     operation = mocker.patch("rohmu.common.statsd.StatsClient.operation")