diff --git a/flow/connectors/mysql/cdc.go b/flow/connectors/mysql/cdc.go index 04acf62d4..846656491 100644 --- a/flow/connectors/mysql/cdc.go +++ b/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 } diff --git a/flow/connectors/mysql/mysql.go b/flow/connectors/mysql/mysql.go index 3a6c6a115..638ad6588 100644 --- a/flow/connectors/mysql/mysql.go +++ b/flow/connectors/mysql/mysql.go @@ -55,6 +55,16 @@ func NewMySqlConnector(ctx context.Context, config *protos.MySqlConfig) (*MySqlC 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 @@ func NewMySqlConnector(ctx context.Context, config *protos.MySqlConfig) (*MySqlC 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 @@ func (c *MySqlConnector) connect(ctx context.Context) (*client.Conn, error) { 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 } diff --git a/flow/connectors/postgres/postgres.go b/flow/connectors/postgres/postgres.go index b89045b25..bffcb67d5 100644 --- a/flow/connectors/postgres/postgres.go +++ b/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() diff --git a/flow/connectors/postgres/postgres_source.go b/flow/connectors/postgres/postgres_source.go index d2c42c0f7..546dabf5d 100644 --- a/flow/connectors/postgres/postgres_source.go +++ b/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) diff --git a/flow/connectors/postgres/qrep_partition_test.go b/flow/connectors/postgres/qrep_partition_test.go index 7038abbff..c3ccf3b5c 100644 --- a/flow/connectors/postgres/qrep_partition_test.go +++ b/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) } diff --git a/flow/connectors/postgres/ssh_wrapped_conn.go b/flow/connectors/postgres/ssh_wrapped_conn.go index a23314688..afd0e988b 100644 --- a/flow/connectors/postgres/ssh_wrapped_conn.go +++ b/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 } diff --git a/flow/connectors/utils/rds.go b/flow/connectors/utils/rds.go index 440ce2ed1..e9c841e8b 100644 --- a/flow/connectors/utils/rds.go +++ b/flow/connectors/utils/rds.go @@ -22,10 +22,11 @@ import ( 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,7 +80,7 @@ func BuildPeerAWSCredentials(awsAuth *protos.AwsAuthenticationConfig) PeerAWSCre 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 != "" { @@ -87,7 +88,7 @@ func GetRDSToken(ctx context.Context, connConfig RDSConnectionConfig, rdsAuth *R 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.