use window functions for efficiency

coder · SasSwart · Dec 14, 2024 · Dec 14, 2024 · Dec 14, 2024 · Dec 17, 2024
commit 0f17038b67da177e5d933a4222e02aa3f6fa390b
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
@@ -2258,6 +2258,10 @@ func TestGroupRemovalTrigger(t *testing.T) {
 func TestGetUserStatusChanges(t *testing.T) {
 	t.Parallel()
 
+	now := dbtime.Now()
+	createdAt := now.Add(-5 * 24 * time.Hour)                           // 5 days ago
+	firstTransitionTime := createdAt.Add(2 * 24 * time.Hour)            // 3 days ago
+	secondTransitionTime := firstTransitionTime.Add(2 * 24 * time.Hour) // 1 days ago
 	t.Run("No Users", func(t *testing.T) {
 		t.Parallel()
 		db, _ := dbtestutil.NewDB(t)
@@ -2307,8 +2311,6 @@ func TestGetUserStatusChanges(t *testing.T) {
 				ctx := testutil.Context(t, testutil.WaitShort)
 
 				// Create a user that's been in the specified status for the past 30 days
-				now := dbtime.Now()
-				createdAt := now.Add(-29 * 24 * time.Hour)
 				dbgen.User(t, db, database.User{
 					Status:    tc.status,
 					CreatedAt: createdAt,
@@ -2324,8 +2326,10 @@ func TestGetUserStatusChanges(t *testing.T) {
 				require.NotEmpty(t, userStatusChanges, "should return results")
 
 				// We should have an entry for each status change
-				require.Len(t, userStatusChanges, 1, "should have 1 status change")
-				require.Equal(t, userStatusChanges[0].NewStatus, tc.status, "should have the correct status")
+				require.Len(t, userStatusChanges, 5, "should have 1 user * 5 days = 5 rows")
+				for _, row := range userStatusChanges {
+					require.Equal(t, row.NewStatus, tc.status, "should have the correct status")
+				}
 			})
 		}
 	})
@@ -2393,20 +2397,17 @@ func TestGetUserStatusChanges(t *testing.T) {
 				ctx := testutil.Context(t, testutil.WaitShort)
 
 				// Create a user that starts with initial status
-				now := dbtime.Now()
-				createdAt := now.Add(-5 * 24 * time.Hour) // 5 days ago
 				user := dbgen.User(t, db, database.User{
 					Status:    tc.initialStatus,
 					CreatedAt: createdAt,
 					UpdatedAt: createdAt,
 				})
 
 				// After 2 days, change status to target status
-				statusChangeTime := createdAt.Add(2 * 24 * time.Hour)
 				user, err := db.UpdateUserStatus(ctx, database.UpdateUserStatusParams{
 					ID:        user.ID,
 					Status:    tc.targetStatus,
-					UpdatedAt: statusChangeTime,
+					UpdatedAt: firstTransitionTime,
 				})
 				require.NoError(t, err)
 
@@ -2418,10 +2419,15 @@ func TestGetUserStatusChanges(t *testing.T) {
 				require.NoError(t, err)
 				require.NotEmpty(t, userStatusChanges, "should return results")
 
-				// We should have an entry for each status change, including the initial status
-				require.Len(t, userStatusChanges, 2, "should have 2 status changes")
-				require.Equal(t, userStatusChanges[0].NewStatus, tc.initialStatus, "should have the initial status")
-				require.Equal(t, userStatusChanges[1].NewStatus, tc.targetStatus, "should have the target status")
+				// We should have an entry for each status (active, dormant, suspended) for each day
+				require.Len(t, userStatusChanges, 5, "should have 1 user * 5 days = 5 rows")
+				for _, row := range userStatusChanges {
+					if row.ChangedAt.Before(firstTransitionTime) {
+						require.Equal(t, row.NewStatus, tc.initialStatus, "should have the initial status")
+					} else {
+						require.Equal(t, row.NewStatus, tc.targetStatus, "should have the target status")
+					}
+				}
 			})
 		}
 	})
@@ -2606,20 +2612,18 @@ func TestGetUserStatusChanges(t *testing.T) {
 				})
 
 				// First transition at 2 days
-				user1TransitionTime := createdAt.Add(2 * 24 * time.Hour)
 				user1, err := db.UpdateUserStatus(ctx, database.UpdateUserStatusParams{
 					ID:        user1.ID,
 					Status:    tc.user1Transition.to,
-					UpdatedAt: user1TransitionTime,
+					UpdatedAt: firstTransitionTime,
 				})
 				require.NoError(t, err)
 
 				// Second transition at 4 days
-				user2TransitionTime := createdAt.Add(4 * 24 * time.Hour)
 				user2, err = db.UpdateUserStatus(ctx, database.UpdateUserStatusParams{
 					ID:        user2.ID,
 					Status:    tc.user2Transition.to,
-					UpdatedAt: user2TransitionTime,
+					UpdatedAt: secondTransitionTime,
 				})
 				require.NoError(t, err)
 
