[usage] store and fetch usage entries

Author: svenefftinge

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

@@ -5,6 +5,7 @@
 package dbtest
 
 import (
+	"context"
 	"testing"
 
 	"github.com/gitpod-io/gitpod/usage/pkg/db"
@@ -67,8 +68,7 @@ func CreateUsageRecords(t *testing.T, conn *gorm.DB, entries ...db.Usage) []db.U
 		ids = append(ids, record.ID.String())
 	}
 
-	require.NoError(t, conn.CreateInBatches(&records, 1000).Error)
-
+	require.NoError(t, db.InsertUsage(context.Background(), conn, entries...))
 	t.Cleanup(func() {
 		require.NoError(t, conn.Where(ids).Delete(&db.Usage{}).Error)
 	})

components/usage/pkg/db/usage.go

@@ -11,6 +11,7 @@ import (
 	"github.com/google/uuid"
 	"gorm.io/datatypes"
 	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
 )
 
 type Usage struct {
@@ -34,12 +35,33 @@ func (u *Usage) TableName() string {
 	return "d_b_usage"
 }
 
+func InsertUsage(ctx context.Context, conn *gorm.DB, records ...Usage) error {
+	return conn.WithContext(ctx).Debug().
+		Clauses(clause.OnConflict{DoNothing: true}).
+		CreateInBatches(records, 1000).Error
+}
+
+func UpdateUsage(ctx context.Context, conn *gorm.DB, record Usage) error {
+	return conn.WithContext(ctx).Save(record).Error
+}
+
+func FindAllDraftUsage(ctx context.Context, conn *gorm.DB) ([]Usage, error) {
+	var usageRecords []Usage
+	result := conn.WithContext(ctx).
+		Where("draft = TRUE").
+		Order("effectiveTime DESC").
+		Find(&usageRecords)
+	if result.Error != nil {
+		return nil, fmt.Errorf("failed to get usage records: %s", result.Error)
+	}
+	return usageRecords, nil
+}
+
 func FindUsage(ctx context.Context, conn *gorm.DB, attributionId AttributionID, from, to VarcharTime, offset int64, limit int64) ([]Usage, error) {
 	db := conn.WithContext(ctx)
 
 	var usageRecords []Usage
 	result := db.
-		WithContext(ctx).
 		Where("attributionId = ?", attributionId).
 		Where("? <= effectiveTime AND effectiveTime < ?", from.String(), to.String()).
 		Order("effectiveTime DESC").

components/usage/pkg/db/usage_test.go

@@ -47,3 +47,93 @@ func TestFindUsageInRange(t *testing.T) {
 	require.Equal(t, 1, len(listResult))
 	require.Equal(t, []db.Usage{entryInside}, listResult)
 }
+
+func TestInsertUsageRecords(t *testing.T) {
+	conn := dbtest.ConnectForTests(t)
+
+	attributionID := db.NewTeamAttributionID(uuid.New().String())
+	start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC)
+
+	usage := dbtest.NewUsage(t, db.Usage{
+		AttributionID: attributionID,
+		EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)),
+		Draft:         true,
+	})
+
+	dbtest.CreateUsageRecords(t, conn, usage)
+	updatedDesc := "Updated Description"
+	usage.Description = updatedDesc
+
+	require.NoError(t, db.InsertUsage(context.Background(), conn, usage))
+
+	drafts, err := db.FindAllDraftUsage(context.Background(), conn)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(drafts))
+	require.NotEqual(t, updatedDesc, drafts[0].Description)
+}
+
+func TestUpdateUsageRecords(t *testing.T) {
+	conn := dbtest.ConnectForTests(t)
+
+	attributionID := db.NewTeamAttributionID(uuid.New().String())
+	start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC)
+
+	usage := dbtest.NewUsage(t, db.Usage{
+		AttributionID: attributionID,
+		EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)),
+		Draft:         true,
+	})
+
+	dbtest.CreateUsageRecords(t, conn, usage)
+	updatedDesc := "Updated Description"
+	usage.Description = updatedDesc
+
+	require.NoError(t, db.UpdateUsage(context.Background(), conn, usage))
+
+	drafts, err := db.FindAllDraftUsage(context.Background(), conn)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(drafts))
+	require.Equal(t, updatedDesc, drafts[0].Description)
+}
+
+func TestFindAllDraftUsage(t *testing.T) {
+	conn := dbtest.ConnectForTests(t)
+
+	attributionID := db.NewTeamAttributionID(uuid.New().String())
+	start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC)
+
+	usage1 := dbtest.NewUsage(t, db.Usage{
+		AttributionID: attributionID,
+		EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)),
+		Draft:         true,
+	})
+	usage2 := dbtest.NewUsage(t, db.Usage{
+		AttributionID: attributionID,
+		EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)),
+		Draft:         true,
+	})
+	usage3 := dbtest.NewUsage(t, db.Usage{
+		AttributionID: attributionID,
+		EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)),
+		Draft:         false,
+	})
+
+	dbtest.CreateUsageRecords(t, conn, usage1, usage2, usage3)
+	drafts, err := db.FindAllDraftUsage(context.Background(), conn)
+	require.NoError(t, err)
+	require.Equal(t, 2, len(drafts))
+	for _, usage := range drafts {
+		require.True(t, usage.Draft)
+	}
+
+	// let's finalize one record
+	usage2.Draft = false
+	require.NoError(t, db.UpdateUsage(context.Background(), conn, usage2))
+
+	drafts, err = db.FindAllDraftUsage(context.Background(), conn)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(drafts))
+	for _, usage := range drafts {
+		require.True(t, usage.Draft)
+	}
+}