fix up

Author: easyCZ

components/usage/pkg/apiv1/billing.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"github.com/gitpod-io/gitpod/usage/pkg/contentservice"
 	"fmt"
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"math"
 	"time"
 
@@ -67,31 +68,74 @@ 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")
 	}
 
-	invoiceID := in.GetInvoiceId()
-
-	invoice, err := s.stripeClient.GetInvoice(ctx, invoiceID)
+	invoice, err := s.stripeClient.GetInvoice(ctx, in.GetInvoiceId())
 	if err != nil {
-		return nil, status.Errorf(codes.NotFound, "Failed to get invoice with ID %s: %s", invoiceID, err.Error())
+		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 {
-		return nil, status.Errorf(codes.NotFound, "Invoice %s does not contain reportID", invoiceID)
+		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.")
 	}
 
-	// need to identify which AttributionID this finalize invoice event is for, how do we do that?
-	fmt.Println(report)
+	for _, session := range report {
+		if session.AttributionID == attributionID {
+
+		}
+	}
+
+	invoicedSessions := report.FilterToAttributionID(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_instance_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!")
+		// TODO(milan): Add metrics and alert on this!
+
+		return nil, status.Errorf(codes.Internal, "Failed to mark %d sessions as billed by stripe.")
+	}
 
 	return &v1.FinalizeInvoiceResponse{}, nil
 }

components/usage/pkg/db/workspace_instance_usage.go

@@ -91,3 +91,14 @@ func ListUsage(ctx context.Context, conn *gorm.DB, attributionId AttributionID,
 }
 
 type UsageReport []WorkspaceInstanceUsage
+
+func (r UsageReport) FilterToAttributionID(attributionID AttributionID) UsageReport {
+	var sessions []WorkspaceInstanceUsage
+	for _, sess := range r {
+		if sess.AttributionID == attributionID {
+			sessions = append(sessions, sess)
+		}
+	}
+
+	return sessions
+}

components/usage/pkg/stripe/stripe.go

@@ -18,6 +18,7 @@ import (
 
 const (
 	ReportIDMetadataKey = "reportId"
+	TeamIDMetadataKey   = "teamId"
 )
 
 type Client struct {
@@ -248,6 +249,7 @@ func (c *Client) GetInvoice(ctx context.Context, invoiceID string) (*stripe.Invo
 	invoice, err := c.sc.Invoices.Get(invoiceID, &stripe.InvoiceParams{
 		Params: stripe.Params{
 			Context: ctx,
+			Expand:  []*string{stripe.String("data.subscriptions")},
 		},
 	})
 	if err != nil {
@@ -267,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 ")
 			}