From b71d7c1e78be25ceb17dfd55396277d1723b3aed Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Thu, 15 Oct 2026 13:42:21 +0200 Subject: [PATCH] refactor(config): share plan writes with caller transactions --- internal/config/store_postgres.go | 73 +-------- internal/config/store_postgres_plan_writes.go | 82 ++++++++++ ...e_postgres_plan_writes_integration_test.go | 153 ++++++++++++++++++ 3 files changed, 238 insertions(+), 70 deletions(-) create mode 100644 internal/config/store_postgres_plan_writes.go create mode 100644 internal/config/store_postgres_plan_writes_integration_test.go diff --git a/internal/config/store_postgres.go b/internal/config/store_postgres.go index 2f3d19be..5dc57909 100644 --- a/internal/config/store_postgres.go +++ b/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 { diff --git a/internal/config/store_postgres_plan_writes.go b/internal/config/store_postgres_plan_writes.go new file mode 100644 index 00000000..6fb99f1c --- /dev/null +++ b/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 +} diff --git a/internal/config/store_postgres_plan_writes_integration_test.go b/internal/config/store_postgres_plan_writes_integration_test.go new file mode 100644 index 00000000..126e1112 --- /dev/null +++ b/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") + } +}