diff --git a/internal/config/store_postgres_recommendation_source_account_test.go b/internal/config/store_postgres_recommendation_source_account_test.go new file mode 100644 index 00000000..4cf2ab80 --- /dev/null +++ b/internal/config/store_postgres_recommendation_source_account_test.go @@ -0,0 +1,50 @@ +//go:build integration + +package config_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/LeanerCloud/cloud-commitments-platform/internal/config" + "github.com/stretchr/testify/require" +) + +func TestPostgresRecommendationSourceEncodingRollback(t *testing.T) { + ctx := context.Background() + store, cleanup := setupRecommendationsStore(ctx, t) + defer cleanup() + at := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + seed := awsRec("seed", "rds", "region", "seed", 10) + seed.SourceAccountID = new(" SourceAbC/é\t") + require.NoError(t, store.UpsertRecommendations(ctx, at, []config.RecommendationRecord{seed}, nil)) + require.NoError(t, store.SetRecommendationsCollectionError(ctx, "preserved")) + before, err := store.GetRecommendationsFreshness(ctx) + require.NoError(t, err) + for _, bad := range []string{"opaque\x00source", string([]byte{'x', 0xff, 'y'})} { + for _, replace := range []bool{false, true} { + recs := make([]config.RecommendationRecord, 501) + for i := range recs { + recs[i] = awsRec(fmt.Sprint(i), "rds", "region", fmt.Sprint(i), 20) + recs[i].SourceAccountID = new("valid") + } + recs[500].SourceAccountID = &bad + if replace { + err = store.ReplaceRecommendations(ctx, at.Add(time.Hour), recs) + } else { + err = store.UpsertRecommendations(ctx, at.Add(time.Hour), recs, []config.SuccessfulCollect{{Provider: "aws"}}) + } + require.EqualError(t, err, "invalid source_account_id encoding") + rows, listErr := store.ListStoredRecommendations(ctx, config.RecommendationFilter{}) + require.NoError(t, listErr) + require.Len(t, rows, 1) + require.Equal(t, "seed", rows[0].ID) + require.Equal(t, seed.SourceAccountID, rows[0].SourceAccountID) + after, freshnessErr := store.GetRecommendationsFreshness(ctx) + require.NoError(t, freshnessErr) + require.Equal(t, before, after) + } + } +} diff --git a/internal/config/store_postgres_recommendations.go b/internal/config/store_postgres_recommendations.go index f3d8b93b..f4ce3672 100644 --- a/internal/config/store_postgres_recommendations.go +++ b/internal/config/store_postgres_recommendations.go @@ -8,6 +8,7 @@ import ( "math" "strings" "time" + "unicode/utf8" "github.com/LeanerCloud/cloud-commitments-go/pkg/logging" "github.com/google/uuid" @@ -193,6 +194,9 @@ func insertRecommendationsBatch(ctx context.Context, tx pgx.Tx, collectedAt time for i := range recs { rec := recs[i] + if rec.SourceAccountID != nil && (!utf8.ValidString(*rec.SourceAccountID) || strings.ContainsRune(*rec.SourceAccountID, 0)) { + return fmt.Errorf("invalid source_account_id encoding") + } payload, err := json.Marshal(rec) if err != nil { return fmt.Errorf("failed to marshal recommendation %d: %w", i, err) diff --git a/internal/config/types.go b/internal/config/types.go index 60d11945..3caf5d3c 100644 --- a/internal/config/types.go +++ b/internal/config/types.go @@ -630,6 +630,7 @@ type RecommendationRecord struct { PurchaseID string `json:"purchase_id,omitempty" dynamodbav:"purchase_id,omitempty"` Error string `json:"error,omitempty" dynamodbav:"error,omitempty"` CloudAccountID *string `json:"cloud_account_id,omitempty" dynamodbav:"cloud_account_id,omitempty"` + SourceAccountID *string `json:"source_account_id,omitempty" dynamodbav:"source_account_id,omitempty"` // SuppressedCount is the cumulative count already committed against // this recommendation's 6-tuple (account, provider, service, region, // resource_type, engine) within the active grace window. The diff --git a/internal/scheduler/recommendation_completeness_integration_test.go b/internal/scheduler/recommendation_completeness_integration_test.go index 126429d2..1902887c 100644 --- a/internal/scheduler/recommendation_completeness_integration_test.go +++ b/internal/scheduler/recommendation_completeness_integration_test.go @@ -30,6 +30,7 @@ import ( type completenessHTTP struct { resourceType string + accounts []any details []string fallback []string failScope bool @@ -80,14 +81,18 @@ func (f *completenessHTTP) Do(req *http.Request) (*http.Response, error) { if input.LookbackPeriodInDays == "THIRTY_DAYS" { quantities = f.fallback } - for _, quantity := range quantities { - details = append(details, map[string]any{ + for i, quantity := range quantities { + detail := map[string]any{ "RecommendedNumberOfInstancesToPurchase": quantity, "EstimatedMonthlySavingsAmount": "10", "EstimatedMonthlyOnDemandCost": "30", "InstanceDetails": map[string]any{"RDSInstanceDetails": map[string]any{ "InstanceType": f.resourceType, "Region": "us-east-1", "DeploymentOption": "Single-AZ", }}, - }) + } + if i < len(f.accounts) { + detail["AccountId"] = f.accounts[i] + } + details = append(details, detail) } if f.failScope && call == 1 { status = http.StatusBadRequest diff --git a/internal/scheduler/recommendation_source_account_integration_test.go b/internal/scheduler/recommendation_source_account_integration_test.go new file mode 100644 index 00000000..d504ebd6 --- /dev/null +++ b/internal/scheduler/recommendation_source_account_integration_test.go @@ -0,0 +1,67 @@ +//go:build integration + +package scheduler + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/LeanerCloud/cloud-commitments-platform/internal/config" + "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/migrations" + "github.com/LeanerCloud/cloud-commitments-platform/internal/database/postgres/testhelpers" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestAWSRecommendationSourceAccountPersistence(t *testing.T) { + ctx := context.Background() + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + t.Cleanup(func() { container.Cleanup(ctx) }) + migrationPath, err := filepath.Abs("../database/postgres/migrations") + require.NoError(t, err) + require.NoError(t, migrations.RunMigrations(ctx, container.DB.Pool(), migrationPath, "", "")) + store := config.NewPostgresStore(container.DB) + collector := uuid.NewString() + require.NoError(t, store.CreateCloudAccount(ctx, &config.CloudAccount{ + ID: collector, Name: "collector", Provider: "aws", ExternalID: "collector-account", + })) + fixture := &completenessHTTP{resourceType: "db.t3.medium", details: []string{"2", "2", "2", "2", "2"}, + accounts: []any{" SourceAbC/é\t", "sourceabc/e\u0301", nil, ""}} + client, err := completenessProvider(t, fixture).GetRecommendationsClient(ctx) + require.NoError(t, err) + raw, err := client.GetAllRecommendations(ctx) + require.NoError(t, err) + require.Len(t, raw, 30) + s := &Scheduler{} + records := s.tagAccount(s.convertRecommendations(raw, "aws"), collector) + require.Len(t, records, 30) + require.Equal(t, new(" SourceAbC/é\t"), records[0].SourceAccountID) + require.Equal(t, new("sourceabc/e\u0301"), records[1].SourceAccountID) + require.Nil(t, records[2].SourceAccountID) + require.Nil(t, records[3].SourceAccountID) + require.Nil(t, records[4].SourceAccountID) + require.Equal(t, &collector, records[0].CloudAccountID) + at := time.Now().UTC().Truncate(time.Second) + require.NoError(t, store.UpsertRecommendations(ctx, at, records[:1], nil)) + rows, err := store.ListStoredRecommendations(ctx, config.RecommendationFilter{}) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, records[0].SourceAccountID, rows[0].SourceAccountID) + require.Equal(t, &collector, rows[0].CloudAccountID) + var key string + require.NoError(t, container.DB.Pool().QueryRow(ctx, "SELECT source_account_key FROM recommendations").Scan(&key)) + require.Equal(t, *records[0].SourceAccountID, key) + require.Error(t, store.UpsertRecommendations(ctx, at.Add(time.Hour), records[:2], nil)) + rows, err = store.ListStoredRecommendations(ctx, config.RecommendationFilter{}) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, records[0].SourceAccountID, rows[0].SourceAccountID) + for _, req := range fixture.requests { + if req.Operation == "AWSInsightsIndexService.GetReservationPurchaseRecommendation" { + require.Equal(t, "LINKED", req.AccountScope) + } + } +} diff --git a/internal/scheduler/recommendation_source_account_test.go b/internal/scheduler/recommendation_source_account_test.go new file mode 100644 index 00000000..3d9d4694 --- /dev/null +++ b/internal/scheduler/recommendation_source_account_test.go @@ -0,0 +1,54 @@ +package scheduler + +import ( + "encoding/json" + "testing" + + "github.com/LeanerCloud/cloud-commitments-go/pkg/common" + "github.com/LeanerCloud/cloud-commitments-platform/internal/config" + "github.com/stretchr/testify/require" +) + +func TestConvertRecommendationsSourceAccount(t *testing.T) { + t.Parallel() + for _, cloud := range []string{"aws", "azure", "gcp"} { + t.Run(cloud, func(t *testing.T) { + s := &Scheduler{} + input := []common.Recommendation{ + {Account: " SourceAbC/é\t", Service: common.ServiceEC2, Term: "1yr"}, + {Account: "sourceabc/e\u0301", Service: common.ServiceEC2, Term: "1yr"}, + {Service: common.ServiceEC2, Term: "1yr"}, + } + records := s.convertRecommendations(input, cloud) + require.Len(t, records, 3) + input[0].Account = "changed" + tagged := s.tagAccount(records, "11111111-1111-4111-8111-111111111111") + require.Len(t, tagged, 3) + for i, want := range []string{" SourceAbC/é\t", "sourceabc/e\u0301", ""} { + raw, err := json.Marshal(tagged[i]) + require.NoError(t, err) + var wire map[string]any + require.NoError(t, json.Unmarshal(raw, &wire)) + require.Equal(t, "11111111-1111-4111-8111-111111111111", wire["cloud_account_id"]) + if want == "" { + require.NotContains(t, wire, "source_account_id") + } else { + require.Equal(t, want, wire["source_account_id"]) + } + } + }) + } +} + +func TestRecommendationSourceAccountLegacyJSON(t *testing.T) { + t.Parallel() + for _, raw := range []string{`{}`, `{"source_account_id":null}`} { + var record config.RecommendationRecord + require.NoError(t, json.Unmarshal([]byte(raw), &record)) + encoded, err := json.Marshal(record) + require.NoError(t, err) + var wire map[string]any + require.NoError(t, json.Unmarshal(encoded, &wire)) + require.NotContains(t, wire, "source_account_id") + } +} diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index f3fa80d6..bb84f48d 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -1639,7 +1639,12 @@ func (s *Scheduler) convertRecommendations(recs []common.Recommendation, provide continue } + var sourceAccountID *string + if rec.Account != "" { + sourceAccountID = new(rec.Account) + } record := config.RecommendationRecord{ + SourceAccountID: sourceAccountID, Provider: providerName, Service: string(rec.Service), Region: rec.Region,