[stripe] Handle invoice finalization

Author: easyCZ

components/usage/pkg/apiv1/billing.go

@@ -6,18 +6,20 @@ package apiv1
 
 import (
 	"context"
-	"github.com/gitpod-io/gitpod/usage/pkg/contentservice"
+
 	"math"
 	"time"
 
 	"github.com/gitpod-io/gitpod/common-go/log"
 	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
+	"github.com/gitpod-io/gitpod/usage/pkg/contentservice"
 	"github.com/gitpod-io/gitpod/usage/pkg/db"
 	"github.com/gitpod-io/gitpod/usage/pkg/stripe"
 	"github.com/google/uuid"
 	stripesdk "github.com/stripe/stripe-go/v72"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"gorm.io/gorm"
 )
 
@@ -31,12 +33,12 @@ func NewBillingService(stripeClient *stripe.Client, billInstancesAfter time.Time
 }
 
 type BillingService struct {
-	conn               *gorm.DB
-	stripeClient       *stripe.Client
-	billInstancesAfter time.Time
-
+	conn           *gorm.DB
+	stripeClient   *stripe.Client
 	contentService contentservice.Interface
 
+	billInstancesAfter time.Time
+
 	v1.UnimplementedBillingServiceServer
 }
 
@@ -66,7 +68,66 @@ func (s *BillingService) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoic
 }
 
 func (s *BillingService) FinalizeInvoice(ctx context.Context, in *v1.FinalizeInvoiceRequest) (*v1.FinalizeInvoiceResponse, error) {
-	log.Infof("Finalizing invoice for invoice %q", in.GetInvoiceId())
+	logger := log.WithField("invoice_id", in.GetInvoiceId())
+
+	if in.GetInvoiceId() == "" {
+		return nil, status.Errorf(codes.InvalidArgument, "Missing InvoiceID")
+	}
+
+	invoice, err := s.stripeClient.GetInvoice(ctx, in.GetInvoiceId())
+	if err != nil {
+		logger.WithError(err).Error("Failed to retrieve invoice from Stripe.")
+		return nil, status.Errorf(codes.NotFound, "Failed to get invoice with ID %s: %s", in.GetInvoiceId(), err.Error())
+	}
+
+	reportID, found := invoice.Metadata[stripe.ReportIDMetadataKey]
+	if !found {
+		logger.Error("Failed to find report ID metadata on invoice from Stripe.")
+		return nil, status.Errorf(codes.NotFound, "Invoice %s does not contain reportID", in.GetInvoiceId())
+	}
+	logger = logger.WithField("report_id", reportID)
+
+	subscription := invoice.Subscription
+	if subscription == nil {
+		logger.Error("No subscription information available for invoice.")
+		return nil, status.Errorf(codes.Internal, "Failed to retrieve subscription details from invoice.")
+	}
+
+	teamID, found := subscription.Metadata[stripe.TeamIDMetadataKey]
+	if !found {
+		logger.Error("Failed to find teamID from subscription metadata.")
+		return nil, status.Errorf(codes.Internal, "Failed to extra teamID from Stripe subscription.")
+	}
+	logger = logger.WithField("team_id", teamID)
+
+	attributionID := db.NewTeamAttributionID(teamID)
+
+	// To support individual `user`s, we'll need to also extract the `userId` from metadata here and handle separately.
+
+	report, err := s.contentService.DownloadUsageReport(ctx, reportID)
+	if err != nil {
+		logger.WithError(err).Error("Failed to retrieve usage report from content service.")
+		return nil, status.Errorf(codes.Internal, "Failed to download usage report.")
+	}
+
+	invoicedSessions := report.GetUsageRecordsForAttributionID(attributionID)
+	var errors []error
+	for _, session := range invoicedSessions {
+		_, err := s.SetBilledSession(ctx, &v1.SetBilledSessionRequest{
+			InstanceId: session.InstanceID.String(),
+			From:       timestamppb.New(session.StartedAt),
+			System:     v1.System_SYSTEM_STRIPE,
+		})
+		if err != nil {
+			logger.WithField("workspace_ignstance_id", session.InstanceID).WithError(err).Error("Failed to mark session as billed by Stripe.")
+			errors = append(errors, err)
+		}
+	}
+
+	if len(errors) != 0 {
+		logger.Errorf("Failed to mark %d sessions as billed. You have to update them manually!", len(errors))
+		return nil, status.Errorf(codes.Internal, "Failed to mark %d sessions as billed by stripe.", len(errors))
+	}
 
 	return &v1.FinalizeInvoiceResponse{}, nil
 }

components/usage/pkg/contentservice/usage_report.go

@@ -25,3 +25,14 @@ type UsageReport struct {
 
 	UsageRecords []db.WorkspaceInstanceUsage `json:"usageRecords"`
 }
