Add AWS creds config to object_store (#98)

Signed-off-by: Rajvaibhav Rahane <rrahane@amazon.com>

Author: Rajrahane

e2e/api/vector_dataset_generator.py

@@ -12,6 +12,7 @@
 from botocore.exceptions import ClientError
 from core.common.models import IndexBuildParameters
 from core.object_store.object_store_factory import ObjectStoreFactory
+from core.object_store.s3.s3_object_store_config import S3ClientConfig
 from core.object_store.types import ObjectStoreType
 import logging
 from tqdm import tqdm
@@ -38,11 +39,11 @@ def initialize_object_store(self):
             doc_count=5,  # Will be set per dataset
         )
         object_store_config = {
-            "retries": s3_config["retries"],
-            "region": s3_config["region"],
-            "S3_ENDPOINT_URL": os.environ.get(
-                "S3_ENDPOINT_URL", "http://localhost:4566"
-            ),
+            "s3_client_config": S3ClientConfig(
+                region_name=s3_config["region"],
+                max_retries=s3_config["retries"],
+                endpoint_url=os.environ.get("S3_ENDPOINT_URL", "http://localhost:4566"),
+            )
         }
         return ObjectStoreFactory.create_object_store(
             index_build_params, object_store_config

remote_vector_index_builder/app/services/index_builder.py

@@ -9,6 +9,7 @@
 import os
 from typing import Optional, Tuple
 from app.models.workflow import BuildWorkflow
+from core.object_store.s3.s3_object_store_config import S3ClientConfig
 from core.tasks import run_tasks
 
 logger = logging.getLogger(__name__)
@@ -36,7 +37,13 @@ def build_index(
         """
         s3_endpoint_url = os.environ.get("S3_ENDPOINT_URL", None)
         result = run_tasks(
-            workflow.index_build_parameters, {"S3_ENDPOINT_URL": s3_endpoint_url}
+            workflow.index_build_parameters,
+            {
+                "s3_client_config": S3ClientConfig(
+                    region_name=os.environ.get("AWS_DEFAULT_REGION", None),
+                    endpoint_url=s3_endpoint_url,
+                ),
+            },
         )
         if not result.file_name:
             return False, None, result.error

remote_vector_index_builder/core/object_store/s3/s3_object_store.py

@@ -12,7 +12,7 @@
 import threading
 from functools import cache
 from io import BytesIO
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 import boto3
 from boto3.s3.transfer import TransferConfig
@@ -21,6 +21,7 @@
 from core.common.exceptions import BlobError
 from core.common.models import IndexBuildParameters
 from core.object_store.object_store import ObjectStore
+from core.object_store.s3.s3_object_store_config import S3ClientConfig
 
 logger = logging.getLogger(__name__)
 
@@ -46,28 +47,24 @@ def get_cpus(factor: float) -> int:
 
 
 @cache
-def get_boto3_client(
-    region: str, retries: int, endpoint_url: Optional[str] = None
-) -> boto3.client:
+def get_boto3_client(s3_client_config: S3ClientConfig) -> boto3.client:
     """Create or retrieve a cached boto3 S3 client.
 
     Args:
-        region (str): AWS region name for the S3 client
-        retries (int): Maximum number of retry attempts for failed requests
-        endpoint_url (str): s3 endpoint URL. Defaults to None, in which case boto3
-            automatically constructs the appropriate URL to use when communicating
-            with a service. During integration testing, this can be set to the endpoint URL
-            for LocalStack S3 service.
+        s3_client_config (S3ClientConfig): Configuration class for creating S3 Boto3 client
 
     Returns:
         boto3.client: Configured S3 client instance
     """
-    config = Config(retries={"max_attempts": retries})
+    config = Config(retries={"max_attempts": s3_client_config.max_retries})
     return boto3.client(
         "s3",
         config=config,
-        region_name=region,
-        endpoint_url=endpoint_url,
+        region_name=s3_client_config.region_name,
+        endpoint_url=s3_client_config.endpoint_url,
+        aws_access_key_id=s3_client_config.aws_access_key_id,
+        aws_secret_access_key=s3_client_config.aws_secret_access_key,
+        aws_session_token=s3_client_config.aws_session_token,
     )
 
 
@@ -102,11 +99,27 @@ def __init__(
         Args:
             index_build_params (IndexBuildParameters): Contains bucket name and other
                 index building parameters
+
             object_store_config (Dict[str, Any]): Configuration dictionary containing:
-                - retries (int): Maximum number of retry attempts (default: 3)
-                - region (str): AWS region name (default: 'us-west-2')
+            Contains:
                 - transfer_config (Dict[str, Any]): s3 TransferConfig parameters
                 - debug: Turns on debug mode (default: False)
+                - s3_client_config (S3ClientConfig) (Required):
+                    Required:
+                        - region_name (str) (required): AWS Region name
+                    Optional:
+                        - endpoint_url (Optional[str]): Custom S3 endpoint URL
+                        - max_retries (int) (default: 3): Maximum number of retry attempts for failed requests
+
+                        AWS Credentials (all optional):
+                            - aws_access_key_id (Optional[str]): AWS Access Key ID
+                            - aws_secret_access_key (Optional[str]): AWS Secret Access Key
+                            - aws_session_token (Optional[str]): Temporary session token for STS credentials
+
+        Note:
+            AWS credentials are optional as boto3 will attempt to find credentials:
+            For more details see boto3 client documentation:
+            https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
         """
 
         self.DEFAULT_DOWNLOAD_TRANSFER_CONFIG = {
@@ -130,14 +143,12 @@ def __init__(
             "ChecksumAlgorithm": "CRC32",
         }
         self.bucket = index_build_params.container_name
-        self.max_retries = object_store_config.get("retries", 3)
-        self.region = object_store_config.get("region", "us-west-2")
 
-        self.s3_client = get_boto3_client(
-            region=self.region,
-            retries=self.max_retries,
-            endpoint_url=object_store_config.get("S3_ENDPOINT_URL"),
-        )
+        s3_client_config: S3ClientConfig = object_store_config.get("s3_client_config")
+        self.max_retries = s3_client_config.max_retries
+        self.region = s3_client_config.region_name
+
+        self.s3_client = get_boto3_client(s3_client_config)
 
         download_transfer_config = object_store_config.get(
             "download_transfer_config", {}

remote_vector_index_builder/core/object_store/s3/s3_object_store_config.py

@@ -0,0 +1,87 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class S3ClientConfig(BaseModel):
+    """
+    Configuration class for creating S3 Boto3 client
+
+    Initialize the S3ObjectStore boto3 client with the following parameters:
+
+    Attributes:
+        region_name (str) (required): AWS Region name
+        endpoint_url (Optional[str]): Custom S3 endpoint URL
+        max_retries (int) (default: 3): Maximum number of retry attempts for failed requests
+
+        AWS Credentials (all optional):
+        aws_access_key_id (Optional[str]): AWS Access Key ID
+        aws_secret_access_key (Optional[str]): AWS Secret Access Key
+        aws_session_token (Optional[str]): Temporary session token for STS credentials
+
+    Note:
+        AWS credentials are optional as boto3 will attempt to find credentials:
+        For more details see boto3 client documentation:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+    """
+
+    region_name: str
+    endpoint_url: Optional[str] = None
+    max_retries: int = 3
+
+    # AWS Credentials parameters
+    # Ref: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+    aws_access_key_id: Optional[str] = None
+    aws_secret_access_key: Optional[str] = None
+    aws_session_token: Optional[str] = None
+
+    def __hash__(self):
+        """
+        Generate a hash value for this configuration.
+
+        Required for @functools.cache to use this class as a dictionary key.
+        All attributes that affect the client configuration must be included
+        in the hash tuple to ensure proper cache behavior.
+
+        Returns:
+            int: Hash value based on all configuration attributes
+        """
+        return hash(
+            (
+                self.region_name,
+                self.max_retries,
+                self.endpoint_url,
+                self.aws_access_key_id,
+                self.aws_secret_access_key,
+                self.aws_session_token,
+            )
+        )
+
+    def __eq__(self, other):
+        """
+        Compare this configuration with another for equality.
+
+        Required for @functools.cache to properly identify cache hits.
+        Two configurations are equal if all their attributes are equal.
+
+        Args:
+            other: Another object to compare with this configuration
+
+        Returns:
+            bool: True if other is an S3ClientConfig with identical attributes,
+                  False otherwise
+
+        Note:
+            Returns NotImplemented for non-S3ClientConfig objects to allow
+            Python to try other comparison methods.
+        """
+        if not isinstance(other, S3ClientConfig):
+            return NotImplemented
+        return (
+            self.region_name == other.region_name
+            and self.max_retries == other.max_retries
+            and self.endpoint_url == other.endpoint_url
+            and self.aws_access_key_id == other.aws_access_key_id
+            and self.aws_secret_access_key == other.aws_secret_access_key
+            and self.aws_session_token == other.aws_session_token
+        )

test_remote_vector_index_builder/conftest.py

@@ -6,6 +6,7 @@
 # compatible open source license.
 
 import pytest
+import os
 from core.common.models.index_build_parameters import (
     AlgorithmParameters,
     IndexBuildParameters,
@@ -31,3 +32,11 @@ def index_build_parameters():
         ),
         repository_type="s3",
     )
+
+
+@pytest.fixture(autouse=True)
+def aws_credentials():
+    """Mocked AWS Credentials for tests."""
+    os.environ["AWS_DEFAULT_REGION"] = "us-west-2"
+    yield
+    os.environ.pop("AWS_DEFAULT_REGION", None)

test_remote_vector_index_builder/test_core/conftest.py

@@ -12,6 +12,7 @@
 from unittest.mock import Mock
 
 from core.common.models import VectorsDataset
+from core.object_store.s3.s3_object_store_config import S3ClientConfig
 
 
 class DeletionTracker:
@@ -197,7 +198,7 @@ def _write_index(self, index, filepath):
 def object_store_config():
     """Create a sample object store configuration for testing"""
     return {
-        "region": "us-west-2",
+        "s3_client_config": S3ClientConfig(region_name="us-west-2", max_retries="4")
     }