[usage] Find running and stopped instances in ledger reconciler

Author: easyCZ

components/usage/pkg/apiv1/usage.go

@@ -181,6 +181,42 @@ func (s *UsageService) GetCostCenter(ctx context.Context, in *v1.GetCostCenterRe
 	}, nil
 }
 
+func (s *UsageService) ReconcileUsageWithLedger(ctx context.Context, req *v1.ReconcileUsageWithLedgerRequest) (*v1.ReconcileUsageWithLedgerResponse, error) {
+	from := req.GetFrom().AsTime()
+	to := req.GetTo().AsTime()
+
+	logger := log.
+		WithField("from", from).
+		WithField("to", to)
+
+	if to.Before(from) {
+		return nil, status.Errorf(codes.InvalidArgument, "To must not be before From")
+	}
+
+	stopped, err := db.FindStoppedWorkspaceInstancesInRange(ctx, s.conn, from, to)
+	if err != nil {
+		logger.WithError(err).Errorf("Failed to find stopped workspace instances.")
+		return nil, status.Errorf(codes.Internal, "failed to query for stopped instances")
+	}
+	logger.Infof("Found %d stopped workspace instances in range.", len(stopped))
+
+	running, err := db.FindRunningWorkspaceInstances(ctx, s.conn)
+	if err != nil {
+		logger.WithError(err).Errorf("Failed to find running workspace instances.")
+		return nil, status.Errorf(codes.Internal, "failed to query for running instances")
+	}
+	logger.Infof("Found %d running workspaces since the beginning of time.", len(running))
+
+	usageDrafts, err := db.FindAllDraftUsage(ctx, s.conn)
+	if err != nil {
+		logger.WithError(err).Errorf("Failed to find all draft usage records.")
+		return nil, status.Errorf(codes.Internal, "failed to find all draft usage records")
+	}
+	logger.Infof("Found %d draft usage records.", len(usageDrafts))
+
+	return &v1.ReconcileUsageWithLedgerResponse{}, nil
+}
+
 func NewUsageService(conn *gorm.DB, reportGenerator *ReportGenerator, contentSvc contentservice.Interface) *UsageService {
 	return &UsageService{
 		conn:            conn,

components/usage/pkg/apiv1/usage_test.go

@@ -556,3 +556,41 @@ func TestReportGenerator_GenerateUsageReportTable(t *testing.T) {
 		})
 	}
 }
+
+func TestUsageService_ReconcileUsageWithLedger(t *testing.T) {
+	dbconn := dbtest.ConnectForTests(t)
+	from := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)
+	to := time.Date(2022, 05, 1, 1, 00, 00, 00, time.UTC)
+
+	// stopped instances
+	dbtest.CreateWorkspaceInstances(t, dbconn, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{
+		StoppingTime: db.NewVarcharTime(from.Add(1 * time.Minute)),
+	}))
+
+	// running instances
+	dbtest.CreateWorkspaceInstances(t, dbconn, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{}))
+
+	// usage drafts
+	dbtest.CreateUsageRecords(t, dbconn, dbtest.NewUsage(t, db.Usage{
+		Kind:  db.WorkspaceInstanceUsageKind,
+		Draft: true,
+	}))
+
+	srv := baseserver.NewForTests(t,
+		baseserver.WithGRPC(baseserver.MustUseRandomLocalAddress(t)),
+	)
+
+	v1.RegisterUsageServiceServer(srv.GRPC(), NewUsageService(dbconn, nil, nil))
+	baseserver.StartServerForTests(t, srv)
+
+	conn, err := grpc.Dial(srv.GRPCAddress(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+
+	client := v1.NewUsageServiceClient(conn)
+
+	_, err = client.ReconcileUsageWithLedger(context.Background(), &v1.ReconcileUsageWithLedgerRequest{
+		From: timestamppb.New(from),
+		To:   timestamppb.New(to),
+	})
+	require.NoError(t, err)
+}

components/usage/pkg/db/dbtest/usage.go

@@ -23,7 +23,7 @@ func NewUsage(t *testing.T, record db.Usage) db.Usage {
 		Description:         "some description",
 		CreditCents:         42,
 		EffectiveTime:       db.VarcharTime{},
-		Kind:                "workspaceinstance",
+		Kind:                db.WorkspaceInstanceUsageKind,
 		WorkspaceInstanceID: uuid.New(),
 	}
 

components/usage/pkg/db/usage.go

@@ -14,13 +14,20 @@ import (
 	"gorm.io/gorm/clause"
 )
 
+type UsageKind string
+
+const (
+	WorkspaceInstanceUsageKind UsageKind = "workspaceinstance"
+	InvoiceUsageKind                     = "invoice"
+)
+
 type Usage struct {
 	ID                  uuid.UUID      `gorm:"primary_key;column:id;type:char;size:36;" json:"id"`
 	AttributionID       AttributionID  `gorm:"column:attributionId;type:varchar;size:255;" json:"attributionId"`
 	Description         string         `gorm:"column:description;type:varchar;size:255;" json:"description"`
 	CreditCents         int64          `gorm:"column:creditCents;type:bigint;" json:"creditCents"`
 	EffectiveTime       VarcharTime    `gorm:"column:effectiveTime;type:varchar;size:255;" json:"effectiveTime"`
-	Kind                string         `gorm:"column:kind;type:char;size:10;" json:"kind"`
+	Kind                UsageKind      `gorm:"column:kind;type:char;size:10;" json:"kind"`
 	WorkspaceInstanceID uuid.UUID      `gorm:"column:workspaceInstanceId;type:char;size:36;" json:"workspaceInstanceId"`
 	Draft               bool           `gorm:"column:draft;type:boolean;" json:"draft"`
 	Metadata            datatypes.JSON `gorm:"column:metadata;type:text;size:65535" json:"metadata"`