[usage] Add usageAttributionID to WorksapceInstance model (in go)

Author: easyCZ

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

@@ -59,9 +59,15 @@ func NewWorkspaceInstance(t *testing.T, instance db.WorkspaceInstance) db.Worksp
 		status = instance.Status
 	}
 
+	attributionID := db.NewUserAttributionID(uuid.New().String())
+	if instance.UsageAttributionID != "" {
+		attributionID = instance.UsageAttributionID
+	}
+
 	return db.WorkspaceInstance{
 		ID:                 id,
 		WorkspaceID:        workspaceID,
+		UsageAttributionID: attributionID,
 		Configuration:      nil,
 		Region:             "",
 		ImageBuildInfo:     sql.NullString{},

components/usage/pkg/db/workspace_instance.go

@@ -11,6 +11,7 @@ import (
 	"github.com/google/uuid"
 	"gorm.io/datatypes"
 	"gorm.io/gorm"
+	"strings"
 	"time"
 )
 
@@ -23,6 +24,7 @@ type WorkspaceInstance struct {
 	IdeURL             string         `gorm:"column:ideUrl;type:varchar;size:255;" json:"ideUrl"`
 	WorkspaceBaseImage string         `gorm:"column:workspaceBaseImage;type:varchar;size:255;" json:"workspaceBaseImage"`
 	WorkspaceImage     string         `gorm:"column:workspaceImage;type:varchar;size:255;" json:"workspaceImage"`
+	UsageAttributionID AttributionID  `gorm:"column:usageAttributionId;type:varchar;size:60;" json:"usageAttributionId"`
 
 	CreationTime VarcharTime `gorm:"column:creationTime;type:varchar;size:255;" json:"creationTime"`
 	StartedTime  VarcharTime `gorm:"column:startedTime;type:varchar;size:255;" json:"startedTime"`
@@ -87,3 +89,33 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to
 
 	return instances, nil
 }
+
+const (
+	AttributionEntity_User = "user"
+	AttributionEntity_Team = "team"
+)
+
+func newAttributionID(entity, identifier string) AttributionID {
+	return AttributionID(fmt.Sprintf("%s:%s", entity, identifier))
+}
+
+func NewUserAttributionID(userID string) AttributionID {
+	return newAttributionID(AttributionEntity_User, userID)
+}
+
+func NewTeamAttributionID(teamID string) AttributionID {
+	return newAttributionID(AttributionEntity_Team, teamID)
+}
+
+// AttributionID consists of an entity, and an identifier in the form:
+// <entity>:<identifier>, e.g. team:a7dcf253-f05e-4dcf-9a47-cf8fccc74717
+type AttributionID string
+
+func (a AttributionID) Values() (entity string, identifier string) {
+	tokens := strings.Split(string(a), ":")
+	if len(tokens) != 2 {
+		return "", ""
+	}
+
+	return tokens[0], tokens[1]
+}

components/usage/pkg/db/workspace_instance_test.go

@@ -217,3 +217,25 @@ func TestListWorkspaceInstancesInRange_InBatches(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, results, len(instances))
 }
+
+func TestAttributionID_Values(t *testing.T) {
+	scenarios := []struct {
+		Input          string
+		ExpectedEntity string
+		ExpectedID     string
+	}{
+		{Input: "team:123", ExpectedEntity: db.AttributionEntity_Team, ExpectedID: "123"},
+		{Input: "user:123", ExpectedEntity: db.AttributionEntity_User, ExpectedID: "123"},
+		{Input: "user:123:invalid", ExpectedEntity: "", ExpectedID: ""},
+		{Input: "invalid:123:", ExpectedEntity: "", ExpectedID: ""},
+		{Input: "", ExpectedEntity: "", ExpectedID: ""},
+	}
+
+	for _, s := range scenarios {
+		t.Run(fmt.Sprintf("attribution in: %s, expecting: %s %s", s.Input, s.ExpectedEntity, s.ExpectedID), func(t *testing.T) {
+			entity, id := db.AttributionID(s.Input).Values()
+			require.Equal(t, s.ExpectedEntity, entity)
+			require.Equal(t, s.ExpectedID, id)
+		})
+	}
+}