feat(clp-package)!: Refactor S3 log-ingestion configuration to use S3 URL as compression script argument (fixes #842). (#852)

Co-authored-by: haiqi96 <14502009+haiqi96@users.noreply.github.com>
Co-authored-by: Kirk Rodrigues <2454684+kirkrodrigues@users.noreply.github.com>

Author: 3 people

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

@@ -507,7 +507,7 @@ def validate_results_cache_config(
 def validate_worker_config(clp_config: CLPConfig):
     clp_config.validate_logs_input_config()
     clp_config.validate_archive_output_config()
-    clp_config.validate_stream_output_dir()
+    clp_config.validate_stream_output_config()
 
 
 def validate_webui_config(

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

@@ -113,7 +113,7 @@ def _validate_s3_input_args(
             f" {StorageEngine.CLP_S}."
         )
     if len(parsed_args.paths) != 1:
-        args_parser.error(f"Only one key prefix can be specified for input type {InputType.S3}.")
+        args_parser.error(f"Only one URL can be specified for input type {InputType.S3}.")
     if parsed_args.path_list is not None:
         args_parser.error(f"Path list file is unsupported for input type {InputType.S3}.")
 
@@ -141,7 +141,9 @@ def main(argv):
     args_parser.add_argument(
         "--no-progress-reporting", action="store_true", help="Disables progress reporting."
     )
-    args_parser.add_argument("paths", metavar="PATH", nargs="*", help="Paths to compress.")
+    args_parser.add_argument(
+        "paths", metavar="PATH", nargs="*", help="Paths or an S3 URL to compress."
+    )
     args_parser.add_argument(
         "-f", "--path-list", dest="path_list", help="A file listing all paths to compress."
     )

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

@@ -12,6 +12,7 @@
 import msgpack
 from clp_py_utils.clp_config import CLPConfig, COMPRESSION_JOBS_TABLE_NAME
 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
 from job_orchestration.scheduler.constants import (
     CompressionJobCompletionStatus,
@@ -136,31 +137,34 @@ def _generate_clp_io_config(
     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(
             paths_to_compress=logs_to_compress,
             timestamp_key=parsed_args.timestamp_key,
             path_prefix_to_remove=str(CONTAINER_INPUT_LOGS_ROOT_DIR),
         )
     elif InputType.S3 == input_type:
-        if len(logs_to_compress) != 1:
-            raise ValueError(f"Too many key prefixes: {len(logs_to_compress)} > 1")
-
-        s3_config = clp_config.logs_input.s3_config
+        if len(logs_to_compress) == 0:
+            raise ValueError(f"No URLs given.")
+        elif len(logs_to_compress) != 1:
+            raise ValueError(f"Too many URLs: {len(logs_to_compress)} > 1")
+
+        s3_url = logs_to_compress[0]
+        region_code, bucket_name, key_prefix = parse_s3_url(s3_url)
+        aws_authentication = clp_config.logs_input.aws_authentication
         return S3InputConfig(
-            region_code=s3_config.region_code,
-            bucket=s3_config.bucket,
-            key_prefix=s3_config.key_prefix + logs_to_compress[0],
-            aws_authentication=s3_config.aws_authentication,
+            region_code=region_code,
+            bucket=bucket_name,
+            key_prefix=key_prefix,
+            aws_authentication=aws_authentication,
             timestamp_key=parsed_args.timestamp_key,
         )
     else:
         raise ValueError(f"Unsupported input type: {input_type}")
 
 
 def _get_logs_to_compress(logs_list_path: pathlib.Path) -> List[str]:
-    # Define the path processing function based on the input type
-    process_path_func: typing.Callable[[str], str]
-
     # Read logs from the input file
     logs_to_compress = []
     with open(logs_list_path, "r") as f:

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

@@ -393,13 +393,13 @@ def validate_bucket(cls, field):
             raise ValueError("bucket cannot be empty")
         return field
 
-    @validator("key_prefix")
-    def validate_key_prefix(cls, field):
-        if "" == field:
-            raise ValueError("key_prefix cannot be empty")
-        if not field.endswith("/"):
-            raise ValueError('key_prefix must end with "/"')
-        return field
+
+class S3IngestionConfig(BaseModel):
+    type: Literal[StorageType.S3.value] = StorageType.S3.value
+    aws_authentication: AwsAuthentication
+
+    def dump_to_primitive_dict(self):
+        return self.dict()
 
 
 class FsStorage(BaseModel):
@@ -424,13 +424,6 @@ def dump_to_primitive_dict(self):
 class S3Storage(BaseModel):
     type: Literal[StorageType.S3.value] = StorageType.S3.value
     s3_config: S3Config
-
-    def dump_to_primitive_dict(self):
-        d = self.dict()
-        return d
-
-
-class OutputS3Storage(S3Storage):
     staging_directory: pathlib.Path
 
     @validator("staging_directory")
@@ -439,16 +432,28 @@ def validate_staging_directory(cls, field):
             raise ValueError("staging_directory cannot be empty")
         return field
 
+    @root_validator
+    def validate_key_prefix(cls, values):
+        s3_config = values.get("s3_config")
+        if not hasattr(s3_config, "key_prefix"):
+            raise ValueError("s3_config must have field key_prefix")
+        key_prefix = s3_config.key_prefix
+        if "" == key_prefix:
+            raise ValueError("s3_config.key_prefix cannot be empty")
+        if not key_prefix.endswith("/"):
+            raise ValueError('s3_config.key_prefix must end with "/"')
+        return values
+
     def make_config_paths_absolute(self, clp_home: pathlib.Path):
         self.staging_directory = make_config_path_absolute(clp_home, self.staging_directory)
 
     def dump_to_primitive_dict(self):
-        d = super().dump_to_primitive_dict()
+        d = self.dict()
         d["staging_directory"] = str(d["staging_directory"])
         return d
 
 
-class InputFsStorage(FsStorage):
+class FsIngestionConfig(FsStorage):
     directory: pathlib.Path = pathlib.Path("/")
 
 
@@ -460,16 +465,16 @@ class StreamFsStorage(FsStorage):
     directory: pathlib.Path = CLP_DEFAULT_DATA_DIRECTORY_PATH / "streams"
 
 
-class ArchiveS3Storage(OutputS3Storage):
+class ArchiveS3Storage(S3Storage):
     staging_directory: pathlib.Path = CLP_DEFAULT_DATA_DIRECTORY_PATH / "staged-archives"
 
 
-class StreamS3Storage(OutputS3Storage):
+class StreamS3Storage(S3Storage):
     staging_directory: pathlib.Path = CLP_DEFAULT_DATA_DIRECTORY_PATH / "staged-streams"
 
 
 def _get_directory_from_storage_config(
-    storage_config: Union[FsStorage, OutputS3Storage],
+    storage_config: Union[FsStorage, S3Storage],
 ) -> pathlib.Path:
     storage_type = storage_config.type
     if StorageType.FS == storage_type:
@@ -481,7 +486,7 @@ def _get_directory_from_storage_config(
 
 
 def _set_directory_for_storage_config(
-    storage_config: Union[FsStorage, OutputS3Storage], directory
+    storage_config: Union[FsStorage, S3Storage], directory
 ) -> None:
     storage_type = storage_config.type
     if StorageType.FS == storage_type:
@@ -603,7 +608,7 @@ def validate_port(cls, field):
 class CLPConfig(BaseModel):
     execution_container: Optional[str] = None
 
-    logs_input: Union[InputFsStorage, S3Storage] = InputFsStorage()
+    logs_input: Union[FsIngestionConfig, S3IngestionConfig] = FsIngestionConfig()
 
     package: Package = Package()
     database: Database = Database()
@@ -638,14 +643,20 @@ def make_config_paths_absolute(self, clp_home: pathlib.Path):
         self._os_release_file_path = make_config_path_absolute(clp_home, self._os_release_file_path)
 
     def validate_logs_input_config(self):
-        if StorageType.FS == self.logs_input.type:
+        logs_input_type = self.logs_input.type
+        if StorageType.FS == logs_input_type:
             # NOTE: This can't be a pydantic validator since input_logs_dir might be a
             # package-relative path that will only be resolved after pydantic validation
             input_logs_dir = self.logs_input.directory
             if not input_logs_dir.exists():
                 raise ValueError(f"logs_input.directory '{input_logs_dir}' doesn't exist.")
             if not input_logs_dir.is_dir():
                 raise ValueError(f"logs_input.directory '{input_logs_dir}' is not a directory.")
+        if StorageType.S3 == logs_input_type and StorageEngine.CLP_S != self.package.storage_engine:
+            raise ValueError(
+                f"logs_input.type = 's3' is only supported with package.storage_engine"
+                f" = '{StorageEngine.CLP_S}'"
+            )
 
     def validate_archive_output_config(self):
         if (
@@ -661,7 +672,15 @@ def validate_archive_output_config(self):
         except ValueError as ex:
             raise ValueError(f"archive_output.storage's directory is invalid: {ex}")
 
-    def validate_stream_output_dir(self):
+    def validate_stream_output_config(self):
+        if (
+            StorageType.S3 == self.stream_output.storage.type
+            and StorageEngine.CLP_S != self.package.storage_engine
+        ):
+            raise ValueError(
+                f"stream_output.storage.type = 's3' is only supported with package.storage_engine"
+                f" = '{StorageEngine.CLP_S}'"
+            )
         try:
             validate_path_could_be_dir(self.stream_output.get_directory())
         except ValueError as ex:
@@ -681,17 +700,16 @@ def validate_logs_dir(self):
 
     def validate_aws_config_dir(self):
         profile_auth_used = False
-        storage_configs = []
+        auth_configs = []
 
         if StorageType.S3 == self.logs_input.type:
-            storage_configs.append(self.logs_input)
+            auth_configs.append(self.logs_input.aws_authentication)
         if StorageType.S3 == self.archive_output.storage.type:
-            storage_configs.append(self.archive_output.storage)
+            auth_configs.append(self.archive_output.storage.s3_config.aws_authentication)
         if StorageType.S3 == self.stream_output.storage.type:
-            storage_configs.append(self.stream_output.storage)
+            auth_configs.append(self.stream_output.storage.s3_config.aws_authentication)
 
-        for storage in storage_configs:
-            auth = storage.s3_config.aws_authentication
+        for auth in auth_configs:
             if AwsAuthType.profile.value == auth.type:
                 profile_auth_used = True
                 break

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

@@ -1,4 +1,5 @@
 import os
+import re
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -7,6 +8,7 @@
 from job_orchestration.scheduler.job_config import S3InputConfig
 
 from clp_py_utils.clp_config import (
+    AwsAuthentication,
     AwsAuthType,
     CLPConfig,
     COMPRESSION_SCHEDULER_COMPONENT_NAME,
@@ -51,16 +53,16 @@ def _get_session_credentials(aws_profile: Optional[str] = None) -> Optional[S3Cr
     )
 
 
-def get_credential_env_vars(config: S3Config) -> Dict[str, str]:
+def get_credential_env_vars(auth: AwsAuthentication) -> Dict[str, str]:
     """
     Generates AWS credential environment variables for tasks.
-    :param config: S3Config or S3InputConfig from which to retrieve credentials.
-    :return: A [str, str] Dict which access key pair and session token if applicable.
-    :raise: ValueError if auth type is not supported
+    :param auth: AwsAuthentication
+    :return: A dictionary containing an access key-pair and optionally, a session token; or an empty
+    dictionary if the AWS-credential environment-variables should've been set already.
+    :raise: ValueError if `auth.type` is not a supported type or fails to authenticate with the
+    given `auth`.
     """
     env_vars: Optional[Dict[str, str]] = None
-    auth = config.aws_authentication
-
     aws_credentials: Optional[S3Credentials] = None
 
     if AwsAuthType.env_vars == auth.type:
@@ -104,20 +106,22 @@ def generate_container_auth_options(
     :return: Tuple of (whether aws config mount is needed, credential env_vars to set).
     :raises: ValueError if environment variables are not set correctly.
     """
-    storages_by_component_type: List[Union[S3Storage, FsStorage]] = []
+    output_storages_by_component_type: List[Union[S3Storage, FsStorage]] = []
+    input_storage_needed = False
+
     if component_type in (
         COMPRESSION_SCHEDULER_COMPONENT_NAME,
         COMPRESSION_WORKER_COMPONENT_NAME,
     ):
-        storages_by_component_type = [clp_config.logs_input, clp_config.archive_output.storage]
+        output_storages_by_component_type = [clp_config.archive_output.storage]
+        input_storage_needed = True
     elif component_type in (LOG_VIEWER_WEBUI_COMPONENT_NAME,):
-        storages_by_component_type = [clp_config.stream_output.storage]
+        output_storages_by_component_type = [clp_config.stream_output.storage]
     elif component_type in (
         QUERY_SCHEDULER_COMPONENT_NAME,
         QUERY_WORKER_COMPONENT_NAME,
     ):
-        storages_by_component_type = [
-            clp_config.logs_input,
+        output_storages_by_component_type = [
             clp_config.archive_output.storage,
             clp_config.stream_output.storage,
         ]
@@ -126,14 +130,21 @@ def generate_container_auth_options(
     config_mount = False
     add_env_vars = False
 
-    for storage in storages_by_component_type:
+    for storage in output_storages_by_component_type:
         if StorageType.S3 == storage.type:
             auth = storage.s3_config.aws_authentication
             if AwsAuthType.profile == auth.type:
                 config_mount = True
             elif AwsAuthType.env_vars == auth.type:
                 add_env_vars = True
 
+    if input_storage_needed and StorageType.S3 == clp_config.logs_input.type:
+        auth = clp_config.logs_input.aws_authentication
+        if AwsAuthType.profile == auth.type:
+            config_mount = True
+        elif AwsAuthType.env_vars == auth.type:
+            add_env_vars = True
+
     credentials_env_vars = []
 
     if add_env_vars:
@@ -186,6 +197,41 @@ def _create_s3_client(s3_config: S3Config, boto3_config: Optional[Config] = None
     return s3_client
 
 
+def parse_s3_url(s3_url: str) -> Tuple[str, str, str]:
+    """
+    Parses the region_code, bucket, and key_prefix from the given S3 URL.
+    :param s3_url: A host-style URL or path-style URL.
+    :return: A tuple of (region_code, bucket, key_prefix).
+    :raise: ValueError if `s3_url` is not a valid host-style URL or path-style URL.
+    """
+
+    host_style_url_regex = re.compile(
+        r"https://(?P<bucket_name>[a-z0-9.-]+)\.s3(\.(?P<region_code>[a-z]+-[a-z]+-[0-9]))"
+        r"\.(?P<endpoint>[a-z0-9.-]+)/(?P<key_prefix>[^?]+).*"
+    )
+    match = host_style_url_regex.match(s3_url)
+
+    if match is None:
+        path_style_url_regex = re.compile(
+            r"https://s3(\.(?P<region_code>[a-z]+-[a-z]+-[0-9]))\.(?P<endpoint>[a-z0-9.-]+)/"
+            r"(?P<bucket_name>[a-z0-9.-]+)/(?P<key_prefix>[^?]+).*"
+        )
+        match = path_style_url_regex.match(s3_url)
+
+    if match is None:
+        raise ValueError(f"Unsupported URL format: {s3_url}")
+
+    region_code = match.group("region_code")
+    bucket_name = match.group("bucket_name")
+    endpoint = match.group("endpoint")
+    key_prefix = match.group("key_prefix")
+
+    if AWS_ENDPOINT != endpoint:
+        raise ValueError(f"Unsupported endpoint: {endpoint}")
+
+    return region_code, bucket_name, key_prefix
+
+
 def generate_s3_virtual_hosted_style_url(
     region_code: str, bucket_name: str, object_key: str
 ) -> str:

components/job-orchestration/job_orchestration/executor/compress/compression_task.py

@@ -208,7 +208,7 @@ def _make_clp_s_command_and_env(
 
     if InputType.S3 == clp_config.input.type:
         compression_env_vars = dict(os.environ)
-        compression_env_vars.update(get_credential_env_vars(clp_config.input))
+        compression_env_vars.update(get_credential_env_vars(clp_config.input.aws_authentication))
         compression_cmd.append("--auth")
         compression_cmd.append("s3")
     else:

components/job-orchestration/job_orchestration/executor/query/extract_stream_task.py

@@ -112,7 +112,7 @@ def _make_clp_s_command_and_env_vars(
         ))
         # fmt: on
         env_vars = dict(os.environ)
-        env_vars.update(get_credential_env_vars(s3_config))
+        env_vars.update(get_credential_env_vars(s3_config.aws_authentication))
     else:
         # fmt: off
         command.extend((