quick refactor for rds auth

flow/connectors/mysql/cdc.go

@@ -311,15 +311,7 @@ func (c *MySqlConnector) startSyncer(ctx context.Context, env map[string]string)
 	config := c.config
 	if c.rdsAuth != nil {
 		c.logger.Info("Setting up IAM auth for MySQL replication")
-		host := c.config.Host
-		if c.config.TlsHost != "" {
-			host = c.config.TlsHost
-		}
-		token, err := utils.GetRDSToken(ctx, utils.RDSConnectionConfig{
-			Host: host,
-			Port: config.Port,
-			User: config.User,
-		}, c.rdsAuth, "MYSQL")
+		token, err := utils.GetRDSToken(ctx, c.rdsAuth, "MYSQL")
 		if err != nil {
 			return nil, err
 		}

flow/connectors/mysql/mysql.go

@@ -55,6 +55,16 @@
 		return nil, fmt.Errorf("failed to create ssh tunnel: %w", err)
 	}
 	logger := internal.LoggerFromCtx(ctx)
+	var rdsConnectionConfig utils.RDSConnectionConfig
+	host := config.Host
+	if config.TlsHost != "" {
+		host = config.TlsHost
+	}
+	rdsConnectionConfig = utils.RDSConnectionConfig{
+		Host: host,
+		Port: uint32(config.Port),
+		User: config.User,
+	}
 	var rdsAuth *utils.RDSAuth
 	if config.AuthType == protos.MySqlAuthType_MYSQL_IAM_AUTH {
 		rdsAuth = &utils.RDSAuth{
@@ -64,7 +74,10 @@
 			logger.Error("failed to verify auth config", slog.Any("error", err))
 			return nil, fmt.Errorf("failed to verify auth config: %w", err)
 		}
+	} else {
+		rdsAuth = &utils.RDSAuth{}
 	}
