add test for template insights

coder · mafredri · Jul 21, 2023 · Jul 14, 2023 · Jul 18, 2023 · Jul 19, 2023
commit ddc86062fbebc913cd18f49cf0cd40a0a63002a4
diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -1917,22 +1917,127 @@ func (q *FakeQuerier) GetTemplateDAUs(_ context.Context, arg database.GetTemplat
 	return rs, nil
 }
 
-func (q *FakeQuerier) GetTemplateDailyInsights(ctx context.Context, arg database.GetTemplateDailyInsightsParams) ([]database.GetTemplateDailyInsightsRow, error) {
+func (q *FakeQuerier) GetTemplateDailyInsights(_ context.Context, arg database.GetTemplateDailyInsightsParams) ([]database.GetTemplateDailyInsightsRow, error) {
 	err := validateDatabaseType(arg)
 	if err != nil {
 		return nil, err
 	}
 
-	panic("not implemented")
+	type dailyStat struct {
+		startTime, endTime time.Time
+		userSet            map[uuid.UUID]struct{}
+		templateIDSet      map[uuid.UUID]struct{}
+	}
+	dailyStats := []dailyStat{{arg.StartTime, arg.StartTime.AddDate(0, 0, 1), make(map[uuid.UUID]struct{}), make(map[uuid.UUID]struct{})}}
+	for dailyStats[len(dailyStats)-1].endTime.Before(arg.EndTime) {
+		dailyStats = append(dailyStats, dailyStat{dailyStats[len(dailyStats)-1].endTime, dailyStats[len(dailyStats)-1].endTime.AddDate(0, 0, 1), make(map[uuid.UUID]struct{}), make(map[uuid.UUID]struct{})})
+	}
+	if dailyStats[len(dailyStats)-1].endTime.After(arg.EndTime) {
+		dailyStats[len(dailyStats)-1].endTime = arg.EndTime
+	}
+
+	for _, s := range q.workspaceAgentStats {
+		if s.CreatedAt.Before(arg.StartTime) || s.CreatedAt.Equal(arg.EndTime) || s.CreatedAt.After(arg.EndTime) {
+			continue
+		}
+		if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, s.TemplateID) {
+			continue
+		}
+		if s.ConnectionCount == 0 {
+			continue
+		}
+
+		for _, ds := range dailyStats {
+			if s.CreatedAt.Before(ds.startTime) || s.CreatedAt.Equal(ds.endTime) || s.CreatedAt.After(ds.endTime) {
+				continue
+			}
+			ds.userSet[s.UserID] = struct{}{}
+			ds.templateIDSet[s.TemplateID] = struct{}{}
+			break
+		}
+	}
+
+	var result []database.GetTemplateDailyInsightsRow
+	for _, ds := range dailyStats {
+		templateIDs := make([]uuid.UUID, 0, len(ds.templateIDSet))
+		for templateID := range ds.templateIDSet {
+			templateIDs = append(templateIDs, templateID)
+		}
+		slices.SortFunc(templateIDs, func(a, b uuid.UUID) bool {
+			return a.String() < b.String()
+		})
+		result = append(result, database.GetTemplateDailyInsightsRow{
+			StartTime:   ds.startTime,
+			EndTime:     ds.endTime,
+			TemplateIDs: templateIDs,
+			ActiveUsers: int64(len(ds.userSet)),
+		})
+	}
+	return result, nil
 }
 
