feat(clp-package): Expose dataset selection in package CLI scripts and web UI.

components/clp-package-utils/clp_package_utils/general.py

@@ -7,28 +7,34 @@
 import subprocess
 import typing
 import uuid
+from contextlib import closing
 from enum import auto
 from typing import List, Optional, Tuple
 
 import yaml
 from clp_py_utils.clp_config import (
     CLP_DEFAULT_CREDENTIALS_FILE_PATH,
+    CLP_DEFAULT_DATASET_NAME,
     CLPConfig,
+    Database,
     DB_COMPONENT_NAME,
     QUEUE_COMPONENT_NAME,
     REDIS_COMPONENT_NAME,
     REDUCER_COMPONENT_NAME,
     RESULTS_CACHE_COMPONENT_NAME,
+    StorageEngine,
     StorageType,
     WEBUI_COMPONENT_NAME,
     WorkerConfig,
 )
+from clp_py_utils.clp_metadata_db_utils import validate_and_cache_dataset
 from clp_py_utils.core import (
     get_config_value,
     make_config_path_absolute,
     read_yaml_config_file,
     validate_path_could_be_dir,
 )
+from clp_py_utils.sql_adapter import SQL_Adapter
 from strenum import KebabCaseStrEnum
 
 # CONSTANTS
@@ -556,3 +562,31 @@ def validate_path_for_container_mount(path: pathlib.Path) -> None:
                 f"Invalid path: `{path}` cannot be under '{prefix}' which may overlap with a path"
                 f" in the container."
             )
+
+
+def validate_dataset(clp_config: CLPConfig, input_dataset: Optional[str]) -> Optional[str]:
+    """
+    Checks if the provided dataset currently exists in the metadata database, or provides a default
+    value if the CLP_S storage engine is used.
+    :param clp_config:
+    :param input_dataset: Dataset from the CLI.
+    :return: The validated dataset to use.
+    :raise: ValueError
+    """
+    storage_engine: StorageEngine = clp_config.package.storage_engine
+    if StorageEngine.CLP_S != storage_engine:
+        if input_dataset is not None:
+            raise ValueError("Dataset selection is only enabled for CLP_S storage engine.")
+        return None
+
+    dataset: str = CLP_DEFAULT_DATASET_NAME if input_dataset is None else input_dataset
+    db_config: Database = clp_config.database
+    sql_adapter: SQL_Adapter = SQL_Adapter(db_config)
+    clp_db_connection_params: dict[str, any] = db_config.get_clp_connection_params_and_type(True)
+    table_prefix: str = clp_db_connection_params["table_prefix"]
+    with closing(sql_adapter.create_connection(True)) as db_conn, closing(
+        db_conn.cursor(dictionary=True)
+    ) as db_cursor:
+        if not validate_and_cache_dataset(db_cursor, table_prefix, dataset):
+            raise ValueError(f"Dataset `{dataset}` does not exist.")
+    return dataset

components/clp-package-utils/clp_package_utils/scripts/archive_manager.py

@@ -18,6 +18,7 @@
     get_clp_home,
     load_config_file,
     validate_and_load_db_credentials_file,
+    validate_dataset,
 )
 
 # Command/Argument Constants
@@ -61,6 +62,12 @@ def main(argv: typing.List[str]) -> int:
         default=str(default_config_file_path),
         help="CLP package configuration file.",
     )
+    args_parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help="The dataset that the archives belong to.",
+    )
 
     # Top-level commands
     subparsers: argparse._SubParsersAction[argparse.ArgumentParser] = args_parser.add_subparsers(
@@ -158,6 +165,8 @@ def main(argv: typing.List[str]) -> int:
         logger.exception("Failed to load config.")
         return -1
 
+    dataset = validate_dataset(clp_config, parsed_args.dataset)
+
     storage_type: StorageType = clp_config.archive_output.storage.type
     if StorageType.FS != storage_type:
         logger.error(f"Archive deletion is not supported for storage type: {storage_type}.")
@@ -196,6 +205,7 @@ def main(argv: typing.List[str]) -> int:
         "python3",
         "-m", "clp_package_utils.scripts.native.archive_manager",
         "--config", str(generated_config_path_on_container),
+        "--dataset", str(dataset),
         str(subcommand),
     ]
     # fmt : on

components/clp-package-utils/clp_package_utils/scripts/compress.py

@@ -4,13 +4,14 @@
 import subprocess
 import sys
 import uuid
