[usage] Serialize metadata into usage records

Author: easyCZ

components/usage/pkg/apiv1/usage.go

@@ -341,7 +341,11 @@ func (s *UsageService) ReconcileUsageWithLedger(ctx context.Context, req *v1.Rec
 logger.Infof("Found %d workspaces instances for usage records in draft.", len(instancesWithUsageInDraft))
 instances = append(instances, instancesWithUsageInDraft...)
 
-inserts, updates := reconcileUsageWithLedger(instances, usageDrafts, s.pricer, now)
+inserts, updates, err := reconcileUsageWithLedger(instances, usageDrafts, s.pricer, now)
+if err != nil {
+logger.WithError(err).Errorf("Failed to reconcile usage with ledger.")
+return nil, status.Errorf(codes.Internal, "Failed to reconcile usage with ledger.")
+}
 logger.Infof("Identified %d inserts and %d updates against usage records.", len(inserts), len(updates))
 
 if len(inserts) > 0 {
@@ -365,7 +369,7 @@ func (s *UsageService) ReconcileUsageWithLedger(ctx context.Context, req *v1.Rec
 return &v1.ReconcileUsageWithLedgerResponse{}, nil
 }
 
-func reconcileUsageWithLedger(instances []db.WorkspaceInstanceForUsage, drafts []db.Usage, pricer *WorkspacePricer, now time.Time) (inserts []db.Usage, updates []db.Usage) {
+func reconcileUsageWithLedger(instances []db.WorkspaceInstanceForUsage, drafts []db.Usage, pricer *WorkspacePricer, now time.Time) (inserts []db.Usage, updates []db.Usage, err error) {
 
 instancesByID := dedupeWorkspaceInstancesForUsage(instances)
 
@@ -376,19 +380,27 @@ func reconcileUsageWithLedger(instances []db.WorkspaceInstanceForUsage, drafts [
 
 for instanceID, instance := range instancesByID {
 if usage, exists := draftsByWorkspaceID[instanceID]; exists {
-updates = append(updates, updateUsageFromInstance(instance, usage, pricer, now))
+updatedUsage, err := updateUsageFromInstance(instance, usage, pricer, now)
+if err != nil {
+return nil, nil, fmt.Errorf("failed to construct updated usage record: %w", err)
+}
+updates = append(updates, updatedUsage)
 continue
 }
 
-inserts = append(inserts, newUsageFromInstance(instance, pricer, now))
+usage, err := newUsageFromInstance(instance, pricer, now)
+if err != nil {
+return nil, nil, fmt.Errorf("failed to construct usage record: %w", err)
+}
+inserts = append(inserts, usage)
 }
 
-return inserts, updates
+return inserts, updates, nil
 }
 
 const usageDescriptionFromController = "Usage collected by automated system."
 
-func newUsageFromInstance(instance db.WorkspaceInstanceForUsage, pricer *WorkspacePricer, now time.Time) db.Usage {
+func newUsageFromInstance(instance db.WorkspaceInstanceForUsage, pricer *WorkspacePricer, now time.Time) (db.Usage, error) {
 draft := true
 if instance.StoppingTime.IsSet() {
 draft = false
@@ -399,7 +411,7 @@ func newUsageFromInstance(instance db.WorkspaceInstanceForUsage, pricer *Workspa
 effectiveTime = instance.StoppingTime.Time()
 }
 
-return db.Usage{
+usage := db.Usage{
 ID: uuid.New(),
 AttributionID: instance.UsageAttributionID,
 Description: usageDescriptionFromController,
@@ -408,17 +420,43 @@ func newUsageFromInstance(instance db.WorkspaceInstanceForUsage, pricer *Workspa
 Kind: db.WorkspaceInstanceUsageKind,
 WorkspaceInstanceID: instance.ID,
 Draft: draft,
-Metadata: nil,
 }
+
+startedTime := ""
+if instance.StartedTime.IsSet() {
+startedTime = db.TimeToISO8601(instance.StartedTime.Time())
+}
+endTime := ""
+if instance.StoppingTime.IsSet() {
+endTime = db.TimeToISO8601(instance.StoppingTime.Time())
+}
+err := usage.SetMetadataWithWorkspaceInstance(db.WorkspaceInstanceUsageData{
+WorkspaceId: instance.WorkspaceID,
+WorkspaceType: instance.Type,
+WorkspaceClass: instance.WorkspaceClass,
+ContextURL: "",
+StartTime: startedTime,
+EndTime: endTime,
+UserName: "",
+UserAvatarURL: "",
+})
+if err != nil {
+return db.Usage{}, fmt.Errorf("failed to serialize workspace instance metadata: %w", err)
+}
+
+return usage, nil
 }
 
-func updateUsageFromInstance(instance db.WorkspaceInstanceForUsage, usage db.Usage, pricer *WorkspacePricer, now time.Time) db.Usage {
+func updateUsageFromInstance(instance db.WorkspaceInstanceForUsage, usage db.Usage, pricer *WorkspacePricer, now time.Time) (db.Usage, error) {
 // We construct a new record to ensure we always take the data from the source of truth - the workspace instance
-updated := newUsageFromInstance(instance, pricer, now)
+updated, err := newUsageFromInstance(instance, pricer, now)
+if err != nil {
+return db.Usage{}, fmt.Errorf("failed to construct updated usage record: %w", err)
+}
 // but we override the ID to the one we already have
 updated.ID = usage.ID
 
-return updated
+return updated, nil
 }
 
 func collectWorkspaceInstanceIDs(usage []db.Usage) []uuid.UUID {

components/usage/pkg/apiv1/usage_test.go

@@ -624,14 +624,16 @@ func TestReconcileWithLedger(t *testing.T) {
 require.NoError(t, err)
 
 t.Run("no action with no instances and no drafts", func(t *testing.T) {
-inserts, updates := reconcileUsageWithLedger(nil, nil, pricer, now)
+inserts, updates, err := reconcileUsageWithLedger(nil, nil, pricer, now)
+require.NoError(t, err)
 require.Len(t, inserts, 0)
 require.Len(t, updates, 0)
 })
 
 t.Run("no action with no instances but existing drafts", func(t *testing.T) {
 drafts := []db.Usage{dbtest.NewUsage(t, db.Usage{})}
-inserts, updates := reconcileUsageWithLedger(nil, drafts, pricer, now)
+inserts, updates, err := reconcileUsageWithLedger(nil, drafts, pricer, now)
+require.NoError(t, err)
 require.Len(t, inserts, 0)
 require.Len(t, updates, 0)
 })
@@ -648,12 +650,14 @@ func TestReconcileWithLedger(t *testing.T) {
 WorkspaceClass: db.WorkspaceClass_Default,
 Type: db.WorkspaceType_Regular,
 UsageAttributionID: db.NewTeamAttributionID(uuid.New().String()),
+StartedTime: db.NewVarcharTime(now.Add(1 * time.Minute)),
 }
 
-inserts, updates := reconcileUsageWithLedger([]db.WorkspaceInstanceForUsage{instance, instance}, nil, pricer, now)
+inserts, updates, err := reconcileUsageWithLedger([]db.WorkspaceInstanceForUsage{instance, instance}, nil, pricer, now)
+require.NoError(t, err)
 require.Len(t, inserts, 1)
 require.Len(t, updates, 0)
-require.Equal(t, db.Usage{
+expectedUsage := db.Usage{
 ID: inserts[0].ID,
 AttributionID: instance.UsageAttributionID,
 Description: usageDescriptionFromController,
@@ -663,7 +667,18 @@ func TestReconcileWithLedger(t *testing.T) {
 WorkspaceInstanceID: instance.ID,
 Draft: true,
 Metadata: nil,
-}, inserts[0])
+}
+require.NoError(t, expectedUsage.SetMetadataWithWorkspaceInstance(db.WorkspaceInstanceUsageData{
+WorkspaceId: instance.WorkspaceID,
+WorkspaceType: instance.Type,
+WorkspaceClass: instance.WorkspaceClass,
+ContextURL: "",
+StartTime: db.TimeToISO8601(instance.StartedTime.Time()),
+EndTime: "",
+UserName: "",
+UserAvatarURL: "",
+}))
+require.EqualValues(t, expectedUsage, inserts[0])
 })
 
 t.Run("updates a usage record when a draft exists", func(t *testing.T) {
@@ -678,6 +693,7 @@ func TestReconcileWithLedger(t *testing.T) {
 WorkspaceClass: db.WorkspaceClass_Default,
 Type: db.WorkspaceType_Regular,
 UsageAttributionID: db.NewTeamAttributionID(uuid.New().String()),
+StartedTime: db.NewVarcharTime(now.Add(1 * time.Minute)),
 }
 
 // the fields in the usage record deliberately do not match the instance, except for the Instance ID.
@@ -694,10 +710,12 @@ func TestReconcileWithLedger(t *testing.T) {
 Metadata: nil,
 })
 
-inserts, updates := reconcileUsageWithLedger([]db.WorkspaceInstanceForUsage{instance}, []db.Usage{draft}, pricer, now)
+inserts, updates, err := reconcileUsageWithLedger([]db.WorkspaceInstanceForUsage{instance}, []db.Usage{draft}, pricer, now)
+require.NoError(t, err)
 require.Len(t, inserts, 0)
 require.Len(t, updates, 1)
-require.Equal(t, db.Usage{
+
+expectedUsage := db.Usage{
 ID: draft.ID,
 AttributionID: instance.UsageAttributionID,
 Description: usageDescriptionFromController,
@@ -707,6 +725,17 @@ func TestReconcileWithLedger(t *testing.T) {
 WorkspaceInstanceID: instance.ID,
 Draft: true,
 Metadata: nil,
-}, updates[0])
+}
+require.NoError(t, expectedUsage.SetMetadataWithWorkspaceInstance(db.WorkspaceInstanceUsageData{
+WorkspaceId: instance.WorkspaceID,
+WorkspaceType: instance.Type,
+WorkspaceClass: instance.WorkspaceClass,
+ContextURL: "",
+StartTime: db.TimeToISO8601(instance.StartedTime.Time()),
+EndTime: "",
+UserName: "",
+UserAvatarURL: "",
+}))
+require.EqualValues(t, expectedUsage, updates[0])
 })
 }

components/usage/pkg/db/usage.go

@@ -7,6 +7,7 @@ package db
 import (
 "context"
 "database/sql"
+"encoding/json"
 "fmt"
 "math"
 "time"
@@ -47,6 +48,26 @@ type Usage struct {
 Metadata datatypes.JSON `gorm:"column:metadata;type:text;size:65535" json:"metadata"`
 }
 
+func (u *Usage) SetMetadataWithWorkspaceInstance(data WorkspaceInstanceUsageData) error {
+b, err := json.Marshal(data)
+if err != nil {
+return fmt.Errorf("failed to serialize workspace instance usage data into json: %w", err)
+}
+
+u.Metadata = b
+return nil
+}
+
+func (u *Usage) GetMetadataAsWorkspaceInstanceData() (WorkspaceInstanceUsageData, error) {
+var data WorkspaceInstanceUsageData
+err := json.Unmarshal(u.Metadata, &data)
+if err != nil {
+return WorkspaceInstanceUsageData{}, fmt.Errorf("failed unmarshal metadata into wokrspace instance data: %w", err)
+}
+
+return data, nil
+}
+
 // WorkspaceInstanceUsageData represents the shape of metadata for usage entries of kind "workspaceinstance"
 // the equivalent TypeScript definition is maintained in `components/gitpod-protocol/src/usage.ts“
 type WorkspaceInstanceUsageData struct {