feat(storage): rebind transaction and pinned-connection SQL exactly once

pkg/storage/internal/sqlstore/applies.go

@@ -64,9 +64,9 @@ type applyStore struct {
 }
 
 type applyWriteTx struct {
-	tx             *sql.Tx
+	tx             *rebindTx
 	locker         namedlock.Locker
-	targetLockConn *sql.Conn
+	targetLockConn *rebindConn
 	targetLockName string
 }
 
@@ -75,7 +75,7 @@ type queryRower interface {
 }
 
 type txBeginner interface {
-	BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
+	BeginTx(ctx context.Context, opts *sql.TxOptions) (*rebindTx, error)
 }
 
 // claimableApplyStates returns active apply states where operator recovery can
@@ -229,7 +229,7 @@ func applyTargetLockName(database, dbType, environment string) string {
 	return "schemabot_apply_" + hex.EncodeToString(sum[:16])
 }
 
-func acquireApplyTargetLockConn(ctx context.Context, db *rebindDB, locker namedlock.Locker, database, dbType, environment string) (*sql.Conn, string, error) {
+func acquireApplyTargetLockConn(ctx context.Context, db *rebindDB, locker namedlock.Locker, database, dbType, environment string) (*rebindConn, string, error) {
 	if !hasApplyTarget(database, dbType, environment) {
 		return nil, "", fmt.Errorf("active apply target is required for %s/%s/%s", database, dbType, environment)
 	}
@@ -243,7 +243,7 @@ func acquireApplyTargetLockConn(ctx context.Context, db *rebindDB, locker namedl
 	}
 
 	lockName := applyTargetLockName(database, dbType, environment)
-	acquired, err := locker.Acquire(ctx, conn, lockName, applyTargetLockWait)
+	acquired, err := locker.Acquire(ctx, conn.raw(), lockName, applyTargetLockWait)
 	if err != nil {
 		slog.WarnContext(ctx, "failed to acquire apply target lock",
 			"database", database,
@@ -268,15 +268,15 @@ func acquireApplyTargetLockConn(ctx context.Context, db *rebindDB, locker namedl
 	return conn, lockName, nil
 }
 
-func releaseApplyTargetLockConn(ctx context.Context, locker namedlock.Locker, conn *sql.Conn, lockName, operation string) {
+func releaseApplyTargetLockConn(ctx context.Context, locker namedlock.Locker, conn *rebindConn, lockName, operation string) {
 	if conn == nil {
 		return
 	}
 
 	releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), applyTargetLockReleaseTimeout)
 	defer cancel()
 
-	released, err := locker.Release(releaseCtx, conn, lockName)
+	released, err := locker.Release(releaseCtx, conn.raw(), lockName)
 	switch {
 	case err != nil:
 		slog.WarnContext(releaseCtx, "failed to release apply target lock; discarding connection",
@@ -297,7 +297,7 @@ func releaseApplyTargetLockConn(ctx context.Context, locker namedlock.Locker, co
 // discardApplyTargetLockConn marks the pinned connection as bad so the pool
 // destroys it instead of reusing a session whose advisory-lock state is
 // uncertain.
-func discardApplyTargetLockConn(ctx context.Context, conn *sql.Conn, lockName, operation string) {
+func discardApplyTargetLockConn(ctx context.Context, conn *rebindConn, lockName, operation string) {
 	if rawErr := conn.Raw(func(any) error { return driver.ErrBadConn }); rawErr != nil && !errors.Is(rawErr, driver.ErrBadConn) {
 		slog.WarnContext(ctx, "failed to discard apply target lock connection",
 			"operation", operation,
@@ -306,7 +306,7 @@ func discardApplyTargetLockConn(ctx context.Context, conn *sql.Conn, lockName, o
 	}
 }
 
-func closeApplyTargetLockConn(ctx context.Context, conn *sql.Conn, lockName, operation string) {
+func closeApplyTargetLockConn(ctx context.Context, conn *rebindConn, lockName, operation string) {
 	if err := conn.Close(); err != nil {
 		slog.WarnContext(ctx, "failed to close apply target lock connection",
 			"operation", operation,
@@ -339,7 +339,7 @@ func dedupeDeployments(deployments []string) []string {
 // apply_operations rows. It is used to build the full target set of an existing
 // apply being started/resumed/reclaimed so the overlap check covers every
 // deployment the apply owns, not just the parent's primary deployment.
-func operationDeploymentsForApply(ctx context.Context, tx *sql.Tx, applyID int64) ([]string, error) {
+func operationDeploymentsForApply(ctx context.Context, tx *rebindTx, applyID int64) ([]string, error) {
 	rows, err := tx.QueryContext(ctx, `SELECT deployment FROM apply_operations WHERE apply_id = ?`, applyID)
 	if err != nil {
 		return nil, fmt.Errorf("list operation deployments for apply %d: %w", applyID, err)
@@ -374,7 +374,7 @@ func operationDeploymentsForApply(ctx context.Context, tx *sql.Tx, applyID int64
 // is matched against both the parent applies.deployment (the primary) and the
 // apply_operations.deployment rows, so single-operation applies (where the two
 // are equal) behave exactly as before.
-func checkNoActiveApplyForTargets(ctx context.Context, tx *sql.Tx, database, dbType, environment string, deployments []string, excludeApplyID int64) error {
+func checkNoActiveApplyForTargets(ctx context.Context, tx *rebindTx, database, dbType, environment string, deployments []string, excludeApplyID int64) error {
 	deployments = dedupeDeployments(deployments)
 	if len(deployments) == 0 {
 		return fmt.Errorf("active apply target requires at least one deployment for %s/%s/%s", database, dbType, environment)
@@ -569,7 +569,7 @@ func (s *applyStore) CreateWithTasks(ctx context.Context, apply *storage.Apply,
 	return s.CreateWithTasksAndOperations(ctx, apply, tasks, nil)
 }
 
-type applyCreateWriter func(ctx context.Context, tx *sql.Tx, apply *storage.Apply, applyID int64) error
+type applyCreateWriter func(ctx context.Context, tx *rebindTx, apply *storage.Apply, applyID int64) error
 
 // CreateWithTasksAndOperations is the unified atomic apply-create path: it
 // inserts the applies row, the initial tasks, and (optionally) the
@@ -582,7 +582,7 @@ func (s *applyStore) CreateWithTasksAndOperations(ctx context.Context, apply *st
 	for _, op := range operations {
 		deployments = append(deployments, op.Deployment)
 	}
-	return s.createWithRows(ctx, apply, opName, deployments, func(ctx context.Context, tx *sql.Tx, apply *storage.Apply, applyID int64) error {
+	return s.createWithRows(ctx, apply, opName, deployments, func(ctx context.Context, tx *rebindTx, apply *storage.Apply, applyID int64) error {
 		return insertApplyTasksAndOperations(ctx, tx, s.identity, apply, applyID, tasks, operations)
 	})
 }
@@ -596,7 +596,7 @@ func (s *applyStore) CreateWithGroupedOperations(ctx context.Context, apply *sto
 			deployments = append(deployments, group.Operation.Deployment)
 		}
 	}
-	return s.createWithRows(ctx, apply, opName, deployments, func(ctx context.Context, tx *sql.Tx, apply *storage.Apply, applyID int64) error {
+	return s.createWithRows(ctx, apply, opName, deployments, func(ctx context.Context, tx *rebindTx, apply *storage.Apply, applyID int64) error {
 		return insertApplyGroupedOperations(ctx, tx, s.identity, apply, applyID, groups)
 	})
 }
@@ -666,7 +666,7 @@ func (s *applyStore) createWithRows(ctx context.Context, apply *storage.Apply, o
 // same transaction that inserts the apply. This closes the gap where a
 // same-owner rollback can replace pending_plan_id after freshness validation
 // but before the forward apply becomes durable.
-func verifyExpectedLockIntent(ctx context.Context, tx *sql.Tx, apply *storage.Apply) error {
+func verifyExpectedLockIntent(ctx context.Context, tx *rebindTx, apply *storage.Apply) error {
 	if apply.ExpectedLockOwner == "" {
 		if apply.ExpectedPendingPlanID != "" {
 			return fmt.Errorf("verify lock intent for %s/%s: expected pending plan ID set without an expected lock owner", apply.Database, apply.DatabaseType)
@@ -695,7 +695,7 @@ func verifyExpectedLockIntent(ctx context.Context, tx *sql.Tx, apply *storage.Ap
 	return nil
 }
 
-func insertApplyTasksAndOperations(ctx context.Context, tx *sql.Tx, identity identityInserter, apply *storage.Apply, applyID int64, tasks []*storage.Task, operations []*storage.ApplyOperation) error {
+func insertApplyTasksAndOperations(ctx context.Context, tx *rebindTx, identity identityInserter, apply *storage.Apply, applyID int64, tasks []*storage.Task, operations []*storage.ApplyOperation) error {
 	// Operations are inserted BEFORE tasks so each task can be persisted
 	// with the apply_operation_id of the operation it belongs to. Today
 	// the config layer hard-blocks multi-entry deployments so there is
@@ -771,7 +771,7 @@ func insertApplyTasksAndOperations(ctx context.Context, tx *sql.Tx, identity ide
 	return nil
 }
 
-func insertApplyGroupedOperations(ctx context.Context, tx *sql.Tx, identity identityInserter, apply *storage.Apply, applyID int64, groups []*storage.ApplyOperationWithTasks) error {
+func insertApplyGroupedOperations(ctx context.Context, tx *rebindTx, identity identityInserter, apply *storage.Apply, applyID int64, groups []*storage.ApplyOperationWithTasks) error {
 	if len(groups) == 0 {
 		return fmt.Errorf("create apply %s: grouped operations are empty", apply.ApplyIdentifier)
 	}
@@ -1746,7 +1746,7 @@ func (o claimOutcome) claimedAndComplete() bool {
 // window between the re-check and the commit. When another active apply owns the
 // target the claim is refused and the pending start control request is failed so
 // the operator sees why.
-func persistApplyClaim(ctx context.Context, db *rebindDB, locker namedlock.Locker, tx *sql.Tx, apply *storage.Apply, owner string) (claimOutcome, error) {
+func persistApplyClaim(ctx context.Context, db *rebindDB, locker namedlock.Locker, tx *rebindTx, apply *storage.Apply, owner string) (claimOutcome, error) {
 	leaseToken := uuid.NewString()
 	leaseAcquiredAt := time.Now()
 	apply.LeaseOwner = owner
@@ -1784,7 +1784,7 @@ func persistApplyClaim(ctx context.Context, db *rebindDB, locker namedlock.Locke
 // invariant, then either transition+commit or refuse+commit. It commits inside
 // the lock so a concurrent create cannot add a second active apply between the
 // re-check and the commit. The lock is released only after the commit.
-func claimStoppedApplyUnderTargetLock(ctx context.Context, db *rebindDB, locker namedlock.Locker, tx *sql.Tx, apply *storage.Apply, owner, leaseToken string) (claimOutcome, error) {
+func claimStoppedApplyUnderTargetLock(ctx context.Context, db *rebindDB, locker namedlock.Locker, tx *rebindTx, apply *storage.Apply, owner, leaseToken string) (claimOutcome, error) {
 	database, dbType, environment, deployment, err := applyTargetForUpdate(ctx, tx, apply)
 	if err != nil {
 		return claimLostRace, err
@@ -1839,7 +1839,7 @@ func isStartingClaim(applyState string) bool {
 // landing between the SELECT and this UPDATE; a zero rows-affected result means
 // another driver already moved the row, so the caller backs off cleanly. Reports
 // false on that lost race.
-func transitionClaimToState(ctx context.Context, tx *sql.Tx, apply *storage.Apply, targetState, owner, leaseToken string) (bool, error) {
+func transitionClaimToState(ctx context.Context, tx *rebindTx, apply *storage.Apply, targetState, owner, leaseToken string) (bool, error) {
 	result, err := tx.ExecContext(ctx, `
 		UPDATE applies
 		SET state = ?, updated_at = NOW(),
@@ -1871,7 +1871,7 @@ func transitionClaimToState(ctx context.Context, tx *sql.Tx, apply *storage.Appl
 // target, and commits that failure inside tx. The start was accepted but cannot
 // be honored, so leaving the request pending would silently strand it; failing
 // it surfaces the reason to the operator. Returns claimRefusedActiveTarget.
-func refuseStoppedClaimForActiveTarget(ctx context.Context, tx *sql.Tx, apply *storage.Apply, owner, database, dbType, environment string) (claimOutcome, error) {
+func refuseStoppedClaimForActiveTarget(ctx context.Context, tx *rebindTx, apply *storage.Apply, owner, database, dbType, environment string) (claimOutcome, error) {
 	reason := fmt.Sprintf("start refused: another active apply exists for %s/%s/%s", database, dbType, environment)
 	if err := failPendingStartControlRequestTx(ctx, tx, apply.ID, reason); err != nil {
 		return claimLostRace, fmt.Errorf("fail pending start control request for apply %d (%s) on %s/%s/%s: %w", apply.ID, apply.ApplyIdentifier, database, dbType, environment, err)
@@ -1893,7 +1893,7 @@ func refuseStoppedClaimForActiveTarget(ctx context.Context, tx *sql.Tx, apply *s
 // an apply failed inside the supplied claim transaction so the refusal and the
 // failed request commit atomically. Mirrors controlRequestStore.FailPending's
 // SQL; this path holds no apply lease, so it is unguarded by lease token.
-func failPendingStartControlRequestTx(ctx context.Context, tx *sql.Tx, applyID int64, reason string) error {
+func failPendingStartControlRequestTx(ctx context.Context, tx *rebindTx, applyID int64, reason string) error {
 	_, err := tx.ExecContext(ctx, `
 		UPDATE apply_control_requests
 		SET status = ?, error_message = ?, completed_at = COALESCE(completed_at, NOW()), updated_at = NOW()
@@ -2199,7 +2199,7 @@ func (s *applyStore) ReapplyFailed(ctx context.Context, applyID int64) (*storage
 	return apply, nil
 }
 
-func rejectNonReappliableOperations(ctx context.Context, tx *sql.Tx, apply *storage.Apply) error {
+func rejectNonReappliableOperations(ctx context.Context, tx *rebindTx, apply *storage.Apply) error {
 	var count int
 	if err := tx.QueryRowContext(ctx, `
 		SELECT COUNT(*)

pkg/storage/internal/sqlstore/apply_operations.go

@@ -1563,7 +1563,7 @@ func (s *applyOperationStore) ReapStranded(ctx context.Context, limit int) ([]*s
 
 	// Do not wait for the lock: whoever holds it is doing this pass's work, and
 	// this instance's next tick is soon enough.
-	acquired, err := s.locker.Acquire(ctx, conn, strandedReaperLockName, 0)
+	acquired, err := s.locker.Acquire(ctx, conn.raw(), strandedReaperLockName, 0)
 	if err != nil {
 		return nil, fmt.Errorf("acquire stranded reaper lock: %w", err)
 	}
@@ -1574,7 +1574,7 @@ func (s *applyOperationStore) ReapStranded(ctx context.Context, limit int) ([]*s
 		// A held lock parks every instance's reaper until this session is
 		// retired, so the two ways it can survive the pass are reported apart:
 		// the release errored, or it ran and reported the lock was not held.
-		released, err := s.locker.Release(context.WithoutCancel(ctx), conn, strandedReaperLockName)
+		released, err := s.locker.Release(context.WithoutCancel(ctx), conn.raw(), strandedReaperLockName)
 		if err != nil {
 			slog.WarnContext(ctx, "failed to release the stranded reaper lock; reapers stay blocked until this session is retired",
 				"lock", strandedReaperLockName, "error", err)

pkg/storage/internal/sqlstore/apply_operations_test.go

@@ -4847,15 +4847,15 @@ func TestApplyOperationStore_ReapStranded_ElectsOneReaperPerPass(t *testing.T) {
 	other := New(testDB)
 	held, err := other.applyOperations.db.Conn(ctx)
 	require.NoError(t, err)
-	acquired, err := namedlock.MySQL{}.Acquire(ctx, held, strandedReaperLockName, 0)
+	acquired, err := namedlock.MySQL{}.Acquire(ctx, held.raw(), strandedReaperLockName, 0)
 	require.NoError(t, err)
 	require.True(t, acquired, "the stand-in instance should take the reaper lock")
 
 	_, err = store.ApplyOperations().ReapStranded(ctx, 10)
 	require.ErrorIs(t, err, storage.ErrStrandedReaperBusy, "a second reaper steps aside rather than re-scanning")
 	assertApplyOperationState(t, store, op.ID, state.ApplyOperation.Pending)
 
-	released, err := namedlock.MySQL{}.Release(ctx, held, strandedReaperLockName)
+	released, err := namedlock.MySQL{}.Release(ctx, held.raw(), strandedReaperLockName)
 	require.NoError(t, err)
 	require.True(t, released)
 	require.NoError(t, held.Close())

pkg/storage/internal/sqlstore/control_requests.go

@@ -193,7 +193,7 @@ func (s *controlRequestStore) FailPending(ctx context.Context, applyID int64, op
 	return nil
 }
 
-func (s *controlRequestStore) getByIDForUpdate(ctx context.Context, tx *sql.Tx, id int64) (*storage.ApplyControlRequest, error) {
+func (s *controlRequestStore) getByIDForUpdate(ctx context.Context, tx *rebindTx, id int64) (*storage.ApplyControlRequest, error) {
 	row := tx.QueryRowContext(ctx, `
 		SELECT `+controlRequestColumns+`
 		FROM apply_control_requests
@@ -209,7 +209,7 @@ func (s *controlRequestStore) getByIDForUpdate(ctx context.Context, tx *sql.Tx,
 
 func (s *controlRequestStore) getByApplyOperationForUpdate(
 	ctx context.Context,
-	tx *sql.Tx,
+	tx *rebindTx,
 	applyID int64,
 	operation storage.ControlOperation,
 ) (*storage.ApplyControlRequest, error) {