From c825b27ff84348f773be985d5d57ea424c6dc6f8 Mon Sep 17 00:00:00 2001 From: Milan Pavlik Date: Mon, 5 Sep 2022 10:18:36 +0000 Subject: [PATCH] [usage] Find running and stopped instances in ledger reconciler --- components/usage/pkg/apiv1/usage.go | 36 ++++++++++++++++++++++ components/usage/pkg/apiv1/usage_test.go | 38 ++++++++++++++++++++++++ components/usage/pkg/db/dbtest/usage.go | 2 +- components/usage/pkg/db/usage.go | 9 +++++- 4 files changed, 83 insertions(+), 2 deletions(-) diff --git a/components/usage/pkg/apiv1/usage.go b/components/usage/pkg/apiv1/usage.go index d704ddc810519a..0cdcb51a06313c 100644 --- a/components/usage/pkg/apiv1/usage.go +++ b/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, diff --git a/components/usage/pkg/apiv1/usage_test.go b/components/usage/pkg/apiv1/usage_test.go index 938f84d50c6f4f..c77b3eba67351a 100644 --- a/components/usage/pkg/apiv1/usage_test.go +++ b/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) +} diff --git a/components/usage/pkg/db/dbtest/usage.go b/components/usage/pkg/db/dbtest/usage.go index a34712b0098249..ea7c40dc6423c1 100644 --- a/components/usage/pkg/db/dbtest/usage.go +++ b/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(), } diff --git a/components/usage/pkg/db/usage.go b/components/usage/pkg/db/usage.go index ee470dd2a7e750..ba74359f356fbc 100644 --- a/components/usage/pkg/db/usage.go +++ b/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"`