Merge pull request #10910 from saraangelmurphy/awsworkloadidentity

feat: add AWS workload identity support
This commit is contained in:
Timothy Jaeryang Baek 2025-02-27 13:50:36 -08:00 committed by GitHub
commit 1d27402dd0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 37 additions and 12 deletions

View File

@ -101,19 +101,33 @@ class LocalStorageProvider(StorageProvider):
class S3StorageProvider(StorageProvider): class S3StorageProvider(StorageProvider):
def __init__(self): def __init__(self):
config = Config(
s3={
"use_accelerate_endpoint": S3_USE_ACCELERATE_ENDPOINT,
"addressing_style": S3_ADDRESSING_STYLE,
},
)
# If access key and secret are provided, use them for authentication
if S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY:
self.s3_client = boto3.client( self.s3_client = boto3.client(
"s3", "s3",
region_name=S3_REGION_NAME, region_name=S3_REGION_NAME,
endpoint_url=S3_ENDPOINT_URL, endpoint_url=S3_ENDPOINT_URL,
aws_access_key_id=S3_ACCESS_KEY_ID, aws_access_key_id=S3_ACCESS_KEY_ID,
aws_secret_access_key=S3_SECRET_ACCESS_KEY, aws_secret_access_key=S3_SECRET_ACCESS_KEY,
config=Config( config=config,
s3={
"use_accelerate_endpoint": S3_USE_ACCELERATE_ENDPOINT,
"addressing_style": S3_ADDRESSING_STYLE,
},
),
) )
else:
# If no explicit credentials are provided, fall back to default AWS credentials
# This supports workload identity (IAM roles for EC2, EKS, etc.)
self.s3_client = boto3.client(
"s3",
region_name=S3_REGION_NAME,
endpoint_url=S3_ENDPOINT_URL,
config=config,
)
self.bucket_name = S3_BUCKET_NAME self.bucket_name = S3_BUCKET_NAME
self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else "" self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""

View File

@ -187,6 +187,17 @@ class TestS3StorageProvider:
assert not (upload_dir / self.filename).exists() assert not (upload_dir / self.filename).exists()
assert not (upload_dir / self.filename_extra).exists() assert not (upload_dir / self.filename_extra).exists()
def test_init_without_credentials(self, monkeypatch):
"""Test that S3StorageProvider can initialize without explicit credentials."""
# Temporarily unset the environment variables
monkeypatch.setattr(provider, "S3_ACCESS_KEY_ID", None)
monkeypatch.setattr(provider, "S3_SECRET_ACCESS_KEY", None)
# Should not raise an exception
storage = provider.S3StorageProvider()
assert storage.s3_client is not None
assert storage.bucket_name == provider.S3_BUCKET_NAME
class TestGCSStorageProvider: class TestGCSStorageProvider:
Storage = provider.GCSStorageProvider() Storage = provider.GCSStorageProvider()