From 7ffd8501b4105c5c68f4966ac181dbaf14a370e6 Mon Sep 17 00:00:00 2001 From: ReguiguiMohamed Date: Tue, 18 Aug 2026 16:37:54 +0100 Subject: [PATCH] postgres_cdc: refresh IAM auth token for each new connection openPgConnectionFromConfig copied the password into a fresh pgconn.Config when the pool was built, so the heartbeat, monitor, snapshotter and server-version check all kept authenticating with the token captured at startup. Once that token expired the heartbeat write failed on every tick for the rest of the pipeline's life, while the replication stream stayed healthy. Rebuild the password through the existing Config.RefreshAuthToken from a BeforeConnect hook, which pgx runs before it opens each physical connection. Fixes #4668 --- .../postgresql/pglogicalstream/connection.go | 25 +++++- .../pglogicalstream/connection_test.go | 83 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) create mode 100644 internal/impl/postgresql/pglogicalstream/connection_test.go 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) + }) +}