-from typing import List
+from typing import List, Optional
 
 from clp_py_utils.clp_config import StorageEngine
 from job_orchestration.scheduler.job_config import InputType
 
 from clp_package_utils.general import (
     CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
+    CLP_DEFAULT_DATASET_NAME,
     CONTAINER_INPUT_LOGS_ROOT_DIR,
     dump_container_config,
     generate_container_config,
@@ -63,6 +64,7 @@ def _generate_logs_list(
 
 def _generate_compress_cmd(
     parsed_args: argparse.Namespace,
+    dataset: Optional[str],
     config_path: pathlib.Path,
     logs_list_path: pathlib.Path,
 ) -> List[str]:
@@ -74,6 +76,9 @@ def _generate_compress_cmd(
         "--config", str(config_path),
     ]
     # fmt: on
+    if dataset is not None:
+        compress_cmd.append("--dataset")
+        compress_cmd.append(dataset)
     if parsed_args.timestamp_key is not None:
         compress_cmd.append("--timestamp-key")
         compress_cmd.append(parsed_args.timestamp_key)
@@ -131,6 +136,12 @@ def main(argv):
         default=str(default_config_file_path),
         help="CLP package configuration file.",
     )
+    args_parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help="The dataset that the archives belong to.",
+    )
     args_parser.add_argument(
         "--timestamp-key",
         help="The path (e.g. x.y) for the field containing the log event's timestamp.",
@@ -163,10 +174,16 @@ def main(argv):
         return -1
 
     input_type = clp_config.logs_input.type
+    storage_engine: StorageEngine = clp_config.package.storage_engine
+
+    dataset: Optional[str] = parsed_args.dataset
+    if StorageEngine.CLP_S == storage_engine and dataset is None:
+        dataset = CLP_DEFAULT_DATASET_NAME
+
     if InputType.FS == input_type:
         _validate_fs_input_args(parsed_args, args_parser)
     elif InputType.S3 == input_type:
-        _validate_s3_input_args(parsed_args, args_parser, clp_config.package.storage_engine)
+        _validate_s3_input_args(parsed_args, args_parser, storage_engine)
     else:
         raise ValueError(f"Unsupported input type: {input_type}.")
 
@@ -198,7 +215,7 @@ def main(argv):
         container_name, necessary_mounts, clp_config.execution_container
     )
     compress_cmd = _generate_compress_cmd(
-        parsed_args, generated_config_path_on_container, logs_list_path_on_container
+        parsed_args, dataset, generated_config_path_on_container, logs_list_path_on_container
     )
     cmd = container_start_cmd + compress_cmd
     subprocess.run(cmd, check=True)

components/clp-package-utils/clp_package_utils/scripts/decompress.py

@@ -5,7 +5,12 @@
 import sys
 from typing import Optional
 
-from clp_py_utils.clp_config import CLPConfig, StorageEngine, StorageType
+from clp_py_utils.clp_config import (
+    CLP_DEFAULT_DATASET_NAME,
+    CLPConfig,
+    StorageEngine,
+    StorageType,
+)
 
 from clp_package_utils.general import (
     CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
@@ -22,6 +27,7 @@
     JobType,
     load_config_file,
     validate_and_load_db_credentials_file,
+    validate_dataset,
     validate_path_could_be_dir,
 )
 
