Merge remote-tracking branch 'origin/main' into retension_period

Author: haiqi96

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

@@ -23,6 +23,7 @@
     CONTROLLER_TARGET_NAME,
     DB_COMPONENT_NAME,
     GARBAGE_COLLECTOR_NAME,
+    get_components_for_target,
     QUERY_JOBS_TABLE_NAME,
     QUERY_SCHEDULER_COMPONENT_NAME,
     QUERY_WORKER_COMPONENT_NAME,
@@ -884,6 +885,7 @@ def start_webui(
 
     client_settings_json_updates = {
         "ClpStorageEngine": clp_config.package.storage_engine,
+        "ClpQueryEngine": clp_config.package.query_engine,
         "MongoDbSearchResultsMetadataCollectionName": clp_config.webui.results_metadata_collection_name,
         "SqlDbClpArchivesTableName": archives_table_name,
         "SqlDbClpDatasetsTableName": get_datasets_table_name(table_prefix),
@@ -1197,6 +1199,15 @@ def main(argv):
         config_file_path = pathlib.Path(parsed_args.config)
         clp_config = load_config_file(config_file_path, default_config_file_path, clp_home)
 
+        runnable_components = clp_config.get_runnable_components()
+        components_to_start = get_components_for_target(target)
+        components_to_start = components_to_start.intersection(runnable_components)
+
+        # Exit early if no components to start
+        if len(components_to_start) == 0:
+            logger.error(f"{target} not available with current configuration")
+            return -1
+
         # Validate and load necessary credentials
         if target in (
             ALL_TARGET_NAME,
@@ -1278,35 +1289,46 @@ def main(argv):
         conf_dir = clp_home / "etc"
 
         # Start components
-        if target in (ALL_TARGET_NAME, DB_COMPONENT_NAME):
+        if DB_COMPONENT_NAME in components_to_start:
             start_db(instance_id, clp_config, conf_dir)
-        if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, DB_COMPONENT_NAME):
+
+        if (
+            target == CONTROLLER_TARGET_NAME and DB_COMPONENT_NAME in runnable_components
+        ) or DB_COMPONENT_NAME in components_to_start:
             create_db_tables(instance_id, clp_config, container_clp_config, mounts)
-        if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, QUEUE_COMPONENT_NAME):
+
+        if QUEUE_COMPONENT_NAME in components_to_start:
             start_queue(instance_id, clp_config)
-        if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, REDIS_COMPONENT_NAME):
+
+        if REDIS_COMPONENT_NAME in components_to_start:
             start_redis(instance_id, clp_config, conf_dir)
-        if target in (ALL_TARGET_NAME, RESULTS_CACHE_COMPONENT_NAME):
+
+        if RESULTS_CACHE_COMPONENT_NAME in components_to_start:
             start_results_cache(instance_id, clp_config, conf_dir)
-        if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, RESULTS_CACHE_COMPONENT_NAME):
+
+        if (
+            target == CONTROLLER_TARGET_NAME and RESULTS_CACHE_COMPONENT_NAME in runnable_components
+        ) or RESULTS_CACHE_COMPONENT_NAME in components_to_start:
             create_results_cache_indices(instance_id, clp_config, container_clp_config, mounts)
-        if target in (
-            ALL_TARGET_NAME,
-            CONTROLLER_TARGET_NAME,
-            COMPRESSION_SCHEDULER_COMPONENT_NAME,
-        ):
+
+        if COMPRESSION_SCHEDULER_COMPONENT_NAME in components_to_start:
             start_compression_scheduler(instance_id, clp_config, container_clp_config, mounts)
-        if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, QUERY_SCHEDULER_COMPONENT_NAME):
+
+        if QUERY_SCHEDULER_COMPONENT_NAME in components_to_start:
             start_query_scheduler(instance_id, clp_config, container_clp_config, mounts)
-        if target in (ALL_TARGET_NAME, COMPRESSION_WORKER_COMPONENT_NAME):
+
+        if COMPRESSION_WORKER_COMPONENT_NAME in components_to_start:
             start_compression_worker(
                 instance_id, clp_config, container_clp_config, num_workers, mounts
             )
