[usage] Save inserted and updated records in the DB

Author: easyCZ

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
 }
 

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) {

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
 }

components/usage/pkg/db/usage.go

@@ -75,8 +75,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) {
@@ -110,7 +117,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)
 	}
@@ -143,7 +150,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)
 	}
@@ -156,7 +163,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)
 	}