diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 10290fc3c1..028d6f590c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -215,6 +215,8 @@ class DynamoDBArgs(LiteLLMBase): aws_policy_arns: Optional[List[str]] = None aws_policy: Optional[str] = None aws_duration_seconds: Optional[int] = None + assume_role_aws_role_name: Optional[str] = None + assume_role_aws_session_name: Optional[str] = None class ConfigGeneralSettings(LiteLLMBase): diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index f315a74be5..78e81c1542 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -61,12 +61,20 @@ class DynamoDBWrapper(CustomDB): import boto3, os sts_client = boto3.client("sts") - assumed_role = sts_client.assume_role_with_web_identity( + + # call 1 + non_used_assumed_role = sts_client.assume_role_with_web_identity( RoleArn=self.database_arguments.aws_role_name, RoleSessionName=self.database_arguments.aws_session_name, WebIdentityToken=self.database_arguments.aws_web_identity_token, ) + # call 2 + assumed_role = sts_client.assume_role( + RoleArn=self.database_arguments.assume_role_aws_role_name, + RoleSessionName=self.database_arguments.assume_role_aws_session_name, + ) + aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"] aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"] aws_session_token = assumed_role["Credentials"]["SessionToken"]