PILOT-4152: Update cli logic to fit dynamic chunk size when uploading

app/services/file_manager/file_upload/models.py

@@ -7,7 +7,8 @@
 from os.path import basename
 from os.path import dirname
 from os.path import getsize
-from typing import List
+from typing import Any
+from typing import Dict
 from typing import Tuple
 
 from tqdm import tqdm
@@ -62,7 +63,7 @@ class FileObject:
     total_chunks: int
 
     # resumable info
-    uploaded_chunks: List[dict]
+    uploaded_chunks: Dict[str, Dict[str, Any]]
 
     # progress bar object
     progress_bar = None

app/services/file_manager/file_upload/upload_client.py

@@ -61,7 +61,7 @@ def __init__(
         self.user = UserConfig()
         self.operator = self.user.username
         self.upload_message = upload_message
-        self.chunk_size = AppConfig.Env.chunk_size  # remove
+        self.chunk_size = AppConfig.Env.chunk_size
         self.base_url = {
             AppConfig.Env.green_zone: AppConfig.Connections.url_upload_greenroom,
             AppConfig.Env.core_zone: AppConfig.Connections.url_upload_core,
@@ -111,7 +111,7 @@ def resume_upload(self, unfinished_file_objects: List[FileObject]) -> List[FileO
         Parameter:
             - unfinished_file_objects(List[FileObject]): the unfinished items that need to be resumed.
         return:
-            - list of FileObject: the infomation retrieved from backend.
+            - list of FileObject: the information retrieved from backend.
                 - resumable_id(str): the unique identifier for multipart upload.
                 - object_path(str): the path in the object storage.
                 - local_path(str): the local path of file.
@@ -121,6 +121,7 @@ def resume_upload(self, unfinished_file_objects: List[FileObject]) -> List[FileO
         headers = {'Authorization': 'Bearer ' + self.user.access_token, 'Session-ID': self.user.session_id}
         url = AppConfig.Connections.url_bff + f'/v1/project/{self.project_code}/files/resumable'
         rid_file_object_map = {x.resumable_id: x for x in unfinished_file_objects}
+
         payload = {
             'bucket': self.bucket,
             'zone': self.zone,
@@ -300,8 +301,11 @@ def stream_upload(self, file_object: FileObject, pool: ThreadPool) -> List[Apply
         # after all the chunks have been uploaded.
         chunk_result = []
         while True:
-            chunk = f.read(self.chunk_size)
-            chunk_etag = file_object.uploaded_chunks.get(str(count + 1))
+            chunk = file_object.uploaded_chunks.get(str(count + 1), {})
+            chunk_etag = chunk.get('etag')
+            chunk_size = chunk.get('chunk_size', self.chunk_size)
+
+            chunk = f.read(chunk_size)
             local_chunk_etag = hashlib.md5(chunk).hexdigest()
             if not chunk:
                 break
@@ -312,11 +316,11 @@ def stream_upload(self, file_object: FileObject, pool: ThreadPool) -> List[Apply
                 if chunk_etag != local_chunk_etag:
                     SrvErrorHandler.customized_handle(ECustomizedError.INVALID_CHUNK_UPLOAD, value=count + 1)
                     raise INVALID_CHUNK_ETAG(count + 1)
-                file_object.update_progress(self.chunk_size)
+                file_object.update_progress(chunk_size)
             else:
                 res = pool.apply_async(
                     self.upload_chunk,
-                    args=(file_object, count + 1, chunk, local_chunk_etag),
+                    args=(file_object, count + 1, chunk, local_chunk_etag, chunk_size),
                 )
                 chunk_result.append(res)
 
@@ -326,7 +330,7 @@ def stream_upload(self, file_object: FileObject, pool: ThreadPool) -> List[Apply
 
         return chunk_result
 
-    def upload_chunk(self, file_object: FileObject, chunk_number: int, chunk: str, etag: str) -> None:
+    def upload_chunk(self, file_object: FileObject, chunk_number: int, chunk: str, etag: str, chunk_size: int) -> None:
         """
         Summary:
             The function is to upload a chunk directly into minio storage.
@@ -353,6 +357,7 @@ def upload_chunk(self, file_object: FileObject, chunk_number: int, chunk: str, e
                 'key': file_object.item_id,
                 'upload_id': file_object.resumable_id,
                 'chunk_number': chunk_number,
+                'chunk_size': chunk_size,
             }
             headers = {
                 'Authorization': 'Bearer ' + self.user.access_token,

tests/app/services/file_manager/file_upload/test_upload_client.py

@@ -73,7 +73,7 @@ def test_chunk_upload(httpx_mock, mocker):
     mocker.patch('app.services.file_manager.file_upload.models.FileObject.generate_meta', return_value=(1, 1))
 
     test_obj = FileObject('test', 'test', 'test', 'test', 'test')
-    res = upload_client.upload_chunk(test_obj, 0, b'1', 'test_etag')
+    res = upload_client.upload_chunk(test_obj, 0, b'1', 'test_etag', 10)
 
     assert test_obj.progress_bar.n == 1
     assert res.status_code == 200