+	rdsAuth.ConnectionConfig = rdsConnectionConfig
 	contexts := make(chan context.Context)
 	c := &MySqlConnector{
 		PostgresMetadata:      pgMetadata,
@@ -198,15 +211,7 @@
 		config := c.config
 		if c.rdsAuth != nil {
 			c.logger.Info("Setting up IAM auth for MySQL")
-			host := c.config.Host
-			if c.config.TlsHost != "" {
-				host = c.config.TlsHost
-			}
-			token, err := utils.GetRDSToken(ctx, utils.RDSConnectionConfig{
-				Host: host,
-				Port: config.Port,
-				User: config.User,
-			}, c.rdsAuth, "MYSQL")
+			token, err := utils.GetRDSToken(ctx, c.rdsAuth, "MYSQL")
 			if err != nil {
 				return nil, err
 			}

flow/connectors/postgres/postgres.go

@@ -99,6 +99,17 @@ func newPostgresConnector(
 		return nil, fmt.Errorf("failed to create ssh tunnel: %w", err)
 	}
 
+	var rdsConnectionConfig utils.RDSConnectionConfig
+	host := connConfig.Host
+	if pgConfig.TlsHost != "" {
+		host = pgConfig.TlsHost
+	}
+	rdsConnectionConfig = utils.RDSConnectionConfig{
+		Host: host,
+		Port: uint32(connConfig.Port),
+		User: connConfig.User,
+	}
+
 	var rdsAuth *utils.RDSAuth
 	if pgConfig.AuthType == protos.PostgresAuthType_POSTGRES_IAM_AUTH {
 		rdsAuth = &utils.RDSAuth{
@@ -108,8 +119,11 @@ func newPostgresConnector(
 			logger.Error("failed to verify auth config", slog.Any("error", err))
 			return nil, fmt.Errorf("failed to verify auth config: %w", err)
 		}
+	} else {
+		rdsAuth = &utils.RDSAuth{}
 	}
-	conn, err := NewPostgresConnFromConfig(ctx, connConfig, pgConfig.TlsHost, rdsAuth, tunnel)
+	rdsAuth.ConnectionConfig = rdsConnectionConfig
+	conn, err := NewPostgresConnFromConfig(ctx, connConfig, rdsAuth, tunnel)
 	if err != nil {
 		tunnel.Close()
 

flow/connectors/postgres/postgres_source.go

@@ -105,7 +105,7 @@ func (c *PostgresConnector) CreateReplConn(ctx context.Context, env map[string]s
 		c.logger.Info("not setting wal_sender_timeout")
 	}
 
-	conn, err := NewPostgresConnFromConfig(ctx, replConfig, c.Config.TlsHost, c.rdsAuth, c.ssh)
+	conn, err := NewPostgresConnFromConfig(ctx, replConfig, c.rdsAuth, c.ssh)
 	if err != nil {
 		internal.LoggerFromCtx(ctx).Error("failed to create replication connection", slog.Any("error", err))
 		return nil, walSenderTimeout{}, fmt.Errorf("failed to create replication connection: %w", err)

flow/connectors/postgres/qrep_partition_test.go

@@ -89,7 +89,7 @@ func setupTestSchema(t *testing.T) (string, *pgx.Conn, string) {
 	}
 	t.Cleanup(func() { tunnel.Close() })
 
-	conn, err := NewPostgresConnFromConfig(t.Context(), config, "", nil, tunnel)
+	conn, err := NewPostgresConnFromConfig(t.Context(), config, nil, tunnel)
 	if err != nil {
 		t.Fatalf("Failed to create connection: %v", err)
 	}

flow/connectors/postgres/ssh_wrapped_conn.go

@@ -15,7 +15,6 @@ import (
 func NewPostgresConnFromConfig(
 	ctx context.Context,
 	connConfig *pgx.ConnConfig,
-	tlsHost string,
 	rdsAuth *utils.RDSAuth,
 	tunnel *utils.SSHTunnel,
 ) (*pgx.Conn, error) {
@@ -29,16 +28,18 @@ func NewPostgresConnFromConfig(
 	}
 	logger := internal.LoggerFromCtx(ctx)
 	if rdsAuth != nil {
-		host := connConfig.Host
-		if tlsHost != "" {
-			host = tlsHost
+		if rdsAuth.ConnectionConfig.Host == "" {
+			rdsAuth.ConnectionConfig.Host = connConfig.Host
 		}
+		if rdsAuth.ConnectionConfig.Port == 0 {
+			rdsAuth.ConnectionConfig.Port = uint32(connConfig.Port)
+		}
+		if rdsAuth.ConnectionConfig.User == "" {
+			rdsAuth.ConnectionConfig.User = connConfig.User
+		}
+
 		logger.Info("Setting up IAM auth for Postgres")
-		token, err := utils.GetRDSToken(ctx, utils.RDSConnectionConfig{
-			Host: host,
-			Port: uint32(connConfig.Port),
-			User: connConfig.User,
-		}, rdsAuth, "POSTGRES")
+		token, err := utils.GetRDSToken(ctx, rdsAuth, "POSTGRES")
 		if err != nil {
 			return nil, err
 		}

flow/connectors/utils/rds.go

@@ -21,11 +21,12 @@
 // RDSAuthTokenTTL is the cache TTL for RDS auth tokens. RDS Tokens Live for 15 minutes by default
 const RDSAuthTokenTTL = 10 * time.Minute
 
 type RDSAuth struct {
-	updateTime    time.Time
-	AwsAuthConfig *protos.AwsAuthenticationConfig
-	token         string
-	lock          sync.Mutex
+	updateTime       time.Time
+	AwsAuthConfig    *protos.AwsAuthenticationConfig
+	ConnectionConfig RDSConnectionConfig
+	token            string
+	lock             sync.Mutex
 }
 
 func (r *RDSAuth) VerifyAuthConfig() error {
@@ -79,15 +80,15 @@
 
 var regionRegex = regexp.MustCompile(`^.*?\..*?\.([a-z0-9-]+)\.rds\.amazonaws\.com$`)
 
-func GetRDSToken(ctx context.Context, connConfig RDSConnectionConfig, rdsAuth *RDSAuth, connectorName string) (string, error) {
+func GetRDSToken(ctx context.Context, rdsAuth *RDSAuth, connectorName string) (string, error) {
 	logger := internal.LoggerFromCtx(ctx)
 	now := time.Now()
 	if rdsAuth.updateTime.Add(RDSAuthTokenTTL).After(now) && rdsAuth.token != "" {
 		logger.Info("Using cached RDS token for connector", slog.String("connector", connectorName))
 		return rdsAuth.token, nil
 	}
 	logger.Info("Generating new RDS token for connector", slog.String("connector", connectorName))
-	return rdsAuth.GetFreshRdsToken(ctx, connConfig, connectorName)
+	return rdsAuth.GetFreshRdsToken(ctx, rdsAuth.ConnectionConfig, connectorName)
 }
 
 // GetFreshRdsToken bypasses the cache TTL and immediately fetches a new RDS IAM token.