From c4038e38e5d0116a1df610a968f3e3b52685b62a Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 12 Feb 2024 19:54:39 -0800 Subject: [PATCH] (fix) dyanmo fixes --- litellm/proxy/_types.py | 2 ++ litellm/proxy/db/dynamo_db.py | 10 +++++++++- 2 files changed, 11 insertions(+), 1 deletion(-) 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"]