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

Co-authored-by: Marco <[email protected]>
Co-authored-by: Haiqi Xu <[email protected]>
Co-authored-by: Kirk Rodrigues <[email protected]>
Co-authored-by: hoophalab <[email protected]>

Author: 5 people

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

@@ -2,6 +2,7 @@
 import errno
 import os
 import pathlib
+import re
 import secrets
 import socket
 import subprocess
@@ -25,7 +26,11 @@
     WEBUI_COMPONENT_NAME,
     WorkerConfig,
 )
-from clp_py_utils.clp_metadata_db_utils import fetch_existing_datasets
+from clp_py_utils.clp_metadata_db_utils import (
+    fetch_existing_datasets,
+    MYSQL_TABLE_NAME_MAX_LEN,
+    TABLE_SUFFIX_MAX_LEN,
+)
 from clp_py_utils.core import (
     get_config_value,
     make_config_path_absolute,
@@ -562,7 +567,7 @@ def validate_path_for_container_mount(path: pathlib.Path) -> None:
             )
 
 
-def validate_dataset(db_config: Database, dataset: str) -> None:
+def validate_dataset_exists(db_config: Database, dataset: str) -> None:
     """
     Validates that `dataset` exists in the metadata database.
 
@@ -578,3 +583,32 @@ def validate_dataset(db_config: Database, dataset: str) -> None:
     ) as db_cursor:
         if dataset not in fetch_existing_datasets(db_cursor, table_prefix):
             raise ValueError(f"Dataset `{dataset}` doesn't exist.")
+
+
+def validate_dataset_name(clp_table_prefix: str, dataset_name: str) -> None:
+    """
+    Validates that the given dataset name abides by the following rules:
+    - Its length won't cause any metadata table names to exceed MySQL's max table name length.
+    - It only contains alphanumeric characters and underscores.
+
+    :param clp_table_prefix:
+    :param dataset_name:
+    :raise: ValueError if the dataset name is invalid.
+    """
+    if re.fullmatch(r"\w+", dataset_name) is None:
+        raise ValueError(
+            f"Invalid dataset name: `{dataset_name}`. Names can only contain alphanumeric"
+            f" characters and underscores."
+        )
+
+    dataset_name_max_len = (
+        MYSQL_TABLE_NAME_MAX_LEN
+        - len(clp_table_prefix)
+        - 1  # For the separator between the dataset name and the table suffix
+        - TABLE_SUFFIX_MAX_LEN
+    )
+    if len(dataset_name) > dataset_name_max_len:
+        raise ValueError(
+            f"Invalid dataset name: `{dataset_name}`. Names can only be a maximum of"
+            f" {dataset_name_max_len} characters long."
+        )

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

@@ -5,7 +5,11 @@
 import typing
 from pathlib import Path
 
-from clp_py_utils.clp_config import StorageType
+from clp_py_utils.clp_config import (
+    CLP_DEFAULT_DATASET_NAME,
+    StorageEngine,
+    StorageType,
+)
 
 from clp_package_utils.general import (
     CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
@@ -18,6 +22,7 @@
     get_clp_home,
     load_config_file,
     validate_and_load_db_credentials_file,
+    validate_dataset_name,
 )
 
 # Command/Argument Constants
@@ -61,6 +66,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(
@@ -163,6 +174,20 @@ def main(argv: typing.List[str]) -> int:
         logger.error(f"Archive deletion is not supported for storage type: {storage_type}.")
         return -1
 
+    storage_engine: StorageEngine = clp_config.package.storage_engine
+    dataset = parsed_args.dataset
+    if StorageEngine.CLP_S == storage_engine:
+        dataset = CLP_DEFAULT_DATASET_NAME if dataset is None else dataset
+        try:
+            clp_db_connection_params = clp_config.database.get_clp_connection_params_and_type(True)
+            validate_dataset_name(clp_db_connection_params["table_prefix"], dataset)
+        except Exception as e:
+            logger.error(e)
+            return -1
+    elif dataset is not None:
+        logger.error(f"Dataset selection is not supported for storage engine: {storage_engine}.")
+        return -1
+
     # Validate input depending on subcommands
     if (DEL_COMMAND == subcommand and DEL_BY_FILTER_SUBCOMMAND == parsed_args.del_subcommand) or (
         FIND_COMMAND == subcommand
@@ -196,9 +221,12 @@ def main(argv: typing.List[str]) -> int:
         "python3",
         "-m", "clp_package_utils.scripts.native.archive_manager",
         "--config", str(generated_config_path_on_container),
-        str(subcommand),
     ]
     # fmt : on
+    if dataset is not None:
+        archive_manager_cmd.append("--dataset")
+        archive_manager_cmd.append(dataset)
+    archive_manager_cmd.append(subcommand)
 
     # Add subcommand-specific arguments
     if DEL_COMMAND == subcommand:

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

@@ -4,9 +4,12 @@
 import subprocess
 import sys
 import uuid
-from typing import List
+from typing import List, Optional
 
-from clp_py_utils.clp_config import StorageEngine
+from clp_py_utils.clp_config import (
+    CLP_DEFAULT_DATASET_NAME,
+    StorageEngine,
+)
 from job_orchestration.scheduler.job_config import InputType
 
 from clp_package_utils.general import (
@@ -20,6 +23,7 @@
     JobType,
     load_config_file,
     validate_and_load_db_credentials_file,
+    validate_dataset_name,
 )
 
 logger = logging.getLogger(__file__)
@@ -63,6 +67,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 +79,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 +139,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.",
@@ -162,11 +176,25 @@ def main(argv):
         logger.exception("Failed to load config.")
         return -1
 
+    storage_engine: StorageEngine = clp_config.package.storage_engine
+    dataset = parsed_args.dataset
+    if StorageEngine.CLP_S == storage_engine:
+        dataset = CLP_DEFAULT_DATASET_NAME if dataset is None else dataset
+        try:
+            clp_db_connection_params = clp_config.database.get_clp_connection_params_and_type(True)
+            validate_dataset_name(clp_db_connection_params["table_prefix"], dataset)
+        except Exception as e:
+            logger.error(e)
+            return -1
+    elif dataset is not None:
+        logger.error(f"Dataset selection is not supported for storage engine: {storage_engine}.")
+        return -1
+
     input_type = clp_config.logs_input.type
     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 +226,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_name,
     validate_path_could_be_dir,
 )
 
@@ -174,6 +180,24 @@ def handle_extract_stream_cmd(
         )
         return -1
 
+    job_command = parsed_args.command
+    if EXTRACT_JSON_CMD == job_command and StorageEngine.CLP_S != storage_engine:
+        logger.error(f"JSON extraction is not supported for storage engine `{storage_engine}`.")
+        return -1
+
+    dataset = parsed_args.dataset
+    if StorageEngine.CLP_S == storage_engine:
+        dataset = CLP_DEFAULT_DATASET_NAME if dataset is None else dataset
+        try:
+            clp_db_connection_params = clp_config.database.get_clp_connection_params_and_type(True)
+            validate_dataset_name(clp_db_connection_params["table_prefix"], dataset)
+        except Exception as e:
+            logger.error(e)
+            return -1
+    elif dataset is not None:
+        logger.error(f"Dataset selection is not supported for 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 +209,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 +230,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))
+        if dataset is not None:
+            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 +293,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=None,
+        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,
@@ -24,6 +20,7 @@
     CLPConfig,
     get_clp_home,
     load_config_file,
+    validate_dataset_exists,
 )
 
 # Command/Argument Constants
@@ -99,6 +96,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,17 +189,20 @@ 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
+    dataset = parsed_args.dataset
+    if dataset is not None:
+        try:
+            validate_dataset_exists(database_config, dataset)
+        except Exception as e:
+            logger.error(e)
+            return -1
+
     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,
@@ -302,7 +308,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: