feat(redshift): add support for assuming an iam role for redshift (#132)

rudderlabs · Jul 22, 2024 · 180a706 · 180a706
1 parent ec10cc4
commit 180a706
Show file tree

Hide file tree

Showing 9 changed files with 397 additions and 276 deletions.
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
@@ -55,6 +55,7 @@ jobs:
         env:
           REDSHIFT_TEST_ENVIRONMENT_CREDENTIALS: ${{ secrets.REDSHIFT_TEST_ENVIRONMENT_CREDENTIALS }}
           REDSHIFT_DATA_TEST_ENVIRONMENT_CREDENTIALS: ${{ secrets.REDSHIFT_DATA_TEST_ENVIRONMENT_CREDENTIALS }}
+          REDSHIFT_DATA_TEST_ENVIRONMENT_ROLE_ARN_CREDENTIALS: ${{ secrets.REDSHIFT_DATA_TEST_ENVIRONMENT_ROLE_ARN_CREDENTIALS }}
           SNOWFLAKE_TEST_ENVIRONMENT_CREDENTIALS: ${{ secrets.SNOWFLAKE_TEST_ENVIRONMENT_CREDENTIALS }}
           BIGQUERY_TEST_ENVIRONMENT_CREDENTIALS: ${{ secrets.BIGQUERY_TEST_ENVIRONMENT_CREDENTIALS }}
           DATABRICKS_TEST_ENVIRONMENT_CREDENTIALS: ${{ secrets.DATABRICKS_TEST_ENVIRONMENT_CREDENTIALS }}

diff --git a/go.mod b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/config v1.27.27
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.27
 	github.com/aws/aws-sdk-go-v2/service/redshiftdata v1.27.3
+	github.com/aws/aws-sdk-go-v2/service/sts v1.30.3
 	github.com/databricks/databricks-sql-go v1.5.7
 	github.com/dlclark/regexp2 v1.11.2
 	github.com/gliderlabs/ssh v0.3.7
@@ -66,7 +67,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.53.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 // indirect
 	github.com/aws/smithy-go v1.20.3 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/containerd/continuity v0.4.3 // indirect

diff --git a/sqlconnect/internal/redshift/config.go b/sqlconnect/internal/redshift/config.go
@@ -30,6 +30,10 @@ type Config struct {
 	SecretAccessKey string `json:"secretAccessKey"`
 	SessionToken    string `json:"sessionToken"`
 
+	RoleARN       string        `json:"roleARN"`
+	ExternalID    string        `json:"externalID"`
+	RoleARNExpiry time.Duration `json:"roleARNExpiry"` // default: 15m
+
 	Timeout          time.Duration `json:"timeout"`          // default: no timeout
 	MinPolling       time.Duration `json:"minPolling"`       // default: 10ms
 	MaxPolling       time.Duration `json:"maxPolling"`       // default: 5s

diff --git a/sqlconnect/internal/redshift/db.go b/sqlconnect/internal/redshift/db.go
@@ -123,6 +123,9 @@ func newRedshiftDataDB(credentialsJSON json.RawMessage) (*sql.DB, error) {
 		SharedConfigProfile: config.SharedConfigProfile,
 		SecretAccessKey:     config.SecretAccessKey,
 		SessionToken:        config.SessionToken,
+		RoleARN:             config.RoleARN,
+		ExternalID:          config.ExternalID,
+		RoleARNExpiry:       config.RoleARNExpiry,
 		Timeout:             config.Timeout,
 		MinPolling:          config.MinPolling,
 		MaxPolling:          config.MaxPolling,

diff --git a/sqlconnect/internal/redshift/driver/connector.go b/sqlconnect/internal/redshift/driver/connector.go
@@ -19,7 +19,11 @@ type redshiftDataConnector struct {
 }
 
 func (c *redshiftDataConnector) Connect(ctx context.Context) (driver.Conn, error) {
-	client, err := newRedshiftDataClient(ctx, c.cfg, c.cfg.LoadOpts()...)
+	loadOpts, err := c.cfg.LoadOpts(ctx)
+	if err != nil {
+		return nil, err
+	}
+	client, err := newRedshiftDataClient(ctx, c.cfg, loadOpts...)
 	if err != nil {
 		return nil, err
 	}