engine/cleanupmgr: Migrate shifts to job queue (#4192)

* cleanup: implement context-aware sleep and cleanup manager functionality

* cleanup: format code for consistency in alert store initialization

* cleanup: remove unused ConfigSource field from SetupArgs struct

* cleanup: refactor alert cleanup logic and add shift cleanup functionality

* cleanup: refactor alert cleanup logic to use whileWork for better control flow

* cleanup: add comment

* cleanup: remove unused cleanup statements from DB and update logic

* test: refactor alert auto-close and cleanup tests for improved reliability

* fix: update ShiftArgs Kind to reflect cleanup-manager-shifts

engine/cleanupmanager/db.go

@@ -27,10 +27,7 @@ type DB struct {
 
 	cleanupAlertLogs *sql.Stmt
 
-	cleanupOverrides   *sql.Stmt
-	cleanupSchedOnCall *sql.Stmt
-	cleanupEPOnCall    *sql.Stmt
-	alertStore         *alert.Store
+	alertStore *alert.Store
 
 	logIndex int
 }
@@ -88,10 +85,6 @@ func NewDB(ctx context.Context, db *sql.DB, alertstore *alert.Store) (*DB, error
 			select id from scope offset 99
 		`),
 
-		cleanupOverrides:   p.P(`DELETE FROM user_overrides WHERE id = ANY(SELECT id FROM user_overrides WHERE end_time < (now() - $1::interval) LIMIT 100 FOR UPDATE SKIP LOCKED)`),
-		cleanupSchedOnCall: p.P(`DELETE FROM schedule_on_call_users WHERE id = ANY(SELECT id FROM schedule_on_call_users WHERE end_time < (now() - $1::interval) LIMIT 100 FOR UPDATE SKIP LOCKED)`),
-		cleanupEPOnCall:    p.P(`DELETE FROM ep_step_on_call_users WHERE id = ANY(SELECT id FROM ep_step_on_call_users WHERE end_time < (now() - $1::interval) LIMIT 100 FOR UPDATE SKIP LOCKED)`),
-
 		alertStore: alertstore,
 	}, p.Err
 }

engine/cleanupmanager/queries.sql

@@ -35,3 +35,45 @@ AND NOT EXISTS (
         AND log.alert_id = a.id)
 LIMIT 100;
 
+-- name: CleanupMgrDeleteOldOverrides :execrows
+-- CleanupMgrDeleteOldOverrides will delete old overrides from the user_overrides table that are older than the given number of days before now.
+DELETE FROM user_overrides
+WHERE id = ANY (
+        SELECT
+            id
+        FROM
+            user_overrides
+        WHERE
+            end_time <(now() - '1 day'::interval * sqlc.arg(cleanup_days))
+        LIMIT 100
+        FOR UPDATE
+            SKIP LOCKED);
+
+-- name: CleanupMgrDeleteOldScheduleShifts :execrows
+-- CleanupMgrDeleteOldScheduleShifts will delete old schedule shifts from the schedule_on_call_users table that are older than the given number of days before now.
+DELETE FROM schedule_on_call_users
+WHERE id = ANY (
+        SELECT
+            id
+        FROM
+            schedule_on_call_users
+        WHERE
+            end_time <(now() - '1 day'::interval * sqlc.arg(cleanup_days))
+        LIMIT 100
+        FOR UPDATE
+            SKIP LOCKED);
+
+-- name: CleanupMgrDeleteOldStepShifts :execrows
+-- CleanupMgrDeleteOldStepShifts will delete old EP step shifts from the ep_step_on_call_users table that are older than the given number of days before now.
+DELETE FROM ep_step_on_call_users
+WHERE id = ANY (
+        SELECT
+            id
+        FROM
+            ep_step_on_call_users
+        WHERE
+            end_time <(now() - '1 day'::interval * sqlc.arg(cleanup_days))
+        LIMIT 100
+        FOR UPDATE
+            SKIP LOCKED);
+

engine/cleanupmanager/shifts.go

@@ -0,0 +1,58 @@
+package cleanupmanager
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+
+	"github.com/riverqueue/river"
+	"github.com/target/goalert/config"
+	"github.com/target/goalert/gadb"
+)
+
+type ShiftArgs struct{}
+
+func (ShiftArgs) Kind() string { return "cleanup-manager-shifts" }
+
+// CleanupShifts will automatically cleanup old shift and override records.
+func (db *DB) CleanupShifts(ctx context.Context, j *river.Job[ShiftArgs]) error {
+	cfg := config.FromContext(ctx)
+	if cfg.Maintenance.ScheduleCleanupDays <= 0 {
+		return nil
+	}
+
+	err := db.whileWork(ctx, func(ctx context.Context, tx *sql.Tx) (done bool, err error) {
+		count, err := gadb.New(tx).CleanupMgrDeleteOldScheduleShifts(ctx, int64(cfg.Maintenance.ScheduleCleanupDays))
+		if err != nil {
+			return false, fmt.Errorf("delete old shifts: %w", err)
+		}
+		return count < 100, nil
+	})
+	if err != nil {
+		return err
+	}
+
+	err = db.whileWork(ctx, func(ctx context.Context, tx *sql.Tx) (done bool, err error) {
+		count, err := gadb.New(tx).CleanupMgrDeleteOldOverrides(ctx, int64(cfg.Maintenance.ScheduleCleanupDays))
+		if err != nil {
+			return false, fmt.Errorf("delete old overrides: %w", err)
+		}
+		return count < 100, nil
+	})
+	if err != nil {
+		return err
+	}
+
+	err = db.whileWork(ctx, func(ctx context.Context, tx *sql.Tx) (done bool, err error) {
+		count, err := gadb.New(tx).CleanupMgrDeleteOldStepShifts(ctx, int64(cfg.Maintenance.ScheduleCleanupDays))
+		if err != nil {
+			return false, fmt.Errorf("delete old step shifts: %w", err)
+		}
+		return count < 100, nil
+	})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

engine/cleanupmanager/update.go

@@ -63,25 +63,6 @@ func (db *DB) update(ctx context.Context) error {
 			return err
 		}
 	}
-	if cfg.Maintenance.ScheduleCleanupDays > 0 {
-		var dur pgtype.Interval
-		dur.Days = int32(cfg.Maintenance.ScheduleCleanupDays)
-		dur.Status = pgtype.Present
-		_, err = tx.StmtContext(ctx, db.cleanupOverrides).ExecContext(ctx, &dur)
-		if err != nil {
-			return fmt.Errorf("cleanup overrides: %w", err)
-		}
-
-		_, err = tx.StmtContext(ctx, db.cleanupSchedOnCall).ExecContext(ctx, &dur)
-		if err != nil {
-			return fmt.Errorf("cleanup schedule on-call: %w", err)
-		}
-
-		_, err = tx.StmtContext(ctx, db.cleanupEPOnCall).ExecContext(ctx, &dur)
-		if err != nil {
-			return fmt.Errorf("cleanup escalation policy on-call: %w", err)
-		}
-	}
 
 	rows, err := tx.StmtContext(ctx, db.schedData).QueryContext(ctx)
 	if err != nil {