[stripe] Use attributionId in stripe queries

Author: easyCZ

components/usage/pkg/apiv1/billing.go

@@ -43,19 +43,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]int64{}
 	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{
-			Credits:  int64(math.Ceil(balance.CreditCents.ToCredits())),
-			ReportID: "no-report",
-		}
+		creditSummaryForTeams[balance.AttributionID] = int64(math.Ceil(balance.CreditCents.ToCredits()))
 	}
 
 	err = s.stripeClient.UpdateUsage(ctx, creditSummaryForTeams)
@@ -86,7 +76,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.")

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 {
@@ -63,19 +64,14 @@ type Invoice struct {
 	Credits        int64
 }
 
-type CreditSummary struct {
-	Credits  int64
-	ReportID string
-}
-
 // 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]int64) 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 +82,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)
+
+			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, creditsPerTeam[teamID])
+			_, 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.")
 
@@ -127,7 +130,7 @@ func (c *Client) findCustomers(ctx context.Context, query string) ([]*stripe.Cus
 	return customers, nil
 }
 
-func (c *Client) updateUsageForCustomer(ctx context.Context, customer *stripe.Customer, summary CreditSummary) (*UsageRecord, error) {
+func (c *Client) updateUsageForCustomer(ctx context.Context, customer *stripe.Customer, credits int64) (*UsageRecord, error) {
 	subscriptions := customer.Subscriptions.Data
 	if len(subscriptions) != 1 {
 		return nil, fmt.Errorf("customer has an unexpected number of subscriptions %v (expected 1, got %d)", subscriptions, len(subscriptions))
@@ -146,15 +149,15 @@ func (c *Client) updateUsageForCustomer(ctx context.Context, customer *stripe.Cu
 			Context: ctx,
 		},
 		SubscriptionItem: stripe.String(subscriptionItemId),
-		Quantity:         stripe.Int64(summary.Credits),
+		Quantity:         stripe.Int64(credits),
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to register usage for customer %q on subscription item %s", customer.Name, subscriptionItemId)
 	}
 
 	return &UsageRecord{
 		SubscriptionItemID: subscriptionItemId,
-		Quantity:           summary.Credits,
+		Quantity:           credits,
 	}, nil
 }
 
@@ -246,19 +249,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))
 		})