add assume role arn to sts

duckdb · Feb 24, 2025 · 487ff71 · 487ff71
1 parent b3050f3
commit 487ff71
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 6 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -16,7 +16,7 @@ build_loadable_extension(${TARGET_NAME} ${PARAMETERS} ${EXTENSION_SOURCES})
 # Weirdly we need to manually to this, otherwise linking against
 # ${AWSSDK_LINK_LIBRARIES} fails for some reason
 find_package(ZLIB REQUIRED)
-find_package(AWSSDK REQUIRED COMPONENTS core sso sts)
+find_package(AWSSDK REQUIRED COMPONENTS core sso sts identity-management)
 
 # Build static lib
 target_include_directories(${EXTENSION_NAME}

diff --git a/src/aws_extension.cpp b/src/aws_extension.cpp
@@ -120,7 +120,6 @@ static void LoadAWSCredentialsFun(ClientContext &context, TableFunctionInput &da
 
 	data.finished = true;
 }
-
 static void LoadInternal(DuckDB &db) {
 	TableFunctionSet function_set("load_aws_credentials");
 	auto base_fun = TableFunction("load_aws_credentials", {}, LoadAWSCredentialsFun, LoadAWSCredentialsBind);

diff --git a/src/aws_secret.cpp b/src/aws_secret.cpp
@@ -8,6 +8,8 @@
 #include <aws/core/auth/SSOCredentialsProvider.h>
 #include <aws/core/auth/STSCredentialsProvider.h>
 #include <aws/core/client/ClientConfiguration.h>
+#include <aws/identity-management/auth/STSAssumeRoleCredentialsProvider.h>
+#include <aws/identity-management/auth/STSProfileCredentialsProvider.h>
 
 namespace duckdb {
 
@@ -35,12 +37,19 @@ static unique_ptr<KeyValueSecret> ConstructBaseS3Secret(vector<string> &prefix_p
 //! Generate a custom credential provider chain for authentication
 class DuckDBCustomAWSCredentialsProviderChain : public Aws::Auth::AWSCredentialsProviderChain {
 public:
-	explicit DuckDBCustomAWSCredentialsProviderChain(const string &credential_chain, const string &profile = "") {
+	explicit DuckDBCustomAWSCredentialsProviderChain(const string &credential_chain, const string &profile = "", const string &assume_role_arn = "") {
 		auto chain_list = StringUtil::Split(credential_chain, ';');
 
 		for (const auto &item : chain_list) {
 			if (item == "sts") {
-				AddProvider(std::make_shared<Aws::Auth::STSAssumeRoleWebIdentityCredentialsProvider>());
+				if (!profile.empty()) {
+					AddProvider(std::make_shared<Aws::Auth::STSProfileCredentialsProvider>(profile));
+				} else if (!assume_role_arn.empty()) {
+					AddProvider(std::make_shared<Aws::Auth::STSAssumeRoleCredentialsProvider>(assume_role_arn));
+				} else {
+					// TODO: I don't think this does anything
+					AddProvider(std::make_shared<Aws::Auth::STSAssumeRoleWebIdentityCredentialsProvider>());
+				}
 			} else if (item == "sso") {
 				if (profile.empty()) {
 					AddProvider(std::make_shared<Aws::Auth::SSOCredentialsProvider>());
@@ -83,11 +92,12 @@ static unique_ptr<BaseSecret> CreateAWSSecretFromCredentialChain(ClientContext &
 	Aws::Auth::AWSCredentials credentials;
 
 	string profile = TryGetStringParam(input, "profile");
+	string assume_role = TryGetStringParam(input, "assume_role_arn");
 
 	if (input.options.find("chain") != input.options.end()) {
 		string chain = TryGetStringParam(input, "chain");
 
-		DuckDBCustomAWSCredentialsProviderChain provider(chain, profile);
+		DuckDBCustomAWSCredentialsProviderChain provider(chain, profile, assume_role);
 		credentials = provider.GetAWSCredentials();
 	} else {
 		if (input.options.find("profile") != input.options.end()) {
@@ -184,6 +194,10 @@ void CreateAwsSecretFunctions::Register(DatabaseInstance &instance) {
 		cred_chain_function.named_parameters["use_ssl"] = LogicalType::BOOLEAN;
 		cred_chain_function.named_parameters["url_compatibility_mode"] = LogicalType::BOOLEAN;
 
+		cred_chain_function.named_parameters["assume_role_arn"] = LogicalType::VARCHAR;
+
+		cred_chain_function.named_parameters["refresh"] = LogicalType::BOOLEAN;
+
 		if (type == "r2") {
 			cred_chain_function.named_parameters["account_id"] = LogicalType::VARCHAR;
 		}

diff --git a/vcpkg.json b/vcpkg.json
@@ -3,7 +3,7 @@
     "zlib",
     {
       "name": "aws-sdk-cpp",
-      "features": [ "sso", "sts" ]
+      "features": [ "sso", "sts" , "identity-management"]
     },
     "openssl"
   ]