-        if target in (ALL_TARGET_NAME, QUERY_WORKER_COMPONENT_NAME):
+
+        if QUERY_WORKER_COMPONENT_NAME in components_to_start:
             start_query_worker(instance_id, clp_config, container_clp_config, num_workers, mounts)
-        if target in (ALL_TARGET_NAME, REDUCER_COMPONENT_NAME):
+
+        if REDUCER_COMPONENT_NAME in components_to_start:
             start_reducer(instance_id, clp_config, container_clp_config, num_workers, mounts)
-        if target in (ALL_TARGET_NAME, WEBUI_COMPONENT_NAME):
+
+        if WEBUI_COMPONENT_NAME in components_to_start:
             start_webui(instance_id, clp_config, container_clp_config, mounts)
         if target in (ALL_TARGET_NAME, GARBAGE_COLLECTOR_NAME):
             start_garbage_collector(instance_id, clp_config, container_clp_config, mounts)

components/clp-py-utils/clp_py_utils/clp_config.py

@@ -1,6 +1,6 @@
 import pathlib
 from enum import auto
-from typing import Literal, Optional, Union
+from typing import Literal, Optional, Set, Union
 
 from dotenv import dotenv_values
 from pydantic import BaseModel, PrivateAttr, root_validator, validator
@@ -28,10 +28,41 @@
 WEBUI_COMPONENT_NAME = "webui"
 GARBAGE_COLLECTOR_NAME = "garbage_collector"
 
+# Component groups
+GENERAL_SCHEDULING_COMPONENTS = {
+    QUEUE_COMPONENT_NAME,
+    REDIS_COMPONENT_NAME,
+}
+COMPRESSION_COMPONENTS = GENERAL_SCHEDULING_COMPONENTS | {
+    DB_COMPONENT_NAME,
+    COMPRESSION_SCHEDULER_COMPONENT_NAME,
+    COMPRESSION_WORKER_COMPONENT_NAME,
+}
+QUERY_COMPONENTS = GENERAL_SCHEDULING_COMPONENTS | {
+    DB_COMPONENT_NAME,
+    QUERY_SCHEDULER_COMPONENT_NAME,
+    QUERY_WORKER_COMPONENT_NAME,
+    REDUCER_COMPONENT_NAME,
+}
+UI_COMPONENTS = {
+    RESULTS_CACHE_COMPONENT_NAME,
+    WEBUI_COMPONENT_NAME,
+}
+ALL_COMPONENTS = COMPRESSION_COMPONENTS | QUERY_COMPONENTS | UI_COMPONENTS
+
 # Target names
 ALL_TARGET_NAME = ""
 CONTROLLER_TARGET_NAME = "controller"
 
+TARGET_TO_COMPONENTS = {
+    ALL_TARGET_NAME: ALL_COMPONENTS,
+    CONTROLLER_TARGET_NAME: GENERAL_SCHEDULING_COMPONENTS
+    | {
+        COMPRESSION_SCHEDULER_COMPONENT_NAME,
+        QUERY_SCHEDULER_COMPONENT_NAME,
+    },
+}
+
 QUERY_JOBS_TABLE_NAME = "query_jobs"
 QUERY_TASKS_TABLE_NAME = "query_tasks"
 COMPRESSION_JOBS_TABLE_NAME = "compression_jobs"
@@ -50,6 +81,12 @@ class StorageEngine(KebabCaseStrEnum):
     CLP_S = auto()
 
 
+class QueryEngine(KebabCaseStrEnum):
+    CLP = auto()
+    CLP_S = auto()
+    PRESTO = auto()
+
+
 class StorageType(LowercaseStrEnum):
     FS = auto()
     S3 = auto()
@@ -63,10 +100,12 @@ class AwsAuthType(LowercaseStrEnum):
 
 
 VALID_STORAGE_ENGINES = [storage_engine.value for storage_engine in StorageEngine]
+VALID_QUERY_ENGINES = [query_engine.value for query_engine in QueryEngine]
 
 
 class Package(BaseModel):
     storage_engine: str = "clp"
+    query_engine: str = "clp"
 
     @validator("storage_engine")
     def validate_storage_engine(cls, field):
@@ -77,6 +116,37 @@ def validate_storage_engine(cls, field):
             )
         return field
 
