feat(recommendations): preserve source account provenance

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)
+		}
+	}
+}

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)

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

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

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)
+		}
+	}
+}

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")
+	}
+}

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,