DBI-609: Add IAM auth support for CDC pipes with AWS Auruora' (#4495)

Author: masterashu

docker-compose-dev.yml

@@ -47,6 +47,8 @@ x-flow-worker-env: &flow-worker-env
   # https://cloud.google.com/storage/docs/authentication/managing-hmackeys
   AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-}
   AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY:-}
+  # Temporary AWS Credentials for local IAM auth testing
+  AWS_SESSION_TOKEN: ${AWS_SESSION_TOKEN:-}
   # For GCS, set this to "auto" without the quotes
   AWS_REGION: ${AWS_REGION:-}
   # For GCS, set this as: https://storage.googleapis.com

flow/connectors/postgres/postgres_source.go

@@ -114,6 +114,22 @@ func (c *PostgresConnector) CreateReplConn(ctx context.Context, env map[string]s
 }
 
 func (c *PostgresConnector) SetupReplConn(ctx context.Context, env map[string]string) error {
+	// For IAM-authenticated connections, force a fresh token before connecting.
+	// This guarantees each activity retry uses a new token rather than a potentially
+	// stale cached one (cache TTL 10 min; token validity 15 min).
+	if c.rdsAuth != nil {
+		host := c.Config.Host
+		if c.Config.TlsHost != "" {
+			host = c.Config.TlsHost
+		}
+		if _, err := c.rdsAuth.GetFreshRdsToken(ctx, utils.RDSConnectionConfig{
+			Host: host,
+			Port: c.Config.Port,
+			User: c.Config.User,
+		}, "POSTGRES"); err != nil {
+			return fmt.Errorf("failed to refresh RDS IAM token before replication setup: %w", err)
+		}
+	}
 	conn, wst, err := c.CreateReplConn(ctx, env)
 	if err != nil {
 		return err

flow/connectors/utils/aws.go

@@ -216,6 +216,7 @@ func getPeerDBAWSEnv(connectorName string, awsKey string) string {
 func LoadPeerDBAWSEnvConfigProvider(connectorName string) *StaticAWSCredentialsProvider {
 	accessKeyId := getPeerDBAWSEnv(connectorName, "AWS_ACCESS_KEY_ID")
 	secretAccessKey := getPeerDBAWSEnv(connectorName, "AWS_SECRET_ACCESS_KEY")
+	sessionToken := getPeerDBAWSEnv(connectorName, "AWS_SESSION_TOKEN")
 	region := getPeerDBAWSEnv(connectorName, "AWS_REGION")
 	endpointUrl := getPeerDBAWSEnv(connectorName, "AWS_ENDPOINT_URL_S3")
 	rootCa := getPeerDBAWSEnv(connectorName, "ROOT_CA")
@@ -238,6 +239,7 @@ func LoadPeerDBAWSEnvConfigProvider(connectorName string) *StaticAWSCredentialsP
 		AWS: aws.Credentials{
 			AccessKeyID:     accessKeyId,
 			SecretAccessKey: secretAccessKey,
+			SessionToken:    sessionToken,
 		},
 		EndpointUrl: endpointUrlPtr,
 	}, region, rootCAs, tlsHost)

flow/connectors/utils/rds.go

@@ -86,23 +86,26 @@ func GetRDSToken(ctx context.Context, connConfig RDSConnectionConfig, rdsAuth *R
 		logger.Info("Using cached RDS token for connector", slog.String("connector", connectorName))
 		return rdsAuth.token, nil
 	}
-	return func() (string, error) {
-		logger.Info("Generating new RDS token for connector", slog.String("connector", connectorName))
-		rdsAuth.lock.Lock()
-		defer rdsAuth.lock.Unlock()
-		newUpdateTime := time.Now()
-		if rdsAuth.updateTime.Add(RDSAuthTokenTTL).After(now) && rdsAuth.token != "" {
-			return rdsAuth.token, nil
-		}
-		peerAWSCredentials := BuildPeerAWSCredentials(rdsAuth.AwsAuthConfig)
-		token, err := buildRdsToken(ctx, connConfig, peerAWSCredentials, connectorName)
-		if err != nil {
-			return "", err
-		}
-		rdsAuth.token = token
-		rdsAuth.updateTime = newUpdateTime
-		return token, nil
-	}()
+	logger.Info("Generating new RDS token for connector", slog.String("connector", connectorName))
+	return rdsAuth.GetFreshRdsToken(ctx, connConfig, connectorName)
+}
+
+// GetFreshRdsToken bypasses the cache TTL and immediately fetches a new RDS IAM token.
+// Use this when a connection break requires a guaranteed-fresh token for reconnect.
+func (r *RDSAuth) GetFreshRdsToken(
+	ctx context.Context,
+	connConfig RDSConnectionConfig,
+	connectorName string,
+) (string, error) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+	token, err := buildRdsToken(ctx, connConfig, BuildPeerAWSCredentials(r.AwsAuthConfig), connectorName)
+	if err != nil {
+		return "", err
+	}
+	r.token = token
+	r.updateTime = time.Now()
+	return token, nil
 }
 
 func buildRdsToken(