Skip to content

Add lock timeout retry support to apply #558

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tianzhou wants to merge 2 commits into from
Closed

Add lock timeout retry support to apply #558

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6cccfc2
Add lock timeout retry support to apply
claude Aug 25, 2026
899cb15
Validate lock timeout retry settings and clarify retry log wording
claude Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 105 additions & 9 deletions cmd/apply/apply.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,14 @@ import (
"bufio"
"context"
"database/sql"
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"time"

"github.com/jackc/pgx/v5/pgconn"
planCmd "github.com/pgplex/pgschema/cmd/plan"
"github.com/pgplex/pgschema/cmd/util"
"github.com/pgplex/pgschema/internal/fingerprint"
Expand All @@ -19,6 +22,13 @@ import (
"github.com/spf13/cobra"
)

// lockNotAvailableSQLState is the PostgreSQL SQLSTATE raised when a statement
// fails to acquire a lock before lock_timeout expires.
const lockNotAvailableSQLState = "55P03"

// maxLockRetryBackoff caps the exponential backoff between lock timeout retries.
const maxLockRetryBackoff = 30 * time.Second

var (
applyHost string
applyPort int
Expand All @@ -31,6 +41,8 @@ var (
applyAutoApprove bool
applyNoColor bool
applyLockTimeout string
applyLockRetries int
applyLockRetryWait time.Duration
applyApplicationName string

// Plan database connection flags (optional - for using external database instead of embedded postgres)
Expand Down Expand Up @@ -72,6 +84,8 @@ func init() {
ApplyCmd.Flags().BoolVar(&applyAutoApprove, "auto-approve", false, "Apply changes without prompting for approval")
ApplyCmd.Flags().BoolVar(&applyNoColor, "no-color", false, "Disable colored output")
ApplyCmd.Flags().StringVar(&applyLockTimeout, "lock-timeout", "", "Maximum time to wait for database locks (e.g., 30s, 5m, 1h)")
ApplyCmd.Flags().IntVar(&applyLockRetries, "lock-timeout-retries", 0, "Number of times to retry a statement that fails with a lock timeout (requires --lock-timeout)")
ApplyCmd.Flags().DurationVar(&applyLockRetryWait, "lock-timeout-retry-wait", 1*time.Second, "Initial wait before retrying a lock timeout failure; doubles after each retry up to 30s")
ApplyCmd.Flags().StringVar(&applyApplicationName, "application-name", "pgschema", "Application name for database connection (visible in pg_stat_activity) (env: PGAPPNAME)")

// Plan database connection flags (optional - for using external database instead of embedded postgres when using --file)
Expand Down Expand Up @@ -103,6 +117,8 @@ type ApplyConfig struct {
NoColor bool
Quiet bool // Suppress plan display and progress messages (useful for tests)
LockTimeout string
LockRetries int // Number of retries on lock timeout (55P03) failures
LockRetryWait time.Duration // Initial backoff before the first retry; doubles thereafter
ApplicationName string
SSLMode string
// Plan database configuration (needed when GeneratePlan checks provider SSL mode)
Expand All @@ -118,6 +134,19 @@ type ApplyConfig struct {
// If config.File is provided, provider is used to generate the plan.
// The caller is responsible for managing the provider lifecycle (creation and cleanup).
func ApplyMigration(config *ApplyConfig, provider postgres.DesiredStateProvider) error {
// Validate lock retry settings here too, not just in RunApply: a direct
// caller of ApplyMigration could otherwise set LockRetries without a
// LockTimeout and block indefinitely instead of getting a retryable error.
if config.LockRetries < 0 {
return fmt.Errorf("lock timeout retries must be non-negative")
}
if config.LockRetries > 0 && config.LockTimeout == "" {
return fmt.Errorf("lock timeout retries require a lock timeout")
}
if config.LockRetries > 0 && config.LockRetryWait <= 0 {
return fmt.Errorf("lock timeout retry wait must be greater than zero")
}

var migrationPlan *plan.Plan
var err error

Expand Down Expand Up @@ -247,13 +276,15 @@ func ApplyMigration(config *ApplyConfig, provider postgres.DesiredStateProvider)
return nil
}

retry := lockRetryConfig{MaxRetries: config.LockRetries, Backoff: config.LockRetryWait}
Comment thread
tianzhou marked this conversation as resolved.

// Execute by groups with wait directive support
for i, group := range migrationPlan.Groups {
if !config.Quiet {
fmt.Printf("\nExecuting group %d/%d...\n", i+1, len(migrationPlan.Groups))
}

err = executeGroup(ctx, conn, group, i+1, config.Quiet)
err = executeGroup(ctx, conn, group, i+1, config.Quiet, retry)
if err != nil {
return err
}
Expand Down Expand Up @@ -316,6 +347,18 @@ func RunApply(cmd *cobra.Command, args []string) error {
return err
}

// Lock timeout retries only make sense alongside a lock timeout: without one,
// statements block indefinitely instead of failing with a retryable error.
if applyLockRetries < 0 {
return fmt.Errorf("--lock-timeout-retries must be non-negative")
}
if applyLockRetries > 0 && applyLockTimeout == "" {
return fmt.Errorf("--lock-timeout-retries requires --lock-timeout to be set")
}
Comment thread
tianzhou marked this conversation as resolved.
if applyLockRetries > 0 && applyLockRetryWait <= 0 {
return fmt.Errorf("--lock-timeout-retry-wait must be greater than zero")
}

// Build configuration
config := &ApplyConfig{
Host: applyHost,
Expand All @@ -327,6 +370,8 @@ func RunApply(cmd *cobra.Command, args []string) error {
AutoApprove: applyAutoApprove,
NoColor: applyNoColor,
LockTimeout: applyLockTimeout,
LockRetries: applyLockRetries,
LockRetryWait: applyLockRetryWait,
ApplicationName: applyApplicationName,
SSLMode: finalSSLMode,
}
Expand Down Expand Up @@ -444,7 +489,7 @@ func validateSchemaFingerprint(migrationPlan *plan.Plan, host string, port int,
}

// executeGroup executes all steps in a group, handling directives separately from SQL statements
func executeGroup(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup, groupNum int, quiet bool) error {
func executeGroup(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup, groupNum int, quiet bool, retry lockRetryConfig) error {
// Check if this group has directives
hasDirectives := false

Expand All @@ -457,15 +502,15 @@ func executeGroup(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup,

if !hasDirectives {
// No directives - concatenate all SQL and execute in implicit transaction
return executeGroupConcatenated(ctx, conn, group, groupNum, quiet)
return executeGroupConcatenated(ctx, conn, group, groupNum, quiet, retry)
} else {
// Has directives - execute statements individually
return executeGroupIndividually(ctx, conn, group, groupNum, quiet)
return executeGroupIndividually(ctx, conn, group, groupNum, quiet, retry)
}
}

// executeGroupConcatenated concatenates all SQL statements and executes them in an implicit transaction
func executeGroupConcatenated(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup, groupNum int, quiet bool) error {
func executeGroupConcatenated(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup, groupNum int, quiet bool, retry lockRetryConfig) error {
var sqlStatements []string

// Collect all SQL statements
Expand All @@ -480,8 +525,9 @@ func executeGroupConcatenated(ctx context.Context, conn *sql.DB, group plan.Exec
fmt.Printf(" Executing %d statements in implicit transaction\n", len(sqlStatements))
}

// Execute all statements in a single call (implicit transaction)
_, err := util.ExecContextWithLogging(ctx, conn, concatenatedSQL, fmt.Sprintf("execute %d statements in group %d", len(sqlStatements), groupNum))
// Execute all statements in a single call (implicit transaction). Since the
// batch is atomic, retrying the whole batch on a lock timeout is safe.
_, err := execWithLockRetry(ctx, conn, concatenatedSQL, fmt.Sprintf("execute %d statements in group %d", len(sqlStatements), groupNum), retry, quiet)
if err != nil {
return fmt.Errorf("failed to execute concatenated statements in group %d: %w", groupNum, err)
}
Expand All @@ -490,7 +536,7 @@ func executeGroupConcatenated(ctx context.Context, conn *sql.DB, group plan.Exec
}

// executeGroupIndividually executes statements individually without transactions
func executeGroupIndividually(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup, groupNum int, quiet bool) error {
func executeGroupIndividually(ctx context.Context, conn *sql.DB, group plan.ExecutionGroup, groupNum int, quiet bool, retry lockRetryConfig) error {
for stepIdx, step := range group.Steps {
if step.Directive != nil {
// Handle directive execution
Expand All @@ -504,7 +550,7 @@ func executeGroupIndividually(ctx context.Context, conn *sql.DB, group plan.Exec
fmt.Printf(" Executing: %s\n", truncateSQL(step.SQL, 80))
}

_, err := util.ExecContextWithLogging(ctx, conn, step.SQL, fmt.Sprintf("execute statement in group %d, step %d", groupNum, stepIdx+1))
_, err := execWithLockRetry(ctx, conn, step.SQL, fmt.Sprintf("execute statement in group %d, step %d", groupNum, stepIdx+1), retry, quiet)
if err != nil {
return fmt.Errorf("failed to execute statement in group %d, step %d: %w", groupNum, stepIdx+1, err)
}
Expand All @@ -513,6 +559,56 @@ func executeGroupIndividually(ctx context.Context, conn *sql.DB, group plan.Exec
return nil
}

// lockRetryConfig controls retry behavior for statements that fail because they
// could not acquire a lock before lock_timeout expired (SQLSTATE 55P03).
type lockRetryConfig struct {
MaxRetries int // Number of additional attempts after the first failure
Backoff time.Duration // Initial wait before the first retry; doubles after each subsequent retry
}

// sqlExecer is satisfied by both *sql.DB and *sql.Conn, matching util.ExecContextWithLogging.
type sqlExecer interface {
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
}

// isLockTimeoutError reports whether err is a PostgreSQL error raised because a
// statement failed to acquire a lock before lock_timeout expired.
func isLockTimeoutError(err error) bool {
var pgErr *pgconn.PgError
if errors.As(err, &pgErr) {
return pgErr.Code == lockNotAvailableSQLState
}
return false
}

// execWithLockRetry executes sqlStmt, retrying with exponential backoff when the
// failure is a lock timeout (SQLSTATE 55P03). Any other error is returned immediately.
func execWithLockRetry(ctx context.Context, conn sqlExecer, sqlStmt, description string, retry lockRetryConfig, quiet bool) (sql.Result, error) {
backoff := retry.Backoff

for attempt := 0; ; attempt++ {
result, err := util.ExecContextWithLogging(ctx, conn, sqlStmt, description)
if err == nil || !isLockTimeoutError(err) || attempt >= retry.MaxRetries {
return result, err
}

if !quiet {
fmt.Printf(" Lock not available, retrying in %s (retry %d/%d)...\n", backoff, attempt+1, retry.MaxRetries)
}

select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(backoff):
}
Comment on lines +599 to +603

@tianzhou tianzhou Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not making this change. time.After in a select/ctx.Done() loop only "leaks" a timer until it fires or becomes unreachable; since Go 1.23 the runtime can garbage-collect an unreferenced, unstopped Timer without waiting for it to fire, so there's no real leak here. This loop also isn't a hot path - it only re-fires on an actual lock-timeout error, bounded by --lock-timeout-retries (a handful of iterations per statement at most), and the existing wait-directive polling loop in cmd/apply/directive.go uses the exact same time.After pattern for consistency. Happy to switch both to time.NewTimer if this ever becomes a tight, high-frequency loop, but for now it'd be defensive coding against a cost that isn't there.

Generated by Claude Code


backoff *= 2
if backoff > maxLockRetryBackoff {
backoff = maxLockRetryBackoff
}
}
}

// truncateSQL truncates a SQL statement for display purposes
func truncateSQL(sql string, maxLen int) string {
// Remove extra whitespace and newlines
Expand Down
77 changes: 77 additions & 0 deletions cmd/apply/apply_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,24 @@ func TestApplyCommand(t *testing.T) {
t.Errorf("Expected default lock-timeout to be empty, got '%s'", lockTimeoutFlag.DefValue)
}

// Test lock-timeout-retries flag
lockRetriesFlag := flags.Lookup("lock-timeout-retries")
if lockRetriesFlag == nil {
t.Error("Expected --lock-timeout-retries flag to be defined")
}
if lockRetriesFlag.DefValue != "0" {
t.Errorf("Expected default lock-timeout-retries to be '0', got '%s'", lockRetriesFlag.DefValue)
}

// Test lock-timeout-retry-wait flag
lockRetryWaitFlag := flags.Lookup("lock-timeout-retry-wait")
if lockRetryWaitFlag == nil {
t.Error("Expected --lock-timeout-retry-wait flag to be defined")
}
if lockRetryWaitFlag.DefValue != "1s" {
t.Errorf("Expected default lock-timeout-retry-wait to be '1s', got '%s'", lockRetryWaitFlag.DefValue)
}

// Test application-name flag
applicationNameFlag := flags.Lookup("application-name")
if applicationNameFlag == nil {
Expand Down Expand Up @@ -208,11 +226,17 @@ func TestApplyCommandFlagValidation(t *testing.T) {
origUser := applyUser
origFile := applyFile
origPlan := applyPlan
origLockTimeout := applyLockTimeout
origLockRetries := applyLockRetries
origLockRetryWait := applyLockRetryWait
defer func() {
applyDB = origDB
applyUser = origUser
applyFile = origFile
applyPlan = origPlan
applyLockTimeout = origLockTimeout
applyLockRetries = origLockRetries
applyLockRetryWait = origLockRetryWait
}()

t.Run("neither file nor plan specified", func(t *testing.T) {
Expand All @@ -231,6 +255,59 @@ func TestApplyCommandFlagValidation(t *testing.T) {
}
})

t.Run("lock timeout retries without lock timeout", func(t *testing.T) {
// Reset flags
applyDB = "testdb"
applyUser = "testuser"
applyFile = "schema.sql"
applyPlan = ""
applyLockTimeout = ""
applyLockRetries = 3

err := RunApply(ApplyCmd, []string{})
if err == nil {
t.Error("Expected error when --lock-timeout-retries is set without --lock-timeout")
}
if err != nil && err.Error() != "--lock-timeout-retries requires --lock-timeout to be set" {
t.Errorf("Expected specific error message, got: %v", err)
}
})

t.Run("negative lock timeout retries", func(t *testing.T) {
applyDB = "testdb"
applyUser = "testuser"
applyFile = "schema.sql"
applyPlan = ""
applyLockTimeout = "5s"
applyLockRetries = -1

err := RunApply(ApplyCmd, []string{})
if err == nil {
t.Error("Expected error when --lock-timeout-retries is negative")
}
if err != nil && err.Error() != "--lock-timeout-retries must be non-negative" {
t.Errorf("Expected specific error message, got: %v", err)
}
})

t.Run("non-positive lock timeout retry wait", func(t *testing.T) {
applyDB = "testdb"
applyUser = "testuser"
applyFile = "schema.sql"
applyPlan = ""
applyLockTimeout = "5s"
applyLockRetries = 3
applyLockRetryWait = 0

err := RunApply(ApplyCmd, []string{})
if err == nil {
t.Error("Expected error when --lock-timeout-retry-wait is non-positive")
}
if err != nil && err.Error() != "--lock-timeout-retry-wait must be greater than zero" {
t.Errorf("Expected specific error message, got: %v", err)
}
})

t.Run("both file and plan specified", func(t *testing.T) {
// Create a test command to test mutual exclusivity
testCmd := &cobra.Command{
Expand Down
Loading