fix(package): Refactor dataset handling in the package.

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

@@ -8,13 +8,15 @@
 from pathlib import Path
 
 from clp_py_utils.clp_config import (
-    ARCHIVE_TAGS_TABLE_SUFFIX,
-    ARCHIVES_TABLE_SUFFIX,
     CLP_DEFAULT_DATASET_NAME,
     Database,
-    FILES_TABLE_SUFFIX,
     StorageEngine,
 )
+from clp_py_utils.clp_metadata_db_utils import (
+    get_archive_tags_table_name,
+    get_archives_table_name,
+    get_files_table_name,
+)
 from clp_py_utils.sql_adapter import SQL_Adapter
 
 from clp_package_utils.general import (
@@ -191,12 +193,15 @@ def main(argv: typing.List[str]) -> int:
         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,
-            storage_engine,
-            CLP_DEFAULT_DATASET_NAME,
+            dataset,
             parsed_args.begin_ts,
             parsed_args.end_ts,
         )
@@ -207,8 +212,7 @@ def main(argv: typing.List[str]) -> int:
             return _delete_archives(
                 archives_dir,
                 database_config,
-                storage_engine,
-                CLP_DEFAULT_DATASET_NAME,
+                dataset,
                 delete_handler,
                 parsed_args.dry_run,
             )
@@ -219,8 +223,7 @@ def main(argv: typing.List[str]) -> int:
             return _delete_archives(
                 archives_dir,
                 database_config,
-                storage_engine,
-                CLP_DEFAULT_DATASET_NAME,
+                dataset,
                 delete_handler,
                 parsed_args.dry_run,
             )
@@ -235,8 +238,7 @@ def main(argv: typing.List[str]) -> int:
 def _find_archives(
     archives_dir: Path,
     database_config: Database,
-    storage_engine: StorageEngine,
-    dataset: str,
+    dataset: typing.Optional[str],
     begin_ts: int,
     end_ts: int = typing.Optional[int],
 ) -> int:
