From 95d8d6dc887c9327b9b5eaf3eea69e7fb7ac231f Mon Sep 17 00:00:00 2001
From: Zoltan Mezei <59565282+zizzencs@users.noreply.github.com>
Date: Thu, 7 Nov 2024 09:09:54 +0100
Subject: [PATCH] feat: Add support for assuming an AWS IAM role from the
 provider. (#486)

This pull request proposes a solution for
https://github.com/cyrilgdn/terraform-provider-postgresql/issues/263

---------

Co-authored-by: Cyril Gaudin <cyril.gaudin@gmail.com>
---
 postgresql/provider.go           | 40 ++++++++++++++++++++++++++++++--
 website/docs/index.html.markdown |  1 +
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/postgresql/provider.go b/postgresql/provider.go
index af2b435f..f0df3ea9 100644
--- a/postgresql/provider.go
+++ b/postgresql/provider.go
@@ -3,6 +3,8 @@ package postgresql
 import (
 	"context"
 	"fmt"
+	"github.com/aws/aws-sdk-go-v2/credentials"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"os"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
@@ -90,6 +92,13 @@ func Provider() *schema.Provider {
 				Description: "AWS region to use for IAM auth",
 			},
 
+			"aws_rds_iam_provider_role_arn": {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Default:     "",
+				Description: "AWS IAM role to assume for IAM auth",
+			},
+
 			"azure_identity_auth": {
 				Type:     schema.TypeBool,
 				Optional: true,
@@ -227,7 +236,7 @@ func validateExpectedVersion(v interface{}, key string) (warnings []string, erro
 	return
 }
 
-func getRDSAuthToken(region string, profile string, username string, host string, port int) (string, error) {
+func getRDSAuthToken(region string, profile string, role string, username string, host string, port int) (string, error) {
 	endpoint := fmt.Sprintf("%s:%d", host, port)
 
 	ctx := context.Background()
@@ -246,6 +255,32 @@ func getRDSAuthToken(region string, profile string, username string, host string
 		return "", err
 	}
 
+	if role != "" {
+		stsClient := sts.NewFromConfig(awscfg)
+		roleInput := &sts.AssumeRoleInput{
+			RoleArn:         aws.String(role),
+			RoleSessionName: aws.String("TerraformPostgresqlProvider"),
+		}
+
+		roleOutput, err := stsClient.AssumeRole(ctx, roleInput)
+		if err != nil {
+			return "", fmt.Errorf("could not assume AWS role: %w", err)
+		}
+
+		awscfg, err = awsConfig.LoadDefaultConfig(ctx,
+			awsConfig.WithCredentialsProvider(
+				aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(
+					*roleOutput.Credentials.AccessKeyId,
+					*roleOutput.Credentials.SecretAccessKey,
+					*roleOutput.Credentials.SessionToken,
+				)),
+			),
+		)
+		if err != nil {
+			return "", fmt.Errorf("could not load AWS default config: %w", err)
+		}
+	}
+
 	token, err := auth.BuildAuthToken(ctx, endpoint, awscfg.Region, username, awscfg.Credentials)
 
 	return token, err
@@ -312,8 +347,9 @@ func providerConfigure(d *schema.ResourceData) (interface{}, error) {
 	if d.Get("aws_rds_iam_auth").(bool) {
 		profile := d.Get("aws_rds_iam_profile").(string)
 		region := d.Get("aws_rds_iam_region").(string)
+		role := d.Get("aws_rds_iam_provider_role_arn").(string)
 		var err error
-		password, err = getRDSAuthToken(region, profile, username, host, port)
+		password, err = getRDSAuthToken(region, profile, role, username, host, port)
 		if err != nil {
 			return nil, err
 		}
diff --git a/website/docs/index.html.markdown b/website/docs/index.html.markdown
index 30959d90..ce178fbc 100644
--- a/website/docs/index.html.markdown
+++ b/website/docs/index.html.markdown
@@ -183,6 +183,7 @@ The following arguments are supported:
   from the environment (or the given profile, see `aws_rds_iam_profile`)
 * `aws_rds_iam_profile` - (Optional) The AWS IAM Profile to use while using AWS RDS IAM Auth.
 * `aws_rds_iam_region` - (Optional) The AWS region to use while using AWS RDS IAM Auth.
+* `aws_rds_iam_provider_role_arn` - (Optional) AWS IAM role to assume while using AWS RDS IAM Auth.
 * `azure_identity_auth` - (Optional) If set to `true`, call the Azure OAuth token endpoint for temporary token
 * `azure_tenant_id` - (Optional) (Required if `azure_identity_auth` is `true`) Azure tenant ID [read more](https://registry.terraform.io/providers/hashicorp/azurerm/latest/docs/data-sources/client_config.html)