[usage] Attribute workspaces to a team

components/usage/pkg/controller/reconciler.go

@@ -41,6 +41,8 @@ type UsageReconcileStatus struct {
 	InvalidWorkspaceInstances int
 
 	Workspaces int
+
+	Teams int
 }
 
 func (u *UsageReconciler) Reconcile() error {
@@ -82,9 +84,60 @@ func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time.
 	}
 	status.Workspaces = len(workspaces)
 
+	// match workspaces to teams
+	teams, err := u.loadTeamsForWorkspaces(ctx, workspaces)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load teams for workspaces: %w", err)
+	}
+	status.Teams = len(teams)
+
 	return status, nil
 }
 
+type teamWithWorkspaces struct {
+	TeamID     uuid.UUID
+	Workspaces []workspaceWithInstances
+}
+
+func (u *UsageReconciler) loadTeamsForWorkspaces(ctx context.Context, workspaces []workspaceWithInstances) ([]teamWithWorkspaces, error) {
+	// find owner IDs of these workspaces
+	var ownerIDs []uuid.UUID
+	for _, workspace := range workspaces {
+		ownerIDs = append(ownerIDs, workspace.Workspace.OwnerID)
+	}
+
+	// Retrieve memberships. This gives a link between an Owner and a Team they belong to.
+	memberships, err := db.ListTeamMembershipsForUserIDs(ctx, u.conn, ownerIDs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list team memberships: %w", err)
+	}
+
+	membershipsByUserID := map[uuid.UUID]db.TeamMembership{}
+	for _, membership := range memberships {
+		// User can belong to multiple teams. For now, we're choosing the membership at random.
+		membershipsByUserID[membership.UserID] = membership
+	}
+
+	// Convert workspaces into a lookup so that we can index into them by Owner ID, needed for joining Teams with Workspaces
+	workspacesByOwnerID := map[uuid.UUID][]workspaceWithInstances{}
+	for _, workspace := range workspaces {
+		workspacesByOwnerID[workspace.Workspace.OwnerID] = append(workspacesByOwnerID[workspace.Workspace.OwnerID], workspace)
+	}
+
+	// Finally, join the datasets
+	// Because we iterate over memberships, and not workspaces, we're in effect ignoring Workspaces which are not in a team.
+	// This is intended as we focus on Team usage for now.
+	var teamsWithWorkspaces []teamWithWorkspaces
+	for userID, membership := range membershipsByUserID {
+		teamsWithWorkspaces = append(teamsWithWorkspaces, teamWithWorkspaces{
+			TeamID:     membership.TeamID,
+			Workspaces: workspacesByOwnerID[userID],
+		})
+	}
+
+	return teamsWithWorkspaces, nil
+}
+
 type workspaceWithInstances struct {
 	Workspace db.Workspace
 	Instances []db.WorkspaceInstance

components/usage/pkg/controller/reconciler_test.go

@@ -19,7 +19,18 @@ func TestUsageReconciler_Reconcile(t *testing.T) {
 	instanceStatus := []byte(`{"phase": "stopped", "conditions": {"deployed": false, "pullingImages": false, "serviceExists": false}}`)
 	startOfMay := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)
 	startOfJune := time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC)
-	workspace := dbtest.NewWorkspace(t, db.Workspace{ID: "gitpodio-gitpod-gyjr82jkfnd"})
+	teamID := uuid.New()
+	userID := uuid.New()
+	membership := db.TeamMembership{
+		ID:     uuid.New(),
+		TeamID: teamID,
+		UserID: userID,
+		Role:   db.TeamMembershipRole_Member,
+	}
+	workspace := dbtest.NewWorkspace(t, db.Workspace{
+		ID:      "gitpodio-gitpod-gyjr82jkfnd",
+		OwnerID: userID,
+	})
 	instances := []db.WorkspaceInstance{
 		// Ran throughout the reconcile period
 		{
@@ -51,6 +62,9 @@ func TestUsageReconciler_Reconcile(t *testing.T) {
 	tx = conn.Create(&workspace)
 	require.NoError(t, tx.Error)
 
+	tx = conn.Create(&membership)
+	require.NoError(t, tx.Error)
+
 	reconciler := NewUsageReconciler(conn)
 
 	status, err := reconciler.ReconcileTimeRange(context.Background(), startOfMay, startOfJune)
@@ -61,5 +75,6 @@ func TestUsageReconciler_Reconcile(t *testing.T) {
 		WorkspaceInstances:        2,
 		InvalidWorkspaceInstances: 1,
 		Workspaces:                1,
+		Teams:                     1,
 	}, status)
 }