@@ -167,13 +173,19 @@ def handle_extract_stream_cmd(
 
     storage_type = clp_config.archive_output.storage.type
     storage_engine = clp_config.package.storage_engine
+    job_command = parsed_args.command
+
     if StorageType.S3 == storage_type and StorageEngine.CLP == storage_engine:
         logger.error(
             f"Stream extraction is not supported for archive storage type `{storage_type}` with"
             f" storage engine `{storage_engine}`."
         )
         return -1
 
+    if EXTRACT_JSON_CMD == job_command and StorageEngine.CLP_S != storage_engine:
+        logger.error(f"Json extraction is only supported with storage engine `{storage_engine}`.")
+        return -1
+
     container_name = generate_container_name(str(JobType.IR_EXTRACTION))
     container_clp_config, mounts = generate_container_config(clp_config, clp_home)
     generated_config_path_on_container, generated_config_path_on_host = dump_container_config(
@@ -185,7 +197,6 @@ def handle_extract_stream_cmd(
     )
 
     # fmt: off
-    job_command = parsed_args.command
     extract_cmd = [
         "python3",
         "-m", "clp_package_utils.scripts.native.decompress",
@@ -207,6 +218,9 @@ def handle_extract_stream_cmd(
             extract_cmd.append(str(parsed_args.target_uncompressed_size))
     elif EXTRACT_JSON_CMD == job_command:
         extract_cmd.append(str(parsed_args.archive_id))
+        dataset = validate_dataset(clp_config, parsed_args.dataset)
+        extract_cmd.append("--dataset")
+        extract_cmd.append(dataset)
         if parsed_args.target_chunk_size:
             extract_cmd.append("--target-chunk-size")
             extract_cmd.append(str(parsed_args.target_chunk_size))
@@ -267,6 +281,12 @@ def main(argv):
     # JSON extraction command parser
     json_extraction_parser = command_args_parser.add_parser(EXTRACT_JSON_CMD)
     json_extraction_parser.add_argument("archive_id", type=str, help="Archive ID")
+    json_extraction_parser.add_argument(
+        "--dataset",
+        type=str,
+        default=CLP_DEFAULT_DATASET_NAME,
+        help="The dataset that the archives belong to.",
+    )
     json_extraction_parser.add_argument(
         "--target-chunk-size",
         type=int,

components/clp-package-utils/clp_package_utils/scripts/native/archive_manager.py

@@ -7,11 +7,7 @@
 from contextlib import closing
 from pathlib import Path
 
-from clp_py_utils.clp_config import (
-    CLP_DEFAULT_DATASET_NAME,
-    Database,
-    StorageEngine,
-)
+from clp_py_utils.clp_config import Database
 from clp_py_utils.clp_metadata_db_utils import (
     get_archive_tags_table_name,
     get_archives_table_name,
@@ -99,6 +95,12 @@ def main(argv: typing.List[str]) -> int:
         default=str(default_config_file_path),
         help="CLP configuration file.",
     )
+    args_parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help="The dataset that the archives belong to.",
+    )
 
     # Top-level commands
     subparsers: argparse._SubParsersAction[argparse.ArgumentParser] = args_parser.add_subparsers(
@@ -186,22 +188,17 @@ def main(argv: typing.List[str]) -> int:
         logger.exception("Failed to load config.")
         return -1
 
-    storage_engine: StorageEngine = clp_config.package.storage_engine
     database_config: Database = clp_config.database
     archives_dir: Path = clp_config.archive_output.get_directory()
     if not archives_dir.exists():
         logger.error("`archive_output.directory` doesn't exist.")
         return -1
 
-    dataset: typing.Optional[str] = None
-    if StorageEngine.CLP_S == storage_engine:
-        dataset = CLP_DEFAULT_DATASET_NAME
-
     if FIND_COMMAND == parsed_args.subcommand:
         return _find_archives(
             archives_dir,
             database_config,
-            dataset,
+            parsed_args.dataset,
             parsed_args.begin_ts,
             parsed_args.end_ts,
         )
@@ -212,7 +209,7 @@ def main(argv: typing.List[str]) -> int:
             return _delete_archives(
                 archives_dir,
                 database_config,
-                dataset,
+                parsed_args.dataset,
                 delete_handler,
                 parsed_args.dry_run,
             )
@@ -223,7 +220,7 @@ def main(argv: typing.List[str]) -> int:
             return _delete_archives(
                 archives_dir,
                 database_config,
-                dataset,
+                parsed_args.dataset,
                 delete_handler,
                 parsed_args.dry_run,
             )
@@ -302,7 +299,7 @@ def _find_archives(
 def _delete_archives(
     archives_dir: Path,
     database_config: Database,
-    dataset: str,
+    dataset: typing.Optional[str],
     delete_handler: DeleteHandler,
     dry_run: bool = False,
 ) -> int:

components/clp-package-utils/clp_package_utils/scripts/native/compress.py

@@ -10,10 +10,8 @@
 import brotli
 import msgpack
 from clp_py_utils.clp_config import (
-    CLP_DEFAULT_DATASET_NAME,
     CLPConfig,
     COMPRESSION_JOBS_TABLE_NAME,
-    StorageEngine,
 )
 from clp_py_utils.pretty_size import pretty_size
 from clp_py_utils.s3_utils import parse_s3_url
@@ -140,18 +138,12 @@ def _generate_clp_io_config(
     logs_to_compress: List[str],
     parsed_args: argparse.Namespace,
 ) -> Union[S3InputConfig, FsInputConfig]:
-    dataset = (
-        CLP_DEFAULT_DATASET_NAME
-        if StorageEngine.CLP_S == clp_config.package.storage_engine
-        else None
-    )
-
     input_type = clp_config.logs_input.type
     if InputType.FS == input_type:
         if len(logs_to_compress) == 0:
             raise ValueError("No input paths given.")
         return FsInputConfig(
-            dataset=dataset,
+            dataset=parsed_args.dataset,
             paths_to_compress=logs_to_compress,
             timestamp_key=parsed_args.timestamp_key,
             path_prefix_to_remove=str(CONTAINER_INPUT_LOGS_ROOT_DIR),
@@ -166,7 +158,7 @@ def _generate_clp_io_config(
         region_code, bucket_name, key_prefix = parse_s3_url(s3_url)
         aws_authentication = clp_config.logs_input.aws_authentication
         return S3InputConfig(
-            dataset=dataset,
+            dataset=parsed_args.dataset,
             region_code=region_code,
             bucket=bucket_name,
             key_prefix=key_prefix,
@@ -203,6 +195,12 @@ def main(argv):
         default=str(default_config_file_path),
         help="CLP package configuration file.",
     )
+    args_parser.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help="The dataset that the archives belong to.",
+    )
     args_parser.add_argument(
         "-f",
         "--logs-list",