diff --git a/components/usage/pkg/db/dbtest/workspace_instance.go b/components/usage/pkg/db/dbtest/workspace_instance.go index 72dc7604d63581..016f1c72f7de2b 100644 --- a/components/usage/pkg/db/dbtest/workspace_instance.go +++ b/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{}, diff --git a/components/usage/pkg/db/workspace_instance.go b/components/usage/pkg/db/workspace_instance.go index fdcf0fe7af49d5..396ac48a759218 100644 --- a/components/usage/pkg/db/workspace_instance.go +++ b/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: +// :, 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] +} diff --git a/components/usage/pkg/db/workspace_instance_test.go b/components/usage/pkg/db/workspace_instance_test.go index be9cf739880457..73760a55c4d159 100644 --- a/components/usage/pkg/db/workspace_instance_test.go +++ b/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) + }) + } +}