refactor(config): share plan writes with caller transactions

internal/config/store_postgres.go

@@ -557,60 +557,7 @@ func (s *PostgresStore) ListServiceConfigs(ctx context.Context) ([]ServiceConfig
 
 // CreatePurchasePlan creates a new purchase plan.
 func (s *PostgresStore) CreatePurchasePlan(ctx context.Context, plan *PurchasePlan) error {
-	// Generate UUID if not provided
-	if plan.ID == "" {
-		plan.ID = uuid.New().String()
-	}
-
-	// Set timestamps
-	now := time.Now()
-	plan.CreatedAt = now
-	plan.UpdatedAt = now
-
-	// Marshal services and ramp_schedule to JSONB
-	servicesJSON, err := json.Marshal(plan.Services)
-	if err != nil {
-		return fmt.Errorf("failed to marshal services: %w", err)
-	}
-
-	rampScheduleJSON, err := json.Marshal(plan.RampSchedule)
-	if err != nil {
-		return fmt.Errorf("failed to marshal ramp_schedule: %w", err)
-	}
-
-	query := `
-		INSERT INTO purchase_plans (
-			id, name, enabled, auto_purchase, notification_days_before,
-			services, ramp_schedule, created_at, updated_at,
-			next_execution_date, last_execution_date, last_notification_sent,
-			max_step_commitment
-		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
-		RETURNING updated_at
-	`
-
-	var updatedAt time.Time
-	err = s.db.QueryRow(ctx, query,
-		plan.ID,
-		plan.Name,
-		plan.Enabled,
-		plan.AutoPurchase,
-		plan.NotificationDaysBefore,
-		servicesJSON,
-		rampScheduleJSON,
-		plan.CreatedAt,
-		plan.UpdatedAt,
-		plan.NextExecutionDate,
-		plan.LastExecutionDate,
-		plan.LastNotificationSent,
-		plan.MaxStepCommitment,
-	).Scan(&updatedAt)
-
-	if err != nil {
-		return fmt.Errorf("failed to create purchase plan: %w", err)
-	}
-	plan.UpdatedAt = updatedAt
-
-	return nil
+	return s.createPurchasePlan(ctx, s.db, plan)
 }
 
 const purchasePlanSelectCols = `
@@ -4123,22 +4070,8 @@ func (s *PostgresStore) SetPlanAccounts(ctx context.Context, planID string, acco
 		}
 	}()
 
-	err = s.validatePlanAccountProvidersTx(ctx, tx, planID, accountIDs)
-	if err != nil {
-		return err
-	}
-
-	if _, err = tx.Exec(ctx, `DELETE FROM plan_accounts WHERE plan_id = $1`, planID); err != nil {
-		return fmt.Errorf("failed to clear plan accounts: %w", err)
-	}
-
-	for _, accountID := range accountIDs {
-		if _, err = tx.Exec(ctx,
-			`INSERT INTO plan_accounts (plan_id, account_id) VALUES ($1, $2) ON CONFLICT DO NOTHING`,
-			planID, accountID,
-		); err != nil {
-			return fmt.Errorf("failed to insert plan account: %w", err)
-		}
+	if setErr := s.setPlanAccountsTx(ctx, tx, planID, accountIDs); setErr != nil {
+		return setErr
 	}
 
 	if err = tx.Commit(ctx); err != nil {

internal/config/store_postgres_plan_writes.go

@@ -0,0 +1,82 @@
+package config
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+)
+
+func (s *PostgresStore) createPurchasePlan(ctx context.Context, executor globalConfigExecutor, plan *PurchasePlan) error {
+	if plan.ID == "" {
+		plan.ID = uuid.New().String()
+	}
+
+	now := time.Now()
+	plan.CreatedAt = now
+	plan.UpdatedAt = now
+
+	servicesJSON, err := json.Marshal(plan.Services)
+	if err != nil {
+		return fmt.Errorf("failed to marshal services: %w", err)
+	}
+
+	rampScheduleJSON, err := json.Marshal(plan.RampSchedule)
+	if err != nil {
+		return fmt.Errorf("failed to marshal ramp_schedule: %w", err)
+	}
+
+	query := `
+		INSERT INTO purchase_plans (
+			id, name, enabled, auto_purchase, notification_days_before,
+			services, ramp_schedule, created_at, updated_at,
+			next_execution_date, last_execution_date, last_notification_sent,
+			max_step_commitment
+		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
+		RETURNING updated_at
+	`
+
+	var updatedAt time.Time
+	err = executor.QueryRow(ctx, query,
+		plan.ID,
+		plan.Name,
+		plan.Enabled,
+		plan.AutoPurchase,
+		plan.NotificationDaysBefore,
+		servicesJSON,
+		rampScheduleJSON,
+		plan.CreatedAt,
+		plan.UpdatedAt,
+		plan.NextExecutionDate,
+		plan.LastExecutionDate,
+		plan.LastNotificationSent,
+		plan.MaxStepCommitment,
+	).Scan(&updatedAt)
+
+	if err != nil {
+		return fmt.Errorf("failed to create purchase plan: %w", err)
+	}
+	plan.UpdatedAt = updatedAt
+	return nil
+}
+
+func (s *PostgresStore) setPlanAccountsTx(ctx context.Context, tx pgx.Tx, planID string, accountIDs []string) error {
+	if err := s.validatePlanAccountProvidersTx(ctx, tx, planID, accountIDs); err != nil {
+		return err
+	}
+	if _, err := tx.Exec(ctx, `DELETE FROM plan_accounts WHERE plan_id = $1`, planID); err != nil {
+		return fmt.Errorf("failed to clear plan accounts: %w", err)
+	}
+	for _, accountID := range accountIDs {
+		if _, err := tx.Exec(ctx,
+			`INSERT INTO plan_accounts (plan_id, account_id) VALUES ($1, $2) ON CONFLICT DO NOTHING`,
+			planID, accountID,
+		); err != nil {
+			return fmt.Errorf("failed to insert plan account: %w", err)
+		}
+	}
+	return nil
+}

internal/config/store_postgres_plan_writes_integration_test.go

@@ -0,0 +1,153 @@
+//go:build integration
+
+package config
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPlanWriteBodiesOuterTransaction(t *testing.T) {
+	ctx := context.Background()
+	store := setupRampStepStore(ctx, t)
+	account := &CloudAccount{Name: "transaction account", Provider: "aws", ExternalID: uuid.NewString()}
+	require.NoError(t, store.CreateCloudAccount(ctx, account))
+	abort := errors.New("abort composed plan write")
+	for _, commit := range []bool{false, true} {
+		t.Run(map[bool]string{false: "rollback", true: "commit"}[commit], func(t *testing.T) {
+			stamp := time.Now().UTC().Truncate(time.Microsecond).Add(48 * time.Hour)
+			cap := 25.5
+			plan := &PurchasePlan{
+				Name: "composed plan", Enabled: true, NotificationDaysBefore: 1,
+				Services:     map[string]ServiceConfig{"aws/ec2": {Provider: "aws", Service: "ec2", Enabled: true, Term: 1}},
+				RampSchedule: PresetRampSchedules["weekly-25pct"], NextExecutionDate: &stamp,
+				LastExecutionDate: &stamp, LastNotificationSent: &stamp, MaxStepCommitment: &cap,
+			}
+			err := store.WithTx(ctx, func(tx pgx.Tx) error {
+				if err := store.createPurchasePlan(ctx, tx, plan); err != nil {
+					return err
+				}
+				if err := store.setPlanAccountsTx(ctx, tx, plan.ID, []string{account.ID, account.ID}); err != nil {
+					return err
+				}
+				var plans, accounts int
+				require.NoError(t, tx.QueryRow(ctx, `SELECT count(*) FROM purchase_plans WHERE id=$1`, plan.ID).Scan(&plans))
+				require.NoError(t, tx.QueryRow(ctx, `SELECT count(*) FROM plan_accounts WHERE plan_id=$1`, plan.ID).Scan(&accounts))
+				require.Equal(t, 1, plans)
+				require.Equal(t, 1, accounts)
+				if !commit {
+					return abort
+				}
+				return nil
+			})
+			if !commit {
+				require.ErrorIs(t, err, abort)
+				var plans, accounts int
+				require.NoError(t, store.db.QueryRow(ctx, `SELECT count(*) FROM purchase_plans WHERE id=$1`, plan.ID).Scan(&plans))
+				require.NoError(t, store.db.QueryRow(ctx, `SELECT count(*) FROM plan_accounts WHERE plan_id=$1`, plan.ID).Scan(&accounts))
+				require.Zero(t, plans, "outer rollback must remove the new plan")
+				require.Zero(t, accounts, "outer rollback must remove the new associations")
+				return
+			}
+			require.NoError(t, err)
+			got, err := store.GetPurchasePlan(ctx, plan.ID)
+			require.NoError(t, err)
+			require.NotEmpty(t, plan.ID)
+			require.Equal(t, plan.Services, got.Services)
+			require.Equal(t, plan.RampSchedule, got.RampSchedule)
+			require.Equal(t, plan.Name, got.Name)
+			require.Equal(t, plan.Enabled, got.Enabled)
+			require.Equal(t, plan.AutoPurchase, got.AutoPurchase)
+			require.Equal(t, plan.NotificationDaysBefore, got.NotificationDaysBefore)
+			require.Equal(t, plan.MaxStepCommitment, got.MaxStepCommitment)
+			require.WithinDuration(t, plan.CreatedAt, got.CreatedAt, time.Microsecond)
+			require.True(t, plan.UpdatedAt.Equal(got.UpdatedAt))
+			for _, date := range []*time.Time{got.NextExecutionDate, got.LastExecutionDate, got.LastNotificationSent} {
+				require.NotNil(t, date)
+				require.True(t, stamp.Equal(*date))
+			}
+			accounts, err := store.GetPlanAccounts(ctx, plan.ID)
+			require.NoError(t, err)
+			require.Len(t, accounts, 1)
+			require.Equal(t, account.ID, accounts[0].ID)
+			var baseline []byte
+			require.NoError(t, store.db.QueryRow(ctx, `SELECT ramp_baseline FROM purchase_plans WHERE id=$1`, plan.ID).Scan(&baseline))
+			require.Nil(t, baseline)
+			duplicate := *plan
+			duplicate.Name = "must not replace original"
+			var pgErr *pgconn.PgError
+			require.ErrorAs(t, store.CreatePurchasePlan(ctx, &duplicate), &pgErr)
+			require.Equal(t, "23505", pgErr.Code)
+			after, err := store.GetPurchasePlan(ctx, plan.ID)
+			require.NoError(t, err)
+			require.Equal(t, got, after)
+		})
+	}
+}
+
+func TestPlanWriteBodiesAssociationCompatibility(t *testing.T) {
+	ctx := context.Background()
+	store := setupRampStepStore(ctx, t)
+	plan := &PurchasePlan{ID: uuid.NewString(), Name: "existing plan", Services: map[string]ServiceConfig{"aws/ec2": {Provider: "aws", Service: "ec2", Enabled: true}}}
+	id := plan.ID
+	require.NoError(t, store.CreatePurchasePlan(ctx, plan))
+	require.Equal(t, id, plan.ID)
+	ids := make([]string, 0, 3)
+	for _, provider := range []string{"aws", "aws", "azure"} {
+		account := &CloudAccount{Name: provider, Provider: provider, ExternalID: uuid.NewString()}
+		require.NoError(t, store.CreateCloudAccount(ctx, account))
+		ids = append(ids, account.ID)
+	}
+	require.NoError(t, store.SetPlanAccounts(ctx, plan.ID, ids[:1]))
+	require.NoError(t, store.WithTx(ctx, func(tx pgx.Tx) error {
+		_, err := store.SetPlanRampBaselineTx(ctx, tx, plan.ID, PlanRampBaseline{baselineKey(t, baselineRec(ids[0], "m5.large", 7)): {Count: 7}})
+		return err
+	}))
+	_, err := store.db.Exec(ctx, `UPDATE purchase_plans SET services=services || '{"future":{"unknown":true}}'::jsonb, ramp_schedule='{"type":"weekly","current_step":2,"future":"retained"}'::jsonb WHERE id=$1`, plan.ID)
+	require.NoError(t, err)
+	var before string
+	require.NoError(t, store.db.QueryRow(ctx, `SELECT row_to_json(p)::text FROM purchase_plans p WHERE id=$1`, plan.ID).Scan(&before))
+	for _, useWrapper := range []bool{true, false} {
+		set := func(planID string, accounts []string) error {
+			if useWrapper {
+				return store.SetPlanAccounts(ctx, planID, accounts)
+			}
+			return store.WithTx(ctx, func(tx pgx.Tx) error { return store.setPlanAccountsTx(ctx, tx, planID, accounts) })
+		}
+		require.NoError(t, set(plan.ID, []string{ids[1], ids[1]}))
+		for _, invalid := range []struct {
+			planID   string
+			accounts []string
+			missing  bool
+		}{
+			{plan.ID, []string{ids[2]}, false},
+			{plan.ID, []string{uuid.NewString()}, true},
+			{uuid.NewString(), nil, true},
+		} {
+			err := set(invalid.planID, invalid.accounts)
+			if invalid.missing {
+				require.ErrorIs(t, err, ErrNotFound)
+			} else {
+				require.ErrorContains(t, err, "plan provider mismatch:")
+			}
+			accounts, err := store.GetPlanAccounts(ctx, plan.ID)
+			require.NoError(t, err)
+			require.Len(t, accounts, 1)
+			require.Equal(t, ids[1], accounts[0].ID)
+		}
+		require.NoError(t, set(plan.ID, nil))
+		accounts, err := store.GetPlanAccounts(ctx, plan.ID)
+		require.NoError(t, err)
+		require.Empty(t, accounts)
+		var after string
+		require.NoError(t, store.db.QueryRow(ctx, `SELECT row_to_json(p)::text FROM purchase_plans p WHERE id=$1`, plan.ID).Scan(&after))
+		require.Equal(t, before, after, "associations must preserve raw plan, baseline and progress")
+	}
+}