diff --git a/app/services/file_manager/file_upload/models.py b/app/services/file_manager/file_upload/models.py index 8d0b481b..aaf36ffb 100644 --- a/app/services/file_manager/file_upload/models.py +++ b/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 diff --git a/app/services/file_manager/file_upload/upload_client.py b/app/services/file_manager/file_upload/upload_client.py index bea34f29..0640f3e5 100644 --- a/app/services/file_manager/file_upload/upload_client.py +++ b/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, diff --git a/tests/app/services/file_manager/file_upload/test_upload_client.py b/tests/app/services/file_manager/file_upload/test_upload_client.py index d1809d0b..53c3cd76 100644 --- a/tests/app/services/file_manager/file_upload/test_upload_client.py +++ b/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