Update config (it's so small so I don't bother to make it a separate PR)

haiqi96 · haiqi96 · commit c6a1ad2dd84f · 2025-01-16T15:44:16.000-05:00
diff --git a/components/clp-py-utils/clp_py_utils/clp_config.py b/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, Union, Tuple
 
 from dotenv import dotenv_values
 from pydantic import BaseModel, PrivateAttr, validator
@@ -310,13 +310,29 @@ class Queue(BaseModel):
     password: Optional[str]
 
 
+class S3Credentials(BaseModel):
+    access_key_id: str
+    secret_access_key: str
+
+    @validator("access_key_id")
+    def validate_access_key_id(cls, field):
+        if field == "":
+            raise ValueError("access_key_id cannot be empty")
+        return field
+
+    @validator("secret_access_key")
+    def validate_secret_access_key(cls, field):
+        if field == "":
+            raise ValueError("secret_access_key cannot be empty")
+        return field
+
+
 class S3Config(BaseModel):
     region_code: str
     bucket: str
     key_prefix: str
 
-    access_key_id: Optional[str] = None
-    secret_access_key: Optional[str] = None
+    credentials: Optional[S3Credentials]
 
     @validator("region_code")
     def validate_region_code(cls, field):
@@ -338,6 +354,11 @@ def validate_key_prefix(cls, field):
             raise ValueError('key_prefix must end with "/"')
         return field
 
+    def get_credentials(self) -> Tuple[Optional[str], Optional[str]]:
+        if self.credentials is None:
+            return None, None
+        return self.credentials.access_key_id, self.credentials.secret_access_key
+
 
 class FsStorage(BaseModel):
     type: Literal[StorageType.FS.value] = StorageType.FS.value
diff --git a/components/clp-py-utils/clp_py_utils/s3_utils.py b/components/clp-py-utils/clp_py_utils/s3_utils.py
@@ -161,12 +161,13 @@ def s3_put(
         )
 
     config = Config(retries=dict(total_max_attempts=total_max_attempts, mode="adaptive"))
+    aws_access_key_id, aws_secret_access_key = s3_config.get_credentials()
 
     my_s3_client = boto3.client(
         "s3",
         region_name=s3_config.region_code,
-        aws_access_key_id=s3_config.access_key_id,
-        aws_secret_access_key=s3_config.secret_access_key,
+        aws_access_key_id=aws_access_key_id,
+        aws_secret_access_key=aws_secret_access_key,
         config=config,
     )
 

Original file line number	Diff line number	Diff line change
`@@ -161,12 +161,13 @@ def s3_put(`
`161`	`161`	`)`
`162`	`162`
`163`	`163`	`config = Config(retries=dict(total_max_attempts=total_max_attempts, mode="adaptive"))`
	`164`	`+ aws_access_key_id, aws_secret_access_key = s3_config.get_credentials()`
`164`	`165`
`165`	`166`	`my_s3_client = boto3.client(`
`166`	`167`	`"s3",`
`167`	`168`	`region_name=s3_config.region_code,`
`168`		`- aws_access_key_id=s3_config.access_key_id,`
`169`		`- aws_secret_access_key=s3_config.secret_access_key,`
	`169`	`+ aws_access_key_id=aws_access_key_id,`
	`170`	`+ aws_secret_access_key=aws_secret_access_key,`
`170`	`171`	`config=config,`
`171`	`172`	`)`
`172`	`173`