+    @validator("query_engine")
+    def validate_query_engine(cls, field):
+        if field not in VALID_QUERY_ENGINES:
+            raise ValueError(
+                f"package.query_engine must be one of the following"
+                f" {'|'.join(VALID_QUERY_ENGINES)}"
+            )
+        return field
+
+    @root_validator
+    def validate_query_engine_package_compatibility(cls, values):
+        query_engine = values.get("query_engine")
+        storage_engine = values.get("storage_engine")
+
+        if query_engine in [QueryEngine.CLP, QueryEngine.CLP_S]:
+            if query_engine != storage_engine:
+                raise ValueError(
+                    f"query_engine '{query_engine}' is only compatible with "
+                    f"storage_engine '{query_engine}'."
+                )
+        elif query_engine == QueryEngine.PRESTO:
+            if storage_engine != StorageEngine.CLP_S:
+                raise ValueError(
+                    f"query_engine '{QueryEngine.PRESTO}' is only compatible with "
+                    f"storage_engine '{StorageEngine.CLP_S}'."
+                )
+        else:
+            raise ValueError(f"Unsupported query_engine '{query_engine}'.")
+
+        return values
+
 
 class Database(BaseModel):
     type: str = "mariadb"
@@ -799,6 +869,12 @@ def load_redis_credentials_from_file(self):
                 f"Credentials file '{self.credentials_file_path}' does not contain key '{ex}'."
             )
 
+    def get_runnable_components(self) -> Set[str]:
+        if QueryEngine.PRESTO == self.package.query_engine:
+            return COMPRESSION_COMPONENTS | UI_COMPONENTS
+        else:
+            return ALL_COMPONENTS
+
     def dump_to_primitive_dict(self):
         d = self.dict()
         d["logs_input"] = self.logs_input.dump_to_primitive_dict()
@@ -833,3 +909,12 @@ def dump_to_primitive_dict(self):
         d["stream_output"] = self.stream_output.dump_to_primitive_dict()
 
         return d
+
+
+def get_components_for_target(target: str) -> Set[str]:
+    if target in TARGET_TO_COMPONENTS:
+        return TARGET_TO_COMPONENTS[target]
+    elif target in ALL_COMPONENTS:
+        return {target}
+    else:
+        return set()

components/package-template/src/etc/clp-config.yml

@@ -12,6 +12,7 @@
 #
 #package:
 #  storage_engine: "clp"
+#  query_engine: "clp"
 #
 #database:
 #  type: "mariadb"  # "mariadb" or "mysql"

taskfile.yaml

@@ -470,6 +470,7 @@ tasks:
       # NOTE: The trailing slash after the source is necessary so that rsync copies
       # `/parents/A` -> `/parents/B` rather than `/parents/A` -> `/parents/B/A`
       - "rsync --archive '{{.G_PACKAGE_BUILD_DIR}}/' '{{.OUTPUT_DIR}}'"
+
       # Set the storage engine for the package
       - task: "utils:misc:replace-text"
         vars:
@@ -481,6 +482,21 @@ tasks:
           FILE_PATH: "{{.OUTPUT_DIR}}/etc/clp-config.yml"
           SED_EXP: >-
             s/(\#[[:space:]]*storage_engine: ")[^"]+"/\1{{.STORAGE_ENGINE}}"/
+
+      # Set the query engine for the package
+      # NOTE: We set `query_engine` to `STORAGE_ENGINE` intentionally since compatible CLP storage
+      # and query engines have the same name.
+      - task: "utils:misc:replace-text"
+        vars:
+          FILE_PATH: "{{.OUTPUT_DIR}}/lib/python3/site-packages/clp_py_utils/clp_config.py"
+          SED_EXP: >-
+            s/([[:space:]]*query_engine: str = ")[^"]+"/\1{{.STORAGE_ENGINE}}"/
+      - task: "utils:misc:replace-text"
+        vars:
+          FILE_PATH: "{{.OUTPUT_DIR}}/etc/clp-config.yml"
+          SED_EXP: >-
+            s/(\#[[:space:]]*query_engine: ")[^"]+"/\1{{.STORAGE_ENGINE}}"/
+
       - >-
         tar czf '{{.OUTPUT_FILE}}'
         --directory '{{.G_BUILD_DIR}}'