diff --git a/app/commands/file.py b/app/commands/file.py index c17ed119..4d4aacca 100644 --- a/app/commands/file.py +++ b/app/commands/file.py @@ -200,9 +200,10 @@ def file_put(**kwargs): # noqa: C901 if source_file: upload_event['valid_source'] = src_file_info - simple_upload(upload_event, num_of_thread=thread, output_path=output_path) + item_ids = simple_upload(upload_event, num_of_thread=thread, output_path=output_path) - srv_manifest.attach_manifest(attribute, result_file, zone) if attribute else None + # since only file upload can attach manifest, take the first file object + srv_manifest.attach_manifest(attribute, item_ids[0], zone) if attribute else None message_handler.SrvOutPutHandler.all_file_uploaded() diff --git a/app/services/file_manager/file_manifests.py b/app/services/file_manager/file_manifests.py index 7756a48f..e6e98093 100644 --- a/app/services/file_manager/file_manifests.py +++ b/app/services/file_manager/file_manifests.py @@ -61,9 +61,9 @@ def validate_template(self, manifest_json): return False, res_json @require_valid_token() - def attach(self, manifest_json, file_name, zone): + def attach(self, manifest_json: dict, item_id: str, zone: str): url = self.app_config.Connections.url_bff + '/v1/manifest/attach' - manifest_json['file_name'] = file_name + manifest_json['item_id'] = item_id manifest_json['zone'] = zone headers = { 'Authorization': 'Bearer ' + self.user.access_token, @@ -151,8 +151,8 @@ def validate_manifest(self, manifest, raise_error=True): validation_error = '' return validation, validation_error - def attach_manifest(self, manifest, file_name, zone): - res = self.attach(manifest, file_name, zone) + def attach_manifest(self, manifest: dict, item_id: str, zone: str): + res = self.attach(manifest, item_id, zone) if res.get('code') != 200: error = res.get('error_msg') if self.interactive: diff --git a/app/services/file_manager/file_upload/file_upload.py b/app/services/file_manager/file_upload/file_upload.py index 8b1b7e2e..eea78d61 100644 --- a/app/services/file_manager/file_upload/file_upload.py +++ b/app/services/file_manager/file_upload/file_upload.py @@ -9,6 +9,7 @@ from multiprocessing.pool import ThreadPool from typing import Any from typing import Dict +from typing import List from typing import Tuple import click @@ -107,7 +108,7 @@ def simple_upload( # noqa: C901 upload_event, num_of_thread: int = 1, output_path: str = None, -): +) -> List[str]: upload_start_time = time.time() my_file = upload_event.get('file') project_code = upload_event.get('project_code') @@ -181,14 +182,20 @@ def simple_upload( # noqa: C901 pool = ThreadPool(num_of_thread + 1) pool.apply_async(upload_client.upload_token_refresh) + on_success_res = [] for file_object in pre_upload_infos: chunk_res = upload_client.stream_upload(file_object, pool) # NOTE: if there is some racing error make the combine chunks # out of thread pool. - pool.apply_async( + res = pool.apply_async( upload_client.on_succeed, args=(file_object, tags, chunk_res), ) + on_success_res.append(res) + + # finish the upload once all on success api return + # otherwise wait for 1 second and check again + [res.wait() for res in on_success_res] upload_client.set_finish_upload() pool.close() @@ -208,6 +215,8 @@ def simple_upload( # noqa: C901 num_of_file = len(upload_file_path) logger.info(f'Upload Time: {time.time() - upload_start_time:.2f}s for {num_of_file:d} files') + return [file_object.item_id for file_object in pre_upload_infos] + def resume_upload( manifest_json: Dict[str, Any], @@ -263,15 +272,20 @@ def resume_upload( pool = ThreadPool(num_of_thread + 1) pool.apply_async(upload_client.upload_token_refresh) + on_success_res = [] for file_object in unfinished_items: - upload_client.stream_upload(file_object, pool) + chunk_res = upload_client.stream_upload(file_object, pool) # NOTE: if there is some racing error make the combine chunks # out of thread pool. - pool.apply_async( + res = pool.apply_async( upload_client.on_succeed, - args=(file_object, manifest_json.get('tags')), + args=(file_object, manifest_json.get('tags'), chunk_res), ) + on_success_res.append(res) + # finish the upload once all on success api return + # otherwise wait for 1 second and check again + [res.wait() for res in on_success_res] upload_client.set_finish_upload() pool.close() diff --git a/tests/app/commands/test_file.py b/tests/app/commands/test_file.py index d560be53..a82fd154 100644 --- a/tests/app/commands/test_file.py +++ b/tests/app/commands/test_file.py @@ -2,11 +2,45 @@ # # Contact Indoc Research for any questions regarding the use of this source code. +import click + +from app.commands.file import file_put from app.commands.file import file_resume +from app.services.file_manager.file_upload.models import FileObject from app.services.output_manager.error_handler import ECustomizedError from app.services.output_manager.error_handler import customized_error_msg +def test_file_upload_command_success_with_attribute(mocker, cli_runner): + project_code = 'test_project' + target_folder = 'admin' + + mocker.patch('app.commands.file.identify_target_folder', return_value=(project_code, target_folder)) + mocker.patch('app.commands.file.validate_upload_event', return_value={'source_file': '', 'attribute': 'test'}) + mocker.patch('app.commands.file.assemble_path', return_value=('test', {'id': 'id'}, True, 'test')) + + mocker.patch('app.services.file_manager.file_upload.models.FileObject.generate_meta', return_value=(1, 1)) + test_obj = FileObject('resumable_id', 'job_id', 'item_id', 'object/path', 'local_path') + + simple_upload_mock = mocker.patch('app.commands.file.simple_upload', return_value=[test_obj]) + attribute_mock = mocker.patch( + 'app.services.file_manager.file_manifests.SrvFileManifests.attach_manifest', return_value=None + ) + + # create a test file + runner = click.testing.CliRunner() + with runner.isolated_filesystem(): + with open('test.txt', 'w') as f: + f.write('test.txt') + + result = cli_runner.invoke( + file_put, ['--project-path', 'test', '--thread', 1, '--attribute', 'test.json', 'test.txt'] + ) + assert result.exit_code == 0 + simple_upload_mock.assert_called_once() + attribute_mock.assert_called_once() + + def test_resumable_upload_command_success(mocker, cli_runner): mocker.patch('os.path.exists', return_value=True) # mock the open function