diff --git a/components/usage/pkg/apiv1/usage.go b/components/usage/pkg/apiv1/usage.go index 3b32b1fd24d722..2db2f0de2d6a2a 100644 --- a/components/usage/pkg/apiv1/usage.go +++ b/components/usage/pkg/apiv1/usage.go @@ -344,6 +344,24 @@ func (s *UsageService) ReconcileUsageWithLedger(ctx context.Context, req *v1.Rec inserts, updates := reconcileUsageWithLedger(instances, usageDrafts, s.pricer, now) logger.WithField("inserts", inserts).WithField("updates", updates).Infof("Identified %d inserts and %d updates against usage records.", len(inserts), len(updates)) + if len(inserts) > 0 { + err = db.InsertUsage(ctx, s.conn, inserts...) + if err != nil { + logger.WithError(err).Errorf("Failed to insert %d usage records into the database.", len(inserts)) + return nil, status.Errorf(codes.Internal, "Failed to insert usage records into the database.") + } + logger.Infof("Inserted %d new Usage records into the database.", len(inserts)) + } + + if len(updates) > 0 { + err = db.UpdateUsage(ctx, s.conn, updates...) + if err != nil { + logger.WithError(err).Error("Failed to update usage records in the database.") + return nil, status.Errorf(codes.Internal, "Failed to update usage records in the database.") + } + logger.Infof("Updated %d Usage records in the database.", len(updates)) + } + return &v1.ReconcileUsageWithLedgerResponse{}, nil } diff --git a/components/usage/pkg/apiv1/usage_test.go b/components/usage/pkg/apiv1/usage_test.go index 8b59875dc635a0..7ba6842a930847 100644 --- a/components/usage/pkg/apiv1/usage_test.go +++ b/components/usage/pkg/apiv1/usage_test.go @@ -561,19 +561,26 @@ 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) + attributionID := db.NewTeamAttributionID(uuid.New().String()) // stopped instances - dbtest.CreateWorkspaceInstances(t, dbconn, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - StoppingTime: db.NewVarcharTime(from.Add(1 * time.Minute)), - })) + instance := dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + UsageAttributionID: attributionID, + CreationTime: db.NewVarcharTime(from), + StoppingTime: db.NewVarcharTime(to.Add(-1 * time.Minute)), + }) + dbtest.CreateWorkspaceInstances(t, dbconn, instance) // 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, + ID: uuid.New(), + AttributionID: attributionID, + WorkspaceInstanceID: instance.ID, + Kind: db.WorkspaceInstanceUsageKind, + Draft: true, })) srv := baseserver.NewForTests(t, @@ -593,6 +600,15 @@ func TestUsageService_ReconcileUsageWithLedger(t *testing.T) { To: timestamppb.New(to), }) require.NoError(t, err) + + usage, err := db.FindUsage(context.Background(), dbconn, &db.FindUsageParams{ + AttributionId: attributionID, + From: from, + To: to, + ExcludeDrafts: false, + }) + require.NoError(t, err) + require.Len(t, usage, 1) } func TestReconcileWithLedger(t *testing.T) { diff --git a/components/usage/pkg/db/dbtest/usage.go b/components/usage/pkg/db/dbtest/usage.go index ea7c40dc6423c1..d9a5b27ff0eb6d 100644 --- a/components/usage/pkg/db/dbtest/usage.go +++ b/components/usage/pkg/db/dbtest/usage.go @@ -73,7 +73,5 @@ func CreateUsageRecords(t *testing.T, conn *gorm.DB, entries ...db.Usage) []db.U require.NoError(t, conn.Where(ids).Delete(&db.Usage{}).Error) }) - t.Logf("stored %d", len(entries)) - return records } diff --git a/components/usage/pkg/db/usage.go b/components/usage/pkg/db/usage.go index dd3b95b268cbdc..963e5df0620ae9 100644 --- a/components/usage/pkg/db/usage.go +++ b/components/usage/pkg/db/usage.go @@ -62,8 +62,15 @@ func InsertUsage(ctx context.Context, conn *gorm.DB, records ...Usage) error { CreateInBatches(records, 1000).Error } -func UpdateUsage(ctx context.Context, conn *gorm.DB, record Usage) error { - return conn.WithContext(ctx).Save(record).Error +func UpdateUsage(ctx context.Context, conn *gorm.DB, records ...Usage) error { + for _, record := range records { + err := conn.WithContext(ctx).Save(record).Error + if err != nil { + return fmt.Errorf("failed to update usage record ID: %s: %w", record.ID, err) + } + } + + return nil } func FindAllDraftUsage(ctx context.Context, conn *gorm.DB) ([]Usage, error) { @@ -97,7 +104,7 @@ func FindUsage(ctx context.Context, conn *gorm.DB, params *FindUsageParams) ([]U db := conn.WithContext(ctx). Where("attributionId = ?", params.AttributionId). - Where("? <= effectiveTime AND effectiveTime < ?", params.From, params.To) + Where("effectiveTime >= ? AND effectiveTime < ?", TimeToISO8601(params.From), TimeToISO8601(params.To)) if params.ExcludeDrafts { db = db.Where("draft = ?", false) } @@ -130,7 +137,7 @@ func GetUsageSummary(ctx context.Context, conn *gorm.DB, attributionId Attributi query1 := db.Table((&Usage{}).TableName()). Select("sum(creditCents) as creditCentsBalanceAtStart"). Where("attributionId = ?", attributionId). - Where("effectiveTime < ?", from) + Where("effectiveTime < ?", TimeToISO8601(from)) if excludeDrafts { query1 = query1.Where("draft = ?", false) } @@ -143,7 +150,7 @@ func GetUsageSummary(ctx context.Context, conn *gorm.DB, attributionId Attributi query2 := db.Table((&Usage{}).TableName()). Select("sum(creditCents) as creditCentsBalanceInPeriod", "count(id) as numRecordsInRange"). Where("attributionId = ?", attributionId). - Where("? <= effectiveTime AND effectiveTime < ?", from, to) + Where("? <= effectiveTime AND effectiveTime < ?", TimeToISO8601(from), TimeToISO8601(to)) if excludeDrafts { query2 = query2.Where("draft = ?", false) }