@@ -245,7 +247,6 @@ def _find_archives(
     `begin_ts <= archive.begin_timestamp` and `archive.end_timestamp <= end_ts`.
     :param archives_dir:
     :param database_config:
-    :param storage_engine:
     :param dataset:
     :param begin_ts:
     :param end_ts:
@@ -259,16 +260,14 @@ def _find_archives(
             database_config.get_clp_connection_params_and_type(True)
         )
         table_prefix: str = clp_db_connection_params["table_prefix"]
-        if StorageEngine.CLP_S == storage_engine:
-            table_prefix = f"{table_prefix}{dataset}_"
 
         with closing(sql_adapter.create_connection(True)) as db_conn, closing(
             db_conn.cursor(dictionary=True)
         ) as db_cursor:
             query_params: typing.List[int] = [begin_ts]
             query: str = (
                 f"""
-                SELECT id FROM `{table_prefix}{ARCHIVES_TABLE_SUFFIX}`
+                SELECT id FROM `{get_archives_table_name(table_prefix, dataset)}`
                 WHERE begin_timestamp >= %s
                 """
             )
@@ -285,9 +284,10 @@ def _find_archives(
                 return 0
 
             logger.info(f"Found {len(archive_ids)} archives within the specified time range.")
+            archive_output_dir = archives_dir / dataset if dataset is not None else archives_dir
             for archive_id in archive_ids:
                 logger.info(archive_id)
-                archive_path: Path = archives_dir / dataset / archive_id
+                archive_path = archive_output_dir / archive_id
                 if not archive_path.is_dir():
                     logger.warning(f"Archive {archive_id} in database not found on disk.")
 
@@ -302,7 +302,6 @@ def _find_archives(
 def _delete_archives(
     archives_dir: Path,
     database_config: Database,
-    storage_engine: StorageEngine,
     dataset: str,
     delete_handler: DeleteHandler,
     dry_run: bool = False,
@@ -312,7 +311,6 @@ def _delete_archives(
 
     :param archives_dir:
     :param database_config:
-    :param storage_engine:
     :param dataset:
     :param delete_handler: Object to handle differences between by-filter and by-ids delete types.
     :param dry_run: If True, no changes will be made to the database or disk.
@@ -327,8 +325,6 @@ def _delete_archives(
             database_config.get_clp_connection_params_and_type(True)
         )
         table_prefix = clp_db_connection_params["table_prefix"]
-        if StorageEngine.CLP_S == storage_engine:
-            table_prefix = f"{table_prefix}{dataset}_"
 
         with closing(sql_adapter.create_connection(True)) as db_conn, closing(
             db_conn.cursor(dictionary=True)
@@ -341,7 +337,7 @@ def _delete_archives(
 
             db_cursor.execute(
                 f"""
-                DELETE FROM `{table_prefix}{ARCHIVES_TABLE_SUFFIX}`
+                DELETE FROM `{get_archives_table_name(table_prefix, dataset)}`
                 WHERE {query_criteria}
                 RETURNING id
                 """,
@@ -360,14 +356,14 @@ def _delete_archives(
 
             db_cursor.execute(
                 f"""
-                DELETE FROM `{table_prefix}{FILES_TABLE_SUFFIX}`
+                DELETE FROM `{get_files_table_name(table_prefix, dataset)}`
                 WHERE archive_id in ({ids_list_string})
                 """
             )
 
             db_cursor.execute(
                 f"""
-                DELETE FROM `{table_prefix}{ARCHIVE_TAGS_TABLE_SUFFIX}`
+                DELETE FROM `{get_archive_tags_table_name(table_prefix, dataset)}`
                 WHERE archive_id in ({ids_list_string})
                 """
             )
@@ -387,8 +383,9 @@ def _delete_archives(
 
     logger.info(f"Finished deleting archives from the database.")
 
+    archive_output_dir: Path = archives_dir / dataset if dataset is not None else archives_dir
     for archive_id in archive_ids:
-        archive_path: Path = archives_dir / dataset / archive_id
+        archive_path = archive_output_dir / archive_id
         if not archive_path.is_dir():
             logger.warning(f"Archive {archive_id} is not a directory. Skipping deletion.")
             continue

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

@@ -4,13 +4,17 @@
 import pathlib
 import sys
 import time
-import typing
 from contextlib import closing
-from typing import List
+from typing import List, Optional, Union
 
 import brotli
 import msgpack
-from clp_py_utils.clp_config import CLPConfig, COMPRESSION_JOBS_TABLE_NAME
+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
 from clp_py_utils.sql_adapter import SQL_Adapter
@@ -132,14 +136,18 @@ def handle_job(sql_adapter: SQL_Adapter, clp_io_config: ClpIoConfig, no_progress
 
 
 def _generate_clp_io_config(
-    clp_config: CLPConfig, logs_to_compress: List[str], parsed_args: argparse.Namespace
-) -> typing.Union[S3InputConfig, FsInputConfig]:
+    clp_config: CLPConfig,
+    logs_to_compress: List[str],
+    parsed_args: argparse.Namespace,
+    dataset: Optional[str],
+) -> Union[S3InputConfig, FsInputConfig]:
     input_type = clp_config.logs_input.type
 
     if InputType.FS == input_type:
         if len(logs_to_compress) == 0:
             raise ValueError(f"No input paths given.")
         return FsInputConfig(
+            dataset=dataset,
             paths_to_compress=logs_to_compress,
             timestamp_key=parsed_args.timestamp_key,
             path_prefix_to_remove=str(CONTAINER_INPUT_LOGS_ROOT_DIR),
@@ -154,6 +162,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,
             region_code=region_code,
             bucket=bucket_name,
             key_prefix=key_prefix,
@@ -224,7 +233,12 @@ def main(argv):
 
     logs_to_compress = _get_logs_to_compress(pathlib.Path(parsed_args.logs_list).resolve())
 
-    clp_input_config = _generate_clp_io_config(clp_config, logs_to_compress, parsed_args)
+    dataset = (
+        CLP_DEFAULT_DATASET_NAME
+        if StorageEngine.CLP_S == clp_config.package.storage_engine
+        else None
+    )
+    clp_input_config = _generate_clp_io_config(clp_config, logs_to_compress, parsed_args, dataset)
     clp_output_config = OutputConfig.parse_obj(clp_config.archive_output)
     if parsed_args.tags:
         tag_list = [tag.strip().lower() for tag in parsed_args.tags.split(",") if tag]

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

@@ -13,8 +13,8 @@
     CLP_DEFAULT_DATASET_NAME,
     CLPConfig,
     Database,
-    FILES_TABLE_SUFFIX,
 )
+from clp_py_utils.clp_metadata_db_utils import get_files_table_name
 from clp_py_utils.sql_adapter import SQL_Adapter
 from job_orchestration.scheduler.constants import QueryJobStatus, QueryJobType
 from job_orchestration.scheduler.job_config import (
@@ -54,8 +54,9 @@ def get_orig_file_id(db_config: Database, path: str) -> Optional[str]:
     with closing(sql_adapter.create_connection(True)) as db_conn, closing(
         db_conn.cursor(dictionary=True)
     ) as db_cursor:
+        files_table_name = get_files_table_name(table_prefix, None)
         db_cursor.execute(
-            f"SELECT orig_file_id FROM `{table_prefix}{FILES_TABLE_SUFFIX}` WHERE path = (%s)",
+            f"SELECT orig_file_id FROM `{files_table_name}` WHERE path = (%s)",
             (path,),
         )
         results = db_cursor.fetchall()

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

@@ -10,7 +10,12 @@
 
 import msgpack
 import pymongo
-from clp_py_utils.clp_config import Database, ResultsCache
+from clp_py_utils.clp_config import (
+    CLP_DEFAULT_DATASET_NAME,
+    Database,
+    ResultsCache,
+    StorageEngine,
+)
 from clp_py_utils.sql_adapter import SQL_Adapter
 from job_orchestration.scheduler.constants import QueryJobStatus, QueryJobType
 from job_orchestration.scheduler.job_config import AggregationConfig, SearchJobConfig
@@ -32,6 +37,7 @@
 def create_and_monitor_job_in_db(
     db_config: Database,
     results_cache: ResultsCache,
+    dataset: str | None,
     wildcard_query: str,
     tags: str | None,
     begin_timestamp: int | None,
@@ -43,6 +49,7 @@ def create_and_monitor_job_in_db(
     count_by_time_bucket_size: int | None,
 ):
     search_config = SearchJobConfig(
+        dataset=dataset,
         query_string=wildcard_query,
         begin_timestamp=begin_timestamp,
         end_timestamp=end_timestamp,
@@ -113,6 +120,7 @@ async def worker_connection_handler(reader: asyncio.StreamReader, writer: asynci
 async def do_search_without_aggregation(
     db_config: Database,
     results_cache: ResultsCache,
+    dataset: str | None,
     wildcard_query: str,
     tags: str | None,
     begin_timestamp: int | None,
@@ -147,6 +155,7 @@ async def do_search_without_aggregation(
             create_and_monitor_job_in_db,
             db_config,
             results_cache,
+            dataset,
             wildcard_query,
             tags,
             begin_timestamp,
@@ -184,6 +193,7 @@ async def do_search_without_aggregation(
 async def do_search(
     db_config: Database,
     results_cache: ResultsCache,
+    dataset: str | None,
     wildcard_query: str,
     tags: str | None,
     begin_timestamp: int | None,
@@ -198,6 +208,7 @@ async def do_search(
         await do_search_without_aggregation(
             db_config,
             results_cache,
+            dataset,
             wildcard_query,
             tags,
             begin_timestamp,
@@ -211,6 +222,7 @@ async def do_search(
             create_and_monitor_job_in_db,
             db_config,
             results_cache,
+            dataset,
             wildcard_query,
             tags,
             begin_timestamp,
@@ -281,11 +293,17 @@ def main(argv):
         logger.exception("Failed to load config.")
         return -1
 
+    dataset = (
+        CLP_DEFAULT_DATASET_NAME
+        if StorageEngine.CLP_S == clp_config.package.storage_engine
+        else None
+    )
     try:
         asyncio.run(
             do_search(
                 clp_config.database,
                 clp_config.results_cache,
+                dataset,
                 parsed_args.wildcard_query,
                 parsed_args.tags,
                 parsed_args.begin_time,

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

@@ -14,7 +14,6 @@
 import yaml
 from clp_py_utils.clp_config import (
     ALL_TARGET_NAME,
-    ARCHIVES_TABLE_SUFFIX,
     AwsAuthType,
     CLP_DEFAULT_DATASET_NAME,
     CLPConfig,
@@ -23,7 +22,6 @@
     COMPRESSION_WORKER_COMPONENT_NAME,
     CONTROLLER_TARGET_NAME,
     DB_COMPONENT_NAME,
-    FILES_TABLE_SUFFIX,
     QUERY_JOBS_TABLE_NAME,
     QUERY_SCHEDULER_COMPONENT_NAME,
     QUERY_WORKER_COMPONENT_NAME,
@@ -35,6 +33,10 @@
     StorageType,
     WEBUI_COMPONENT_NAME,
 )
+from clp_py_utils.clp_metadata_db_utils import (
+    get_archives_table_name,
+    get_files_table_name,
+)
 from clp_py_utils.s3_utils import generate_container_auth_options
 from job_orchestration.scheduler.constants import QueueName
 from pydantic import BaseModel
@@ -868,13 +870,14 @@ def start_webui(
     # Read, update, and write back client's and server's settings.json
     clp_db_connection_params = clp_config.database.get_clp_connection_params_and_type(True)
     table_prefix = clp_db_connection_params["table_prefix"]
+    dataset: Optional[str] = None
     if StorageEngine.CLP_S == clp_config.package.storage_engine:
-        table_prefix = f"{table_prefix}{CLP_DEFAULT_DATASET_NAME}_"
+        dataset = CLP_DEFAULT_DATASET_NAME
     client_settings_json_updates = {
         "ClpStorageEngine": clp_config.package.storage_engine,
         "MongoDbSearchResultsMetadataCollectionName": clp_config.webui.results_metadata_collection_name,
-        "SqlDbClpArchivesTableName": f"{table_prefix}{ARCHIVES_TABLE_SUFFIX}",
-        "SqlDbClpFilesTableName": f"{table_prefix}{FILES_TABLE_SUFFIX}",
+        "SqlDbClpArchivesTableName": get_archives_table_name(table_prefix, dataset),
+        "SqlDbClpFilesTableName": get_files_table_name(table_prefix, dataset),
         "SqlDbCompressionJobsTableName": COMPRESSION_JOBS_TABLE_NAME,
     }
     client_settings_json = read_and_update_settings_json(
@@ -884,6 +887,7 @@ def start_webui(
         client_settings_json_file.write(json.dumps(client_settings_json))
 
     server_settings_json_updates = {
+        "ClpStorageEngine": clp_config.package.storage_engine,
         "SqlDbHost": clp_config.database.host,
         "SqlDbPort": clp_config.database.port,
         "SqlDbName": clp_config.database.name,