diff --git a/internal/impl/postgresql/pglogicalstream/connection.go b/internal/impl/postgresql/pglogicalstream/connection.go index 698a638eba..1ea19dc02a 100644 --- a/internal/impl/postgresql/pglogicalstream/connection.go +++ b/internal/impl/postgresql/pglogicalstream/connection.go @@ -9,17 +9,24 @@ package pglogicalstream import ( + "context" "database/sql" "fmt" "regexp" "strconv" + "sync" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/pgx/v5/stdlib" ) var re = regexp.MustCompile(`^(\d+)`) +// Every pool refreshes the token through the same Config.DBConfig, so the +// refresh and the read that follows it have to be serialised. +var refreshAuthTokenMu sync.Mutex + func openPgConnectionFromConfig(cfg *Config) (*sql.DB, error) { parsedCfg, err := pgxpool.ParseConfig(cfg.DBRawDSN) if err != nil { @@ -27,7 +34,23 @@ func openPgConnectionFromConfig(cfg *Config) (*sql.DB, error) { } parsedCfg.ConnConfig.Password = cfg.DBConfig.Password parsedCfg.ConnConfig.TLSConfig = cfg.TLSConfig - return stdlib.OpenDB(*parsedCfg.ConnConfig), nil + return stdlib.OpenDB(*parsedCfg.ConnConfig, stdlib.OptionBeforeConnect( + func(ctx context.Context, connCfg *pgx.ConnConfig) error { + if cfg.RefreshAuthToken == nil { + return nil + } + // IAM tokens expire long before the pipeline does, so rebuild the + // password for every new connection rather than reusing the one + // captured when the pool was opened. + refreshAuthTokenMu.Lock() + defer refreshAuthTokenMu.Unlock() + if err := cfg.RefreshAuthToken(ctx); err != nil { + return err + } + connCfg.Password = cfg.DBConfig.Password + return nil + }, + )), nil } func getPostgresVersion(cfg *Config) (int, error) { diff --git a/internal/impl/postgresql/pglogicalstream/connection_test.go b/internal/impl/postgresql/pglogicalstream/connection_test.go new file mode 100644 index 0000000000..206919b8e9 --- /dev/null +++ b/internal/impl/postgresql/pglogicalstream/connection_test.go @@ -0,0 +1,83 @@ +// Copyright 2025 Redpanda Data, Inc. +// +// Licensed as a Redpanda Enterprise file under the Redpanda Community +// License (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// https://github.com/redpanda-data/connect/v4/blob/main/licenses/rcl.md + +package pglogicalstream + +import ( + "context" + "errors" + "fmt" + "sync/atomic" + "testing" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Nothing listens on port 1, so every attempt fails at the dial. That is late +// enough for these tests, which only care about what happens before it. +const unreachableDSN = "postgres://user:startup@127.0.0.1:1/db?sslmode=disable" + +func TestOpenPgConnectionFromConfigAuthToken(t *testing.T) { + t.Run("refreshes the token for every new connection", func(t *testing.T) { + dbConf, err := pgconn.ParseConfig(unreachableDSN) + require.NoError(t, err) + + var refreshes atomic.Int64 + db, err := openPgConnectionFromConfig(&Config{ + DBRawDSN: unreachableDSN, + DBConfig: dbConf, + RefreshAuthToken: func(context.Context) error { + dbConf.Password = fmt.Sprintf("token-%d", refreshes.Add(1)) + return nil + }, + }) + require.NoError(t, err) + defer db.Close() + + for range 2 { + require.Error(t, db.PingContext(t.Context())) + } + assert.Equal(t, int64(2), refreshes.Load()) + assert.Equal(t, "token-2", dbConf.Password) + }) + + t.Run("surfaces a failed refresh", func(t *testing.T) { + dbConf, err := pgconn.ParseConfig(unreachableDSN) + require.NoError(t, err) + + errRefresh := errors.New("token expired") + db, err := openPgConnectionFromConfig(&Config{ + DBRawDSN: unreachableDSN, + DBConfig: dbConf, + RefreshAuthToken: func(context.Context) error { + return errRefresh + }, + }) + require.NoError(t, err) + defer db.Close() + + assert.ErrorIs(t, db.PingContext(t.Context()), errRefresh) + }) + + t.Run("keeps the parsed password when no refresh is configured", func(t *testing.T) { + dbConf, err := pgconn.ParseConfig(unreachableDSN) + require.NoError(t, err) + + db, err := openPgConnectionFromConfig(&Config{ + DBRawDSN: unreachableDSN, + DBConfig: dbConf, + }) + require.NoError(t, err) + defer db.Close() + + require.Error(t, db.PingContext(t.Context())) + assert.Equal(t, "startup", dbConf.Password) + }) +}