diff --git a/components/clp-package-utils/clp_package_utils/general.py b/components/clp-package-utils/clp_package_utils/general.py index 4e0a24745..401e0c77b 100644 --- a/components/clp-package-utils/clp_package_utils/general.py +++ b/components/clp-package-utils/clp_package_utils/general.py @@ -8,7 +8,7 @@ import typing import uuid from enum import auto -from typing import List, Tuple +from typing import List, Optional, Tuple import yaml from clp_py_utils.clp_config import ( @@ -286,7 +286,7 @@ def dump_container_config( def generate_container_start_cmd( - container_name: str, container_mounts: List[CLPDockerMounts], container_image: str + container_name: str, container_mounts: List[Optional[DockerMount]], container_image: str ) -> List[str]: """ Generates the command to start a container with the given mounts and name. diff --git a/components/clp-package-utils/clp_package_utils/scripts/decompress.py b/components/clp-package-utils/clp_package_utils/scripts/decompress.py index 45c4f6dff..5daca3f1a 100755 --- a/components/clp-package-utils/clp_package_utils/scripts/decompress.py +++ b/components/clp-package-utils/clp_package_utils/scripts/decompress.py @@ -3,7 +3,6 @@ import pathlib import subprocess import sys - from typing import Optional from clp_py_utils.clp_config import CLPConfig @@ -60,7 +59,9 @@ def validate_and_load_config( return None -def handle_decompression_command(parsed_args, clp_home: pathlib.Path, default_config_file_path: pathlib.Path): +def handle_decompression_command( + parsed_args, clp_home: pathlib.Path, default_config_file_path: pathlib.Path +): paths_to_decompress_file_path = None if parsed_args.files_from: paths_to_decompress_file_path = pathlib.Path(parsed_args.files_from) @@ -75,7 +76,9 @@ def handle_decompression_command(parsed_args, clp_home: pathlib.Path, default_co extraction_dir.mkdir(exist_ok=True) # Validate and load config file - clp_config = validate_and_load_config(clp_home, pathlib.Path(parsed_args.config), clp_home) + clp_config = validate_and_load_config( + clp_home, pathlib.Path(parsed_args.config), default_config_file_path + ) if not clp_config: return -1 @@ -134,7 +137,9 @@ def handle_decompression_command(parsed_args, clp_home: pathlib.Path, default_co def handle_extraction(parsed_args, clp_home: pathlib.Path, default_config_file_path: pathlib.Path): # Validate and load config file - clp_config = validate_and_load_config(clp_home, pathlib.Path(parsed_args.config), clp_home) + clp_config = validate_and_load_config( + clp_home, pathlib.Path(parsed_args.config), default_config_file_path + ) if not clp_config: return -1 @@ -164,8 +169,8 @@ def handle_extraction(parsed_args, clp_home: pathlib.Path, default_config_file_p else: extract_cmd.append("--path") extract_cmd.append(str(parsed_args.path)) - if parsed_args.target_size: - extract_cmd.append("--target-size") + if parsed_args.target_uncompressed_size: + extract_cmd.append("--target-uncompressed-size") extract_cmd.append(str(parsed_args.target_size)) cmd = container_start_cmd + extract_cmd subprocess.run(cmd, check=True) @@ -201,7 +206,9 @@ def main(argv): # IR extraction command parser ir_extraction_parser = command_args_parser.add_parser(IR_EXTRACTION_COMMAND) ir_extraction_parser.add_argument("msg_ix", type=int, help="Message index.") - ir_extraction_parser.add_argument("--target-size", type=int, help="Target IR size.") + ir_extraction_parser.add_argument( + "--target-uncompressed-size", type=int, help="Target uncompressed IR size." + ) group = ir_extraction_parser.add_mutually_exclusive_group(required=True) group.add_argument("--orig-file-id", type=str, help="Original file ID.") diff --git a/components/clp-package-utils/clp_package_utils/scripts/native/decompress.py b/components/clp-package-utils/clp_package_utils/scripts/native/decompress.py index 71e29246a..52c1bcf81 100755 --- a/components/clp-package-utils/clp_package_utils/scripts/native/decompress.py +++ b/components/clp-package-utils/clp_package_utils/scripts/native/decompress.py @@ -6,7 +6,6 @@ import sys import uuid from contextlib import closing - from typing import Optional import yaml @@ -44,9 +43,10 @@ def get_orig_file_id(db_config: Database, path: str) -> Optional[str]: Returns the original file id of the file with the given path If multiple files have the same path, this method returns the first file returned from the archive metadata database - :param db_config: config of the archive metadata database - :param path: original path of the file - :return: orig_file_id. None if no file matches with input path + :param db_config: + :param path: + :return: The original file id of a file with the given path. + If no file matches with input path, returns None """ sql_adapter = SQL_Adapter(db_config) with closing(sql_adapter.create_connection(True)) as db_conn, closing( @@ -72,16 +72,24 @@ def get_orig_file_id(db_config: Database, path: str) -> Optional[str]: return results[0]["orig_file_id"] -def create_and_monitor_ir_extraction_job_in_db( +def submit_and_monitor_ir_extraction_job_in_db( db_config: Database, orig_file_id: str, msg_ix: int, - target_size: int | None, -): + target_uncompressed_size: int | None, +) -> None: + """ + Submits an IR extraction job to the scheduler database and + waits until the job finishes. + :param db_config: + :param orig_file_id: + :param msg_ix: + :param target_uncompressed_size: + """ extract_ir_config = ExtractIrJobConfig( orig_file_id=orig_file_id, msg_ix=msg_ix, - target_size=target_size, + target_uncompressed_size=target_uncompressed_size, ) sql_adapter = SQL_Adapter(db_config) @@ -98,18 +106,27 @@ async def do_extract( db_config: Database, orig_file_id: str, msg_ix: int, - target_size: int | None, + target_uncompressed_size: int | None, ): await run_function_in_process( - create_and_monitor_ir_extraction_job_in_db, db_config, orig_file_id, msg_ix, target_size + submit_and_monitor_ir_extraction_job_in_db, + db_config, + orig_file_id, + msg_ix, + target_uncompressed_size, ) def handle_ir_extraction_command( - parsed_args: argparse.Namespace, - clp_home: pathlib.Path, - default_config_file_path: pathlib.Path + parsed_args: argparse.Namespace, clp_home: pathlib.Path, default_config_file_path: pathlib.Path ) -> int: + """ + Handles the IR extraction command. + :param parsed_args: + :param clp_home: + :param default_config_file_path: + :return: 0 on success, -1 otherwise. + """ # Validate and load config file clp_config = validate_and_load_config_file( clp_home, pathlib.Path(parsed_args.config), default_config_file_path @@ -129,13 +146,15 @@ def handle_ir_extraction_command( clp_config.database, orig_file_id, parsed_args.msg_ix, - parsed_args.target_size, + parsed_args.target_uncompressed_size, ) ) except asyncio.CancelledError: logger.error("Extraction cancelled.") return -1 + return 0 + def validate_and_load_config_file( clp_home: pathlib.Path, @@ -261,7 +280,9 @@ def main(argv): # IR extraction command parser ir_extraction_parser = command_args_parser.add_parser(IR_EXTRACTION_COMMAND) ir_extraction_parser.add_argument("msg_ix", type=int, help="Message index.") - ir_extraction_parser.add_argument("--target-size", type=int, help="Target IR size.") + ir_extraction_parser.add_argument( + "--target-uncompressed-size", type=int, help="Target uncompressed IR size." + ) group = ir_extraction_parser.add_mutually_exclusive_group(required=True) group.add_argument("--orig-file-id", type=str, help="Original file ID.") diff --git a/components/clp-package-utils/clp_package_utils/scripts/native/search.py b/components/clp-package-utils/clp_package_utils/scripts/native/search.py index 869670242..7dd247fa5 100755 --- a/components/clp-package-utils/clp_package_utils/scripts/native/search.py +++ b/components/clp-package-utils/clp_package_utils/scripts/native/search.py @@ -25,11 +25,6 @@ submit_query_job, wait_for_query_job, ) -from clp_package_utils.scripts.native.utils import ( - run_function_in_process, - submit_query_job, - wait_for_query_job, -) # Setup logging # Create logger