fix: wait for pooler after init temp role (#4018)

* Revert "fix: retry pooler connection after password change (#3949)"

This reverts commit 2b6e22b.

* fix: wait for pooler after init temp role

Author: sweatybridge

internal/bootstrap/bootstrap.go

@@ -9,7 +9,6 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
-	"time"
 
 	"github.com/cenkalti/backoff/v4"
 	"github.com/go-errors/errors"
@@ -88,12 +87,12 @@ func Run(ctx context.Context, starter StarterTemplate, fsys afero.Fs, options ..
 	}
 	// 3. Get api keys
 	var keys []api.ApiKeyResponse
-	policy := newBackoffPolicy(ctx)
+	policy := utils.NewBackoffPolicy(ctx)
 	if err := backoff.RetryNotify(func() error {
 		fmt.Fprintln(os.Stderr, "Linking project...")
 		keys, err = apiKeys.RunGetApiKeys(ctx, flags.ProjectRef)
 		return err
-	}, policy, newErrorCallback()); err != nil {
+	}, policy, utils.NewErrorCallback()); err != nil {
 		return err
 	}
 	// 4. Link project
@@ -109,7 +108,7 @@ func Run(ctx context.Context, starter StarterTemplate, fsys afero.Fs, options ..
 	if err := backoff.RetryNotify(func() error {
 		fmt.Fprintln(os.Stderr, "Checking project health...")
 		return checkProjectHealth(ctx)
-	}, policy, newErrorCallback()); err != nil {
+	}, policy, utils.NewErrorCallback()); err != nil {
 		return err
 	}
 	// 6. Push migrations
@@ -120,7 +119,7 @@ func Run(ctx context.Context, starter StarterTemplate, fsys afero.Fs, options ..
 	policy.Reset()
 	if err := backoff.RetryNotify(func() error {
 		return push.Run(ctx, false, false, true, true, config, fsys)
-	}, policy, newErrorCallback()); err != nil {
+	}, policy, utils.NewErrorCallback()); err != nil {
 		return err
 	}
 	// 7. TODO: deploy functions
@@ -171,32 +170,6 @@ func checkProjectHealth(ctx context.Context) error {
 	return nil
 }
 
-const maxRetries = 8
-
-func newBackoffPolicy(ctx context.Context) backoff.BackOffContext {
-	b := backoff.ExponentialBackOff{
-		InitialInterval:     3 * time.Second,
-		RandomizationFactor: backoff.DefaultRandomizationFactor,
-		Multiplier:          backoff.DefaultMultiplier,
-		MaxInterval:         backoff.DefaultMaxInterval,
-		MaxElapsedTime:      backoff.DefaultMaxElapsedTime,
-		Stop:                backoff.Stop,
-		Clock:               backoff.SystemClock,
-	}
-	b.Reset()
-	return backoff.WithContext(backoff.WithMaxRetries(&b, maxRetries), ctx)
-}
-
-func newErrorCallback() backoff.Notify {
-	failureCount := 0
-	logger := utils.GetDebugLogger()
-	return func(err error, d time.Duration) {
-		failureCount += 1
-		fmt.Fprintln(logger, err)
-		fmt.Fprintf(os.Stderr, "Retry (%d/%d): ", failureCount, maxRetries)
-	}
-}
-
 const (
 	SUPABASE_SERVICE_ROLE_KEY = "SUPABASE_SERVICE_ROLE_KEY"
 	SUPABASE_ANON_KEY         = "SUPABASE_ANON_KEY"

internal/utils/flags/db_url.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"text/template"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/go-errors/errors"
 	"github.com/jackc/pgconn"
 	"github.com/spf13/afero"
@@ -133,6 +134,7 @@ func NewDbConfigWithPassword(ctx context.Context, projectRef string) pgconn.Conf
 			// Special handling for pooler username
 			if suffix := "." + projectRef; strings.HasSuffix(config.User, suffix) {
 				newRole.User += suffix
+				defer tryPooler(ctx, &config)
 			}
 			config.User = newRole.User
 			return config
@@ -148,6 +150,18 @@ func NewDbConfigWithPassword(ctx context.Context, projectRef string) pgconn.Conf
 	return config
 }
 
+func tryPooler(ctx context.Context, config *pgconn.Config) {
+	if err := backoff.RetryNotify(func() error {
+		conn, err := pgconn.ConnectConfig(ctx, config)
+		if err != nil {
+			return errors.Errorf("failed to connect as temp role: %w", err)
+		}
+		return conn.Close(ctx)
+	}, utils.NewBackoffPolicy(ctx), utils.NewErrorCallback()); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+	}
+}
+
 var (
 	//go:embed queries/role.sql
 	initRoleEmbed    string

internal/utils/retry.go

@@ -0,0 +1,30 @@
+package utils
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+)
+
+const maxRetries = 8
+
+func NewBackoffPolicy(ctx context.Context) backoff.BackOffContext {
+	b := backoff.NewExponentialBackOff(backoff.WithInitialInterval(3 * time.Second))
+	return backoff.WithContext(backoff.WithMaxRetries(b, maxRetries), ctx)
+}
+
+func NewErrorCallback() backoff.Notify {
+	failureCount := 0
+	logger := GetDebugLogger()
+	return func(err error, d time.Duration) {
+		failureCount += 1
+		if failureCount*3 > maxRetries {
+			logger = os.Stderr
+		}
+		fmt.Fprintln(logger, err)
+		fmt.Fprintf(logger, "Retry (%d/%d): ", failureCount, maxRetries)
+	}
+}

pkg/pgxv5/connect.go

@@ -5,9 +5,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
-	"time"
 
-	"github.com/cenkalti/backoff/v4"
 	"github.com/go-errors/errors"
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgx/v4"
@@ -28,31 +26,19 @@ func Connect(ctx context.Context, connString string, options ...func(*pgx.ConnCo
 	config.OnNotice = func(pc *pgconn.PgConn, n *pgconn.Notice) {
 		fmt.Fprintf(os.Stderr, "%s (%s): %s\n", n.Severity, n.Code, n.Message)
 	}
-	maxRetries := uint64(0)
 	if strings.HasPrefix(config.User, CLI_LOGIN_ROLE) {
 		config.AfterConnect = func(ctx context.Context, pgconn *pgconn.PgConn) error {
 			return pgconn.Exec(ctx, SET_SESSION_ROLE).Close()
 		}
-		// Add retry to allow enough time for password change to propagate to pooler
-		if len(config.User) > len(CLI_LOGIN_ROLE) {
-			maxRetries = 3
-		}
 	}
 	// Apply config overrides
 	for _, op := range options {
 		op(config)
 	}
 	// Connect to database
-	connect := func() (*pgx.Conn, error) {
-		conn, err := pgx.ConnectConfig(ctx, config)
-		if err != nil {
-			return nil, errors.Errorf("failed to connect to postgres: %w", err)
-		}
-		return conn, nil
+	conn, err := pgx.ConnectConfig(ctx, config)
+	if err != nil {
+		return nil, errors.Errorf("failed to connect to postgres: %w", err)
 	}
-	policy := backoff.WithContext(backoff.WithMaxRetries(backoff.NewExponentialBackOff(
-		backoff.WithInitialInterval(3*time.Second)),
-		maxRetries),
-		ctx)
-	return backoff.RetryWithData(connect, policy)
+	return conn, nil
 }