[stripe] Use attributionId in stripe queries

Author: easyCZ

components/usage/pkg/apiv1/billing.go

@@ -52,7 +52,7 @@ func (s *BillingService) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoic
 		return nil, status.Errorf(codes.Internal, "Failed to download usage report with ID: %s", in.GetReportId())
 	}
 
-	credits, err := s.creditSummaryForTeams(report.UsageRecords, in.GetReportId())
+	credits, err := s.creditSummary(report.UsageRecords, in.GetReportId())
 	if err != nil {
 		log.Log.WithError(err).Errorf("Failed to compute credit summary.")
 		return nil, status.Errorf(codes.InvalidArgument, "failed to compute credit summary")
@@ -74,16 +74,9 @@ func (s *BillingService) ReconcileInvoices(ctx context.Context, in *v1.Reconcile
 		return nil, status.Errorf(codes.Internal, "Failed to reconcile invoices.")
 	}
 
-	creditSummaryForTeams := map[string]stripe.CreditSummary{}
+	creditSummaryForTeams := map[db.AttributionID]stripe.CreditSummary{}
 	for _, balance := range balances {
-		entity, id := balance.AttributionID.Values()
-
-		// TODO: Support updating of user attribution IDs
-		if entity != db.AttributionEntity_Team {
-			continue
-		}
-
-		creditSummaryForTeams[id] = stripe.CreditSummary{
+		creditSummaryForTeams[balance.AttributionID] = stripe.CreditSummary{
 			Credits:  int64(math.Ceil(balance.CreditCents.ToCredits())),
 			ReportID: "no-report",
 		}
@@ -117,7 +110,7 @@ func (s *BillingService) FinalizeInvoice(ctx context.Context, in *v1.FinalizeInv
 		return nil, status.Errorf(codes.Internal, "Failed to retrieve subscription details from invoice.")
 	}
 
-	teamID, found := subscription.Metadata[stripe.TeamIDMetadataKey]
+	teamID, found := subscription.Metadata[stripe.AttributionIDMetadataKey]
 	if !found {
 		logger.Error("Failed to find teamID from subscription metadata.")
 		return nil, status.Errorf(codes.Internal, "Failed to extra teamID from Stripe subscription.")
@@ -196,29 +189,24 @@ func (s *BillingService) GetUpcomingInvoice(ctx context.Context, in *v1.GetUpcom
 	}, nil
 }
 
-func (s *BillingService) creditSummaryForTeams(sessions []db.WorkspaceInstanceUsage, reportID string) (map[string]stripe.CreditSummary, error) {
-	creditsPerTeamID := map[string]float64{}
+func (s *BillingService) creditSummary(sessions []db.WorkspaceInstanceUsage, reportID string) (map[db.AttributionID]stripe.CreditSummary, error) {
+	creditsPerAttributionID := map[db.AttributionID]float64{}
 
 	for _, session := range sessions {
 		if session.StartedAt.Before(s.billInstancesAfter) {
 			continue
 		}
 
-		entity, id := session.AttributionID.Values()
-		if entity != db.AttributionEntity_Team {
-			continue
-		}
-
-		if _, ok := creditsPerTeamID[id]; !ok {
-			creditsPerTeamID[id] = 0
+		if _, ok := creditsPerAttributionID[session.AttributionID]; !ok {
+			creditsPerAttributionID[session.AttributionID] = 0
 		}
 
-		creditsPerTeamID[id] += session.CreditsUsed
+		creditsPerAttributionID[session.AttributionID] += session.CreditsUsed
 	}
 
-	rounded := map[string]stripe.CreditSummary{}
-	for teamID, credits := range creditsPerTeamID {
-		rounded[teamID] = stripe.CreditSummary{
+	rounded := map[db.AttributionID]stripe.CreditSummary{}
+	for attributionID, credits := range creditsPerAttributionID {
+		rounded[attributionID] = stripe.CreditSummary{
 			Credits:  int64(math.Ceil(credits)),
 			ReportID: reportID,
 		}

components/usage/pkg/apiv1/billing_test.go

@@ -121,7 +121,7 @@ func TestCreditSummaryForTeams(t *testing.T) {
 	for _, s := range scenarios {
 		t.Run(s.Name, func(t *testing.T) {
 			svc := NewBillingService(&stripe.Client{}, s.BillSessionsAfter, &gorm.DB{}, nil)
-			actual, err := svc.creditSummaryForTeams(s.Sessions, reportID)
+			actual, err := svc.creditSummary(s.Sessions, reportID)
 			require.NoError(t, err)
 			require.Equal(t, s.Expected, actual)
 		})

components/usage/pkg/stripe/stripe.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/gitpod-io/gitpod/usage/pkg/db"
 	"os"
 	"strings"
 
@@ -17,8 +18,8 @@ import (
 )
 
 const (
-	ReportIDMetadataKey = "reportId"
-	TeamIDMetadataKey   = "teamId"
+	ReportIDMetadataKey      = "reportId"
+	AttributionIDMetadataKey = "attributionId"
 )
 
 type Client struct {
@@ -70,12 +71,12 @@ type CreditSummary struct {
 
 // UpdateUsage updates teams' Stripe subscriptions with usage data
 // `usageForTeam` is a map from team name to total workspace seconds used within a billing period.
-func (c *Client) UpdateUsage(ctx context.Context, creditsPerTeam map[string]CreditSummary) error {
-	teamIds := make([]string, 0, len(creditsPerTeam))
-	for k := range creditsPerTeam {
-		teamIds = append(teamIds, k)
+func (c *Client) UpdateUsage(ctx context.Context, creditsPerAttributionID map[db.AttributionID]CreditSummary) error {
+	attributionIDs := make([]db.AttributionID, 0, len(creditsPerAttributionID))
+	for k := range creditsPerAttributionID {
+		attributionIDs = append(attributionIDs, k)
 	}
-	queries := queriesForCustomersWithTeamIds(teamIds)
+	queries := queriesForCustomersWithAttributionIDs(attributionIDs)
 
 	for _, query := range queries {
 		log.Infof("Searching customers in Stripe with query: %q", query)
@@ -86,14 +87,21 @@ func (c *Client) UpdateUsage(ctx context.Context, creditsPerTeam map[string]Cred
 		}
 
 		for _, customer := range customers {
-			teamID := customer.Metadata["teamId"]
-			log.Infof("Found customer %q for teamId %q", customer.Name, teamID)
+			attributionIDRaw := customer.Metadata[AttributionIDMetadataKey]
+			log.Infof("Found customer %q for attribution ID %q", customer.Name, attributionIDRaw)
 
-			_, err := c.updateUsageForCustomer(ctx, customer, creditsPerTeam[teamID])
+			attributionID, err := db.ParseAttributionID(attributionIDRaw)
+			if err != nil {
+				log.WithError(err).Error("Failed to parse attribution ID from Stripe metadata.")
+				continue
+			}
+
+			_, err = c.updateUsageForCustomer(ctx, customer, creditsPerAttributionID[attributionID])
 			if err != nil {
 				log.WithField("customer_id", customer.ID).
 					WithField("customer_name", customer.Name).
 					WithField("subscriptions", customer.Subscriptions).
+					WithField("attribution_id", attributionID).
 					WithError(err).
 					Errorf("Failed to update usage.")
 
@@ -246,19 +254,19 @@ func (c *Client) GetInvoice(ctx context.Context, invoiceID string) (*stripe.Invo
 	return invoice, nil
 }
 
-// queriesForCustomersWithTeamIds constructs Stripe query strings to find the Stripe Customer for each teamId
+// queriesForCustomersWithAttributionIDs constructs Stripe query strings to find the Stripe Customer for each teamId
 // It returns multiple queries, each being a big disjunction of subclauses so that we can process multiple teamIds in one query.
 // `clausesPerQuery` is a limit enforced by the Stripe API.
-func queriesForCustomersWithTeamIds(teamIds []string) []string {
+func queriesForCustomersWithAttributionIDs(attributionIDs []db.AttributionID) []string {
 	const clausesPerQuery = 10
 	var queries []string
 	sb := strings.Builder{}
 
-	for i := 0; i < len(teamIds); i += clausesPerQuery {
+	for i := 0; i < len(attributionIDs); i += clausesPerQuery {
 		sb.Reset()
-		for j := 0; j < clausesPerQuery && i+j < len(teamIds); j++ {
-			sb.WriteString(fmt.Sprintf("metadata['%s']:'%s'", TeamIDMetadataKey, teamIds[i+j]))
-			if j < clausesPerQuery-1 && i+j < len(teamIds)-1 {
+		for j := 0; j < clausesPerQuery && i+j < len(attributionIDs); j++ {
+			sb.WriteString(fmt.Sprintf("metadata['%s']:'%s'", AttributionIDMetadataKey, attributionIDs[i+j]))
+			if j < clausesPerQuery-1 && i+j < len(attributionIDs)-1 {
 				sb.WriteString(" OR ")
 			}
 		}

components/usage/pkg/stripe/stripe_test.go

@@ -6,32 +6,33 @@ package stripe
 
 import (
 	"fmt"
+	"github.com/gitpod-io/gitpod/usage/pkg/db"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
-func TestCustomerQueriesForTeamIds_SingleQuery(t *testing.T) {
+func TestQueriesForCustomersWithAttributionID_Single(t *testing.T) {
 	testCases := []struct {
 		Name            string
-		TeamIds         []string
+		AttributionIDs  []db.AttributionID
 		ExpectedQueries []string
 	}{
 		{
 			Name:            "1 team id",
-			TeamIds:         []string{"abcd-123"},
-			ExpectedQueries: []string{"metadata['teamId']:'abcd-123'"},
+			AttributionIDs:  []db.AttributionID{db.NewTeamAttributionID("abcd-123")},
+			ExpectedQueries: []string{"metadata['attributionId']:'team:abcd-123'"},
 		},
 		{
-			Name:            "2 team ids",
-			TeamIds:         []string{"abcd-123", "abcd-456"},
-			ExpectedQueries: []string{"metadata['teamId']:'abcd-123' OR metadata['teamId']:'abcd-456'"},
+			Name:            "1 team id, 1 user id",
+			AttributionIDs:  []db.AttributionID{db.NewTeamAttributionID("abcd-123"), db.NewUserAttributionID("abcd-456")},
+			ExpectedQueries: []string{"metadata['attributionId']:'team:abcd-123' OR metadata['attributionId']:'user:abcd-456'"},
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.Name, func(t *testing.T) {
-			actualQueries := queriesForCustomersWithTeamIds(tc.TeamIds)
+			actualQueries := queriesForCustomersWithAttributionIDs(tc.AttributionIDs)
 
 			require.Equal(t, tc.ExpectedQueries, actualQueries)
 		})
@@ -66,18 +67,18 @@ func TestCustomerQueriesForTeamIds_MultipleQueries(t *testing.T) {
 		},
 	}
 
-	buildTeamIds := func(numberOfTeamIds int) []string {
-		var teamIds []string
+	buildTeamIds := func(numberOfTeamIds int) []db.AttributionID {
+		var attributionIDs []db.AttributionID
 		for i := 0; i < numberOfTeamIds; i++ {
-			teamIds = append(teamIds, fmt.Sprintf("abcd-%d", i))
+			attributionIDs = append(attributionIDs, db.NewTeamAttributionID(fmt.Sprintf("abcd-%d", i)))
 		}
-		return teamIds
+		return attributionIDs
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.Name, func(t *testing.T) {
 			teamIds := buildTeamIds(tc.NumberOfTeamIds)
-			actualQueries := queriesForCustomersWithTeamIds(teamIds)
+			actualQueries := queriesForCustomersWithAttributionIDs(teamIds)
 
 			require.Equal(t, tc.ExpectedNumberOfQueries, len(actualQueries))
 		})