+
+func (r *UsageReport) GetUsageRecordsForAttributionID(attributionID db.AttributionID) []db.WorkspaceInstanceUsage {
+	var sessions []db.WorkspaceInstanceUsage
+	for _, sess := range r.UsageRecords {
+		if sess.AttributionID == attributionID {
+			sessions = append(sessions, sess)
+		}
+	}
+
+	return sessions
+}

components/usage/pkg/contentservice/usage_report_test.go

@@ -0,0 +1,44 @@
+// Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+// Licensed under the GNU Affero General Public License (AGPL).
+// See License-AGPL.txt in the project root for license information.
+
+package contentservice
+
+import (
+	"github.com/gitpod-io/gitpod/usage/pkg/db"
+	"github.com/gitpod-io/gitpod/usage/pkg/db/dbtest"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"testing"
+	"time"
+)
+
+func TestUsageReport_GetUsageRecordsForAttributionID(t *testing.T) {
+	attributionID_A, attributionID_B := db.NewTeamAttributionID(uuid.New().String()), db.NewTeamAttributionID(uuid.New().String())
+
+	report := UsageReport{
+		GenerationTime: time.Now(),
+		From:           time.Now(),
+		To:             time.Now(),
+		UsageRecords: []db.WorkspaceInstanceUsage{
+			dbtest.NewWorkspaceInstanceUsage(t, db.WorkspaceInstanceUsage{
+				AttributionID: attributionID_A,
+			}),
+			dbtest.NewWorkspaceInstanceUsage(t, db.WorkspaceInstanceUsage{
+				AttributionID: attributionID_A,
+			}),
+			dbtest.NewWorkspaceInstanceUsage(t, db.WorkspaceInstanceUsage{
+				AttributionID: attributionID_B,
+			}),
+		},
+	}
+
+	filteredToAttributionA := report.GetUsageRecordsForAttributionID(attributionID_A)
+	require.Equal(t, []db.WorkspaceInstanceUsage{report.UsageRecords[0], report.UsageRecords[1]}, filteredToAttributionA)
+
+	filteredToAttributionB := report.GetUsageRecordsForAttributionID(attributionID_B)
+	require.Equal(t, []db.WorkspaceInstanceUsage{report.UsageRecords[2]}, filteredToAttributionB)
+
+	filteredToAbsentAttribution := report.GetUsageRecordsForAttributionID(db.NewTeamAttributionID(uuid.New().String()))
+	require.Len(t, filteredToAbsentAttribution, 0)
+}

components/usage/pkg/db/workspace_instance_usage.go

@@ -89,5 +89,3 @@ func ListUsage(ctx context.Context, conn *gorm.DB, attributionId AttributionID,
 	}
 	return usageRecords, nil
 }
-
-type UsageReport []WorkspaceInstanceUsage

components/usage/pkg/stripe/stripe.go

@@ -17,7 +17,8 @@ import (
 )
 
 const (
-	reportIDMetadataKey = "reportId"
+	ReportIDMetadataKey = "reportId"
+	TeamIDMetadataKey   = "teamId"
 )
 
 type Client struct {
@@ -157,7 +158,7 @@ func (c *Client) updateUsageForCustomer(ctx context.Context, customer *stripe.Cu
 	}
 
 	_, err = c.UpdateInvoiceMetadata(ctx, invoice.ID, map[string]string{
-		reportIDMetadataKey: summary.ReportID,
+		ReportIDMetadataKey: summary.ReportID,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to udpate invoice %s metadata with report ID: %w", invoice.ID, err)
@@ -240,6 +241,23 @@ func (c *Client) UpdateInvoiceMetadata(ctx context.Context, invoiceID string, me
 	return invoice, nil
 }
 
+func (c *Client) GetInvoice(ctx context.Context, invoiceID string) (*stripe.Invoice, error) {
+	if invoiceID == "" {
+		return nil, fmt.Errorf("no invoice ID specified")
+	}
+
+	invoice, err := c.sc.Invoices.Get(invoiceID, &stripe.InvoiceParams{
+		Params: stripe.Params{
+			Context: ctx,
+			Expand:  []*string{stripe.String("data.subscriptions")},
+		},
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get invoice %s: %w", invoiceID, err)
+	}
+	return invoice, nil
+}
+
 // queriesForCustomersWithTeamIds 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.
@@ -251,7 +269,7 @@ func queriesForCustomersWithTeamIds(teamIds []string) []string {
 	for i := 0; i < len(teamIds); i += clausesPerQuery {
 		sb.Reset()
 		for j := 0; j < clausesPerQuery && i+j < len(teamIds); j++ {
-			sb.WriteString(fmt.Sprintf("metadata['teamId']:'%s'", teamIds[i+j]))
+			sb.WriteString(fmt.Sprintf("metadata['%s']:'%s'", TeamIDMetadataKey, teamIds[i+j]))
 			if j < clausesPerQuery-1 && i+j < len(teamIds)-1 {
 				sb.WriteString(" OR ")
 			}