-func (q *FakeQuerier) GetTemplateInsights(ctx context.Context, arg database.GetTemplateInsightsParams) (database.GetTemplateInsightsRow, error) {
+func (q *FakeQuerier) GetTemplateInsights(_ context.Context, arg database.GetTemplateInsightsParams) (database.GetTemplateInsightsRow, error) {
 	err := validateDatabaseType(arg)
 	if err != nil {
 		return database.GetTemplateInsightsRow{}, err
 	}
 
-	panic("not implemented")
+	templateIDSet := make(map[uuid.UUID]struct{})
+	appUsageIntervalsByUser := make(map[uuid.UUID]map[time.Time]*database.GetTemplateInsightsRow)
+	for _, s := range q.workspaceAgentStats {
+		if s.CreatedAt.Before(arg.StartTime) || s.CreatedAt.Equal(arg.EndTime) || s.CreatedAt.After(arg.EndTime) {
+			continue
+		}
+		if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, s.TemplateID) {
+			continue
+		}
+		if s.ConnectionCount == 0 {
+			continue
+		}
+
+		templateIDSet[s.TemplateID] = struct{}{}
+		if appUsageIntervalsByUser[s.UserID] == nil {
+			appUsageIntervalsByUser[s.UserID] = make(map[time.Time]*database.GetTemplateInsightsRow)
+		}
+		t := s.CreatedAt.Truncate(5 * time.Minute)
+		if _, ok := appUsageIntervalsByUser[s.UserID][t]; !ok {
+			appUsageIntervalsByUser[s.UserID][t] = &database.GetTemplateInsightsRow{}
+		}
+
+		if s.SessionCountJetBrains > 0 {
+			appUsageIntervalsByUser[s.UserID][t].UsageJetbrainsSeconds = 300
+		}
+		if s.SessionCountVSCode > 0 {
+			appUsageIntervalsByUser[s.UserID][t].UsageVscodeSeconds = 300
+		}
+		if s.SessionCountReconnectingPTY > 0 {
+			appUsageIntervalsByUser[s.UserID][t].UsageReconnectingPtySeconds = 300
+		}
+		if s.SessionCountSSH > 0 {
+			appUsageIntervalsByUser[s.UserID][t].UsageSshSeconds = 300
+		}
+	}
+
+	templateIDs := make([]uuid.UUID, 0, len(templateIDSet))
+	for templateID := range templateIDSet {
+		templateIDs = append(templateIDs, templateID)
+	}
+	slices.SortFunc(templateIDs, func(a, b uuid.UUID) bool {
+		return a.String() < b.String()
+	})
+	result := database.GetTemplateInsightsRow{
+		TemplateIDs: templateIDs,
+		ActiveUsers: int64(len(appUsageIntervalsByUser)),
+	}
+	for _, intervals := range appUsageIntervalsByUser {
+		for _, interval := range intervals {
+			result.UsageJetbrainsSeconds += interval.UsageJetbrainsSeconds
+			result.UsageVscodeSeconds += interval.UsageVscodeSeconds
+			result.UsageReconnectingPtySeconds += interval.UsageReconnectingPtySeconds
+			result.UsageSshSeconds += interval.UsageSshSeconds
+		}
+	}
+	return result, nil
 }
 
 func (q *FakeQuerier) GetTemplateVersionByID(ctx context.Context, templateVersionID uuid.UUID) (database.TemplateVersion, error) {
@@ -2182,7 +2287,7 @@ func (q *FakeQuerier) GetUserCount(_ context.Context) (int64, error) {
 	return existing, nil
 }
 
-func (q *FakeQuerier) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
+func (q *FakeQuerier) GetUserLatencyInsights(_ context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
 	err := validateDatabaseType(arg)
 	if err != nil {
 		return nil, err
@@ -2222,9 +2327,9 @@ func (q *FakeQuerier) GetUserLatencyInsights(ctx context.Context, arg database.G
 	var rows []database.GetUserLatencyInsightsRow
 	for userID, latencies := range latenciesByUserID {
 		sort.Float64s(latencies)
-		templateSet := seenTemplatesByUserID[userID]
-		templateIDs := make([]uuid.UUID, 0, len(templateSet))
-		for templateID := range templateSet {
+		templateIDSet := seenTemplatesByUserID[userID]
+		templateIDs := make([]uuid.UUID, 0, len(templateIDSet))
+		for templateID := range templateIDSet {
 			templateIDs = append(templateIDs, templateID)
 		}
 		slices.SortFunc(templateIDs, func(a, b uuid.UUID) bool {

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/insights.sql b/coderd/database/queries/insights.sql
@@ -38,21 +38,27 @@ WITH d AS (
 		CASE WHEN SUM(was.session_count_reconnecting_pty) > 0 THEN ts.seconds ELSE 0 END AS usage_reconnecting_pty_seconds,
 		CASE WHEN SUM(was.session_count_ssh) > 0 THEN ts.seconds ELSE 0 END AS usage_ssh_seconds
 	FROM ts
-	JOIN workspace_agent_stats was ON (was.created_at >= ts.from_ AND was.created_at < ts.to_)
-	WHERE
-		was.connection_count > 0
+	JOIN workspace_agent_stats was ON (
+		was.created_at >= ts.from_
+		AND was.created_at < ts.to_
+		AND was.connection_count > 0
 		AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN was.template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
+	)
 	GROUP BY ts.from_, ts.to_, ts.seconds, was.user_id
+), template_ids AS (
+	SELECT array_agg(DISTINCT template_id) AS ids
+	FROM usage_by_user, unnest(template_ids) template_id
+	WHERE template_id IS NOT NULL
 )
 
 SELECT
+	COALESCE((SELECT ids FROM template_ids), '{}')::uuid[] AS template_ids,
 	COUNT(DISTINCT user_id) AS active_users,
-	COALESCE(array_agg(DISTINCT template_id), '{}')::uuid[] AS template_ids,
 	COALESCE(SUM(usage_vscode_seconds), 0)::bigint AS usage_vscode_seconds,
 	COALESCE(SUM(usage_jetbrains_seconds), 0)::bigint AS usage_jetbrains_seconds,
 	COALESCE(SUM(usage_reconnecting_pty_seconds), 0)::bigint AS usage_reconnecting_pty_seconds,
 	COALESCE(SUM(usage_ssh_seconds), 0)::bigint AS usage_ssh_seconds
-FROM usage_by_user, unnest(template_ids) as template_id;
+FROM usage_by_user;
 
 -- name: GetTemplateDailyInsights :many
 WITH d AS (
@@ -69,17 +75,23 @@ WITH d AS (
 		was.user_id,
 		array_agg(was.template_id) AS template_ids
 	FROM ts
-	LEFT JOIN workspace_agent_stats was ON (was.created_at >= ts.from_ AND was.created_at < ts.to_)
-	WHERE
-		was.connection_count > 0
+	LEFT JOIN workspace_agent_stats was ON (
+		was.created_at >= ts.from_
+		AND was.created_at < ts.to_
+		AND was.connection_count > 0
 		AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN was.template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
+	)
 	GROUP BY ts.from_, ts.to_, was.user_id
+), template_ids AS (
+	SELECT array_agg(DISTINCT template_id) AS ids
+	FROM usage_by_day, unnest(template_ids) template_id
+	WHERE template_id IS NOT NULL
 )
 
 SELECT
 	from_ AS start_time,
 	to_ AS end_time,
-	COUNT(DISTINCT user_id) AS active_users,
-	array_agg(DISTINCT template_id)::uuid[] AS template_ids
+	COALESCE((SELECT ids FROM template_ids), '{}')::uuid[] AS template_ids,
+	COUNT(DISTINCT user_id) AS active_users
 FROM usage_by_day, unnest(template_ids) as template_id
 GROUP BY from_, to_;
diff --git a/coderd/insights.go b/coderd/insights.go
@@ -129,7 +129,6 @@ func (api *API) insightsUserLatency(rw http.ResponseWriter, r *http.Request) {
 	for _, user := range users {
 		userLatency, ok := usersWithLatencyByID[user.ID]
 		if !ok {
-			// TODO(mafredri): Other cases?
 			// We only include deleted/inactive users if they were
 			// active as part of the requested timeframe.
 			if user.Deleted || user.Status != database.UserStatusActive {
@@ -262,31 +261,31 @@ func (api *API) insightsTemplates(rw http.ResponseWriter, r *http.Request) {
 			ActiveUsers: usage.ActiveUsers,
 			AppsUsage: []codersdk.TemplateAppUsage{
 				{
-					TemplateIDs: usage.TemplateIDs, // TODO(mafredri): Update query to return template IDs/app?
+					TemplateIDs: usage.TemplateIDs,
 					Type:        codersdk.TemplateAppsTypeBuiltin,
 					DisplayName: "Visual Studio Code",
 					Slug:        "vscode",
 					Icon:        "/icons/code.svg",
 					Seconds:     usage.UsageVscodeSeconds,
 				},
 				{
-					TemplateIDs: usage.TemplateIDs, // TODO(mafredri): Update query to return template IDs/app?
+					TemplateIDs: usage.TemplateIDs,
 					Type:        codersdk.TemplateAppsTypeBuiltin,
 					DisplayName: "JetBrains",
 					Slug:        "jetbrains",
 					Icon:        "/icons/intellij.svg",
 					Seconds:     usage.UsageJetbrainsSeconds,
 				},
 				{
-					TemplateIDs: usage.TemplateIDs, // TODO(mafredri): Update query to return template IDs/app?
+					TemplateIDs: usage.TemplateIDs,
 					Type:        codersdk.TemplateAppsTypeBuiltin,
 					DisplayName: "Web Terminal",
 					Slug:        "reconnecting-pty",
 					Icon:        "/icons/terminal.svg",
 					Seconds:     usage.UsageReconnectingPtySeconds,
 				},
 				{
-					TemplateIDs: usage.TemplateIDs, // TODO(mafredri): Update query to return template IDs/app?
+					TemplateIDs: usage.TemplateIDs,
 					Type:        codersdk.TemplateAppsTypeBuiltin,
 					DisplayName: "SSH",
 					Slug:        "ssh",