Replace S3InputConfig.input.aws_* with S3Credentials

y-scope · Jan 21, 2025 · 8ab90df · 8ab90df
1 parent ee3839c
commit 8ab90df
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/components/clp-package-utils/clp_package_utils/scripts/native/compress.py b/components/clp-package-utils/clp_package_utils/scripts/native/compress.py
@@ -10,7 +10,7 @@
 
 import brotli
 import msgpack
-from clp_py_utils.clp_config import COMPRESSION_JOBS_TABLE_NAME
+from clp_py_utils.clp_config import COMPRESSION_JOBS_TABLE_NAME, S3Credentials
 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
@@ -148,8 +148,10 @@ def _generate_clp_io_config(
             region_code=region_code,
             bucket=bucket_name,
             key_prefix=key_prefix,
-            aws_access_key_id=parsed_args.aws_access_key_id,
-            aws_secret_access_key=parsed_args.aws_secret_access_key,
+            credentials=S3Credentials(
+                access_key_id=parsed_args.aws_access_key_id,
+                secret_access_key=parsed_args.aws_secret_access_key,
+            ),
             timestamp_key=parsed_args.timestamp_key,
         )
     else:

diff --git a/components/clp-py-utils/clp_py_utils/s3_utils.py b/components/clp-py-utils/clp_py_utils/s3_utils.py
@@ -77,8 +77,8 @@ def s3_get_object_metadata(s3_input_config: S3InputConfig) -> List[FileMetadata]
     s3_client = boto3.client(
         "s3",
         region_name=s3_input_config.region_code,
-        aws_access_key_id=s3_input_config.aws_access_key_id,
-        aws_secret_access_key=s3_input_config.aws_secret_access_key,
+        aws_access_key_id=s3_input_config.credentials.access_key_id,
+        aws_secret_access_key=s3_input_config.credentials.secret_access_key,
     )
 
     file_metadata_list: List[FileMetadata] = list()

diff --git a/components/job-orchestration/job_orchestration/executor/compress/compression_task.py b/components/job-orchestration/job_orchestration/executor/compress/compression_task.py
@@ -187,8 +187,8 @@ def make_clp_s_command_and_env(
     if InputType.S3 == clp_config.input.type:
         compression_env_vars = {
             **os.environ,
-            "AWS_ACCESS_KEY_ID": clp_config.input.aws_access_key_id,
-            "AWS_SECRET_ACCESS_KEY": clp_config.input.aws_secret_access_key,
+            "AWS_ACCESS_KEY_ID": clp_config.input.credentials.access_key_id,
+            "AWS_SECRET_ACCESS_KEY": clp_config.input.credentials.secret_access_key,
         }
         compression_cmd.append("--auth")
         compression_cmd.append("s3")

diff --git a/components/job-orchestration/job_orchestration/scheduler/job_config.py b/components/job-orchestration/job_orchestration/scheduler/job_config.py
@@ -3,6 +3,7 @@
 import typing
 from enum import auto
 
+from clp_py_utils.clp_config import S3Credentials
 from pydantic import BaseModel, validator
 from strenum import LowercaseStrEnum
 
@@ -34,8 +35,7 @@ class S3InputConfig(BaseModel):
     bucket: str
     key_prefix: str
 
-    aws_access_key_id: str
-    aws_secret_access_key: str
+    credentials: S3Credentials
 
 
 class OutputConfig(BaseModel):