@@ -2631,16 +2635,17 @@ func TestGetUserStatusChanges(t *testing.T) {
 				require.NoError(t, err)
 				require.NotEmpty(t, userStatusChanges)
 
-				// We should have an entry with the correct status changes for each user, including the initial status
-				require.Len(t, userStatusChanges, 4, "should have 4 status changes")
-				require.Equal(t, userStatusChanges[0].UserID, user1.ID, "should have the first user")
-				require.Equal(t, userStatusChanges[0].NewStatus, tc.user1Transition.from, "should have the first user's initial status")
-				require.Equal(t, userStatusChanges[1].UserID, user1.ID, "should have the first user")
-				require.Equal(t, userStatusChanges[1].NewStatus, tc.user1Transition.to, "should have the first user's target status")
-				require.Equal(t, userStatusChanges[2].UserID, user2.ID, "should have the second user")
-				require.Equal(t, userStatusChanges[2].NewStatus, tc.user2Transition.from, "should have the second user's initial status")
-				require.Equal(t, userStatusChanges[3].UserID, user2.ID, "should have the second user")
-				require.Equal(t, userStatusChanges[3].NewStatus, tc.user2Transition.to, "should have the second user's target status")
+				// Expected counts before, between and after the transitions should match:
+				for _, row := range userStatusChanges {
+					switch {
+					case row.ChangedAt.Before(firstTransitionTime):
+						require.Equal(t, row.Count, tc.expectedCounts["initial"][row.NewStatus], "should have the correct count before the first transition")
+					case row.ChangedAt.Before(secondTransitionTime):
+						require.Equal(t, row.Count, tc.expectedCounts["between"][row.NewStatus], "should have the correct count between the transitions")
+					case row.ChangedAt.Before(now):
+						require.Equal(t, row.Count, tc.expectedCounts["final"][row.NewStatus], "should have the correct count after the second transition")
+					}
+				}
 			})
 		}
 	})

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
@@ -773,36 +773,64 @@ JOIN workspace_build_parameters wbp ON (utp.workspace_build_ids @> ARRAY[wbp.wor
 GROUP BY utp.num, utp.template_ids, utp.name, utp.type, utp.display_name, utp.description, utp.options, wbp.value;
 
 -- name: GetUserStatusChanges :many
-WITH last_status_per_day AS (
-    -- First get the last status change for each user for each day
-    SELECT DISTINCT ON (date_trunc('day', changed_at), user_id)
-        date_trunc('day', changed_at)::timestamptz AS date,
+WITH dates AS (
+    SELECT generate_series(
+        date_trunc('day', @start_time::timestamptz),
+        date_trunc('day', @end_time::timestamptz),
+        '1 day'::interval
+    )::timestamptz AS date
+),
+latest_status_before_range AS (
+    -- Get the last status change for each user before our date range
+    SELECT DISTINCT ON (user_id)
+        user_id,
         new_status,
-        user_id
+        changed_at
     FROM user_status_changes
-    WHERE changed_at >= @start_time::timestamptz
-        AND changed_at < @end_time::timestamptz
-    ORDER BY
-        date_trunc('day', changed_at),
+    WHERE changed_at < (SELECT MIN(date) FROM dates)
+    ORDER BY user_id, changed_at DESC
+),
+all_status_changes AS (
+    -- Combine status changes before and during our range
+    SELECT
+        user_id,
+        new_status,
+        changed_at
+    FROM latest_status_before_range
+
+    UNION ALL
+
+    SELECT
         user_id,
-        changed_at DESC  -- This ensures we get the last status for each day
+        new_status,
+        changed_at
+    FROM user_status_changes
+    WHERE changed_at < @end_time::timestamptz
 ),
 daily_counts AS (
-    -- Then count unique users per status per day
     SELECT
-        date,
-        new_status,
-        COUNT(*) AS count
-    FROM last_status_per_day
-    GROUP BY
-        date,
-        new_status
+        d.date,
+        asc1.new_status,
+        -- For each date and status, count users whose most recent status change
+        -- (up to that date) matches this status
+        COUNT(*) FILTER (
+            WHERE asc1.changed_at = (
+                SELECT MAX(changed_at)
+                FROM all_status_changes asc2
+                WHERE asc2.user_id = asc1.user_id
+                AND asc2.changed_at <= d.date
+            )
+        )::bigint AS count
+    FROM dates d
+    CROSS JOIN all_status_changes asc1
+    GROUP BY d.date, asc1.new_status
 )
 SELECT
-    date::timestamptz AS changed_at,
+    date AS changed_at,
     new_status,
-    count::bigint
+    count
 FROM daily_counts
+WHERE count > 0
 ORDER BY
     new_status ASC,
     date ASC;