[usage] Implement BillingService

Author: easyCZ

components/usage/pkg/apiv1/billing.go

@@ -0,0 +1,73 @@
+// 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 apiv1
+
+import (
+	"context"
+	"fmt"
+	"github.com/gitpod-io/gitpod/common-go/log"
+	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
+	"github.com/gitpod-io/gitpod/usage/pkg/db"
+	"github.com/gitpod-io/gitpod/usage/pkg/stripe"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"math"
+)
+
+func NewBillingService(stripeClient *stripe.Client) *BillingService {
+	return &BillingService{
+		stripeClient: stripeClient,
+	}
+}
+
+type BillingService struct {
+	stripeClient *stripe.Client
+
+	v1.UnimplementedBillingServiceServer
+}
+
+func (s *BillingService) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoicesRequest) (*v1.UpdateInvoicesResponse, error) {
+	credits, err := creditSummaryForTeams(in.GetSessions())
+	if err != nil {
+		log.Log.WithError(err).Errorf("Failed to compute credit summary.")
+		return nil, status.Errorf(codes.InvalidArgument, "failed to compute credit summary")
+	}
+
+	err = s.stripeClient.UpdateUsage(credits)
+	if err != nil {
+		log.Log.WithError(err).Errorf("Failed to update stripe invoices.")
+		return nil, status.Errorf(codes.Internal, "failed to update stripe invoices")
+	}
+
+	return &v1.UpdateInvoicesResponse{}, nil
+}
+
+func creditSummaryForTeams(sessions []*v1.BilledSession) (map[string]int64, error) {
+	creditsPerTeamID := map[string]int64{}
+
+	for _, session := range sessions {
+		attributionID, err := db.ParseAttributionID(session.AttributionId)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse attribution ID: %w", err)
+		}
+
+		entity, id := attributionID.Values()
+		if entity != db.AttributionEntity_Team {
+			continue
+		}
+
+		if _, ok := creditsPerTeamID[id]; !ok {
+			creditsPerTeamID[id] = 0
+		}
+
+		creditsPerTeamID[id] += session.GetCredits()
+	}
+
+	for teamID, credits := range creditsPerTeamID {
+		creditsPerTeamID[teamID] = int64(math.Ceil(float64(credits)))
+	}
+
+	return creditsPerTeamID, nil
+}

components/usage/pkg/apiv1/billing_noop.go

@@ -0,0 +1,21 @@
+// 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 apiv1
+
+import (
+	"context"
+	"github.com/gitpod-io/gitpod/common-go/log"
+	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
+)
+
+// BillingServiceNoop is used for Self-Hosted installations
+type BillingServiceNoop struct {
+	v1.UnimplementedBillingServiceServer
+}
+
+func (s *BillingServiceNoop) UpdateInvoices(_ context.Context, _ *v1.UpdateInvoicesRequest) (*v1.UpdateInvoicesResponse, error) {
+	log.Log.Infof("UpdateInvoices RPC invoked in no-op mode, no invoices will be updated.")
+	return &v1.UpdateInvoicesResponse{}, nil
+}

components/usage/pkg/db/workspace_instance.go

@@ -120,6 +120,22 @@ func (a AttributionID) Values() (entity string, identifier string) {
 	return tokens[0], tokens[1]
 }
 
+func ParseAttributionID(s string) (AttributionID, error) {
+	tokens := strings.Split(s, ":")
+	if len(tokens) != 2 {
+		return "", fmt.Errorf("attribution ID (%s) does not have two parts", s)
+	}
+
+	switch tokens[0] {
+	case AttributionEntity_Team:
+		return NewTeamAttributionID(tokens[1]), nil
+	case AttributionEntity_User:
+		return NewUserAttributionID(tokens[1]), nil
+	default:
+		return "", fmt.Errorf("unknown attribution ID type: %s", s)
+	}
+}
+
 const (
 	WorkspaceClass_Default = "default"
 )

components/usage/pkg/server/server.go

@@ -54,6 +54,7 @@ func Start(cfg Config) error {
 
 	var billingController controller.BillingController = &controller.NoOpBillingController{}
 
+	var stripeClient *stripe.Client
 	if cfg.StripeCredentialsFile != "" {
 		config, err := stripe.ReadConfigFromFile(cfg.StripeCredentialsFile)
 		if err != nil {
@@ -65,7 +66,8 @@ func Start(cfg Config) error {
 			return fmt.Errorf("failed to initialize stripe client: %w", err)
 		}
 
-		billingController = controller.NewStripeBillingController(c)
+		stripeClient = c
+		billingController = controller.NewStripeBillingController(c, pricer)
 	}
 
 	schedule, err := time.ParseDuration(cfg.ControllerSchedule)
@@ -97,7 +99,7 @@ func Start(cfg Config) error {
 	if err != nil {
 		return fmt.Errorf("failed to initialize usage server: %w", err)
 	}
-	err = registerGRPCServices(srv, conn)
+	err = registerGRPCServices(srv, conn, stripeClient)
 	if err != nil {
 		return fmt.Errorf("failed to register gRPC services: %w", err)
 	}
@@ -115,7 +117,12 @@ func Start(cfg Config) error {
 	return nil
 }
 
-func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB) error {
+func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *stripe.Client) error {
 	v1.RegisterUsageServiceServer(srv.GRPC(), apiv1.NewUsageService(conn))
+	if stripeClient == nil {
+		v1.RegisterBillingServiceServer(srv.GRPC(), &apiv1.BillingServiceNoop{})
+	} else {
+		v1.RegisterBillingServiceServer(srv.GRPC(), apiv1.NewBillingService(stripeClient))
+	}
 	return nil
 }