feat: purge expired api keys in dbpurge (#20863)

closes https://github.com/coder/coder/issues/19889

This is in response to a migration in v2.27 that takes very long on deployments with large `api_key` tables.

coderd/database/dbauthz/dbauthz.go

@@ -1641,6 +1641,15 @@ func (q *querier) DeleteCustomRole(ctx context.Context, arg database.DeleteCusto
 	return q.db.DeleteCustomRole(ctx, arg)
 }
 
+func (q *querier) DeleteExpiredAPIKeys(ctx context.Context, arg database.DeleteExpiredAPIKeysParams) (int64, error) {
+	// Requires DELETE across all API keys.
+	if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceApiKey); err != nil {
+		return 0, err
+	}
+
+	return q.db.DeleteExpiredAPIKeys(ctx, arg)
+}
+
 func (q *querier) DeleteExternalAuthLink(ctx context.Context, arg database.DeleteExternalAuthLinkParams) error {
 	return fetchAndExec(q.log, q.auth, policy.ActionUpdatePersonal, func(ctx context.Context, arg database.DeleteExternalAuthLinkParams) (database.ExternalAuthLink, error) {
 		//nolint:gosimple

coderd/database/dbauthz/dbauthz_test.go

@@ -216,6 +216,14 @@ func (s *MethodTestSuite) TestAPIKey() {
 		dbm.EXPECT().DeleteAPIKeyByID(gomock.Any(), key.ID).Return(nil).AnyTimes()
 		check.Args(key.ID).Asserts(key, policy.ActionDelete).Returns()
 	}))
+	s.Run("DeleteExpiredAPIKeys", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		args := database.DeleteExpiredAPIKeysParams{
+			Before:     time.Date(2025, 11, 21, 0, 0, 0, 0, time.UTC),
+			LimitCount: 1000,
+		}
+		dbm.EXPECT().DeleteExpiredAPIKeys(gomock.Any(), args).Return(int64(0), nil).AnyTimes()
+		check.Args(args).Asserts(rbac.ResourceApiKey, policy.ActionDelete).Returns(int64(0))
+	}))
 	s.Run("GetAPIKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		key := testutil.Fake(s.T(), faker, database.APIKey{})
 		dbm.EXPECT().GetAPIKeyByID(gomock.Any(), key.ID).Return(key, nil).AnyTimes()

coderd/database/dbgen/dbgen.go

@@ -175,6 +175,13 @@ func APIKey(t testing.TB, db database.Store, seed database.APIKey, munge ...func
 		}
 	}
 
+	// It does not make sense for the created_at to be after the expires_at.
+	// So if expires is set, change the default created_at to be 24 hours before.
+	var createdAt time.Time
+	if !seed.ExpiresAt.IsZero() && seed.CreatedAt.IsZero() {
+		createdAt = seed.ExpiresAt.Add(-24 * time.Hour)
+	}
+
 	params := database.InsertAPIKeyParams{
 		ID: takeFirst(seed.ID, id),
 		// 0 defaults to 86400 at the db layer
@@ -184,7 +191,7 @@ func APIKey(t testing.TB, db database.Store, seed database.APIKey, munge ...func
 		UserID:          takeFirst(seed.UserID, uuid.New()),
 		LastUsed:        takeFirst(seed.LastUsed, dbtime.Now()),
 		ExpiresAt:       takeFirst(seed.ExpiresAt, dbtime.Now().Add(time.Hour)),
-		CreatedAt:       takeFirst(seed.CreatedAt, dbtime.Now()),
+		CreatedAt:       takeFirst(seed.CreatedAt, createdAt, dbtime.Now()),
 		UpdatedAt:       takeFirst(seed.UpdatedAt, dbtime.Now()),
 		LoginType:       takeFirst(seed.LoginType, database.LoginTypePassword),
 		Scopes:          takeFirstSlice([]database.APIKeyScope(seed.Scopes), []database.APIKeyScope{database.ApiKeyScopeCoderAll}),

coderd/database/dbmetrics/querymetrics.go

@@ -312,6 +312,13 @@ func (m queryMetricsStore) DeleteCustomRole(ctx context.Context, arg database.De
 	return r0
 }
 
+func (m queryMetricsStore) DeleteExpiredAPIKeys(ctx context.Context, arg database.DeleteExpiredAPIKeysParams) (int64, error) {
+	start := time.Now()
+	r0, r1 := m.s.DeleteExpiredAPIKeys(ctx, arg)
+	m.queryLatencies.WithLabelValues("DeleteExpiredAPIKeys").Observe(time.Since(start).Seconds())
+	return r0, r1
+}
+
 func (m queryMetricsStore) DeleteExternalAuthLink(ctx context.Context, arg database.DeleteExternalAuthLinkParams) error {
 	start := time.Now()
 	r0 := m.s.DeleteExternalAuthLink(ctx, arg)

coderd/database/dbmock/dbmock.go

@@ -554,6 +554,21 @@ func (mr *MockStoreMockRecorder) DeleteCustomRole(ctx, arg any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteCustomRole", reflect.TypeOf((*MockStore)(nil).DeleteCustomRole), ctx, arg)
 }
 
+// DeleteExpiredAPIKeys mocks base method.
+func (m *MockStore) DeleteExpiredAPIKeys(ctx context.Context, arg database.DeleteExpiredAPIKeysParams) (int64, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DeleteExpiredAPIKeys", ctx, arg)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// DeleteExpiredAPIKeys indicates an expected call of DeleteExpiredAPIKeys.
+func (mr *MockStoreMockRecorder) DeleteExpiredAPIKeys(ctx, arg any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteExpiredAPIKeys", reflect.TypeOf((*MockStore)(nil).DeleteExpiredAPIKeys), ctx, arg)
+}
+
 // DeleteExternalAuthLink mocks base method.
 func (m *MockStore) DeleteExternalAuthLink(ctx context.Context, arg database.DeleteExternalAuthLinkParams) error {
 	m.ctrl.T.Helper()

coderd/database/dbpurge/dbpurge.go

@@ -78,6 +78,19 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
 			if err := tx.ExpirePrebuildsAPIKeys(ctx, dbtime.Time(start)); err != nil {
 				return xerrors.Errorf("failed to expire prebuilds user api keys: %w", err)
 			}
+			expiredAPIKeys, err := tx.DeleteExpiredAPIKeys(ctx, database.DeleteExpiredAPIKeysParams{
+				// Leave expired keys for a week to allow the backend to know the difference
+				// between a 404 and an expired key. This purge code is just to bound the size of
+				// the table to something more reasonable.
+				Before: dbtime.Time(start.Add(time.Hour * 24 * 7 * -1)),
+				// There could be a lot of expired keys here, so set a limit to prevent this
+				// taking too long.
+				// This runs every 10 minutes, so it deletes ~1.5m keys per day at most.
+				LimitCount: 10000,
+			})
+			if err != nil {
+				return xerrors.Errorf("failed to delete expired api keys: %w", err)
+			}
 			deleteOldTelemetryLocksBefore := start.Add(-maxTelemetryHeartbeatAge)
 			if err := tx.DeleteOldTelemetryLocks(ctx, deleteOldTelemetryLocksBefore); err != nil {
 				return xerrors.Errorf("failed to delete old telemetry locks: %w", err)
@@ -93,13 +106,16 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
 
 			deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value())
 			// nolint:gocritic // Needs to run as aibridge context.
-			count, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
+			purgedAIBridgeRecords, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
 			if err != nil {
 				return xerrors.Errorf("failed to delete old aibridge records: %w", err)
 			}
-			logger.Debug(ctx, "purged aibridge entries", slog.F("count", count), slog.F("since", deleteAIBridgeRecordsBefore.Format(time.RFC3339)))
 
-			logger.Debug(ctx, "purged old database entries", slog.F("duration", clk.Since(start)))
+			logger.Debug(ctx, "purged old database entries",
+				slog.F("expired_api_keys", expiredAPIKeys),
+				slog.F("aibridge_records", purgedAIBridgeRecords),
+				slog.F("duration", clk.Since(start)),
+			)
 
 			return nil
 		}, database.DefaultTXOptions().WithID("db_purge")); err != nil {

coderd/database/querier.go

@@ -91,6 +91,7 @@ type sqlcQuerier interface {
 	DeleteCoordinator(ctx context.Context, id uuid.UUID) error
 	DeleteCryptoKey(ctx context.Context, arg DeleteCryptoKeyParams) (CryptoKey, error)
 	DeleteCustomRole(ctx context.Context, arg DeleteCustomRoleParams) error
+	DeleteExpiredAPIKeys(ctx context.Context, arg DeleteExpiredAPIKeysParams) (int64, error)
 	DeleteExternalAuthLink(ctx context.Context, arg DeleteExternalAuthLinkParams) error
 	DeleteGitSSHKey(ctx context.Context, userID uuid.UUID) error
 	DeleteGroupByID(ctx context.Context, id uuid.UUID) error

coderd/database/querier_test.go

@@ -7835,3 +7835,81 @@ func TestUpdateAIBridgeInterceptionEnded(t *testing.T) {
 		}
 	})
 }
+
+func TestDeleteExpiredAPIKeys(t *testing.T) {
+	t.Parallel()
+	db, _ := dbtestutil.NewDB(t)
+
+	// Constant time for testing
+	now := time.Date(2025, 11, 20, 12, 0, 0, 0, time.UTC)
+	expiredBefore := now.Add(-time.Hour) // Anything before this is expired
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	user := dbgen.User(t, db, database.User{})
+
+	expiredTimes := []time.Time{
+		expiredBefore.Add(-time.Hour * 24 * 365),
+		expiredBefore.Add(-time.Hour * 24),
+		expiredBefore.Add(-time.Hour),
+		expiredBefore.Add(-time.Minute),
+		expiredBefore.Add(-time.Second),
+	}
+	for _, exp := range expiredTimes {
+		// Expired api keys
+		dbgen.APIKey(t, db, database.APIKey{UserID: user.ID, ExpiresAt: exp})
+	}
+
+	unexpiredTimes := []time.Time{
+		expiredBefore.Add(time.Hour * 24 * 365),
+		expiredBefore.Add(time.Hour * 24),
+		expiredBefore.Add(time.Hour),
+		expiredBefore.Add(time.Minute),
+		expiredBefore.Add(time.Second),
+	}
+	for _, unexp := range unexpiredTimes {
+		// Unexpired api keys
+		dbgen.APIKey(t, db, database.APIKey{UserID: user.ID, ExpiresAt: unexp})
+	}
+
+	// All keys are present before deletion
+	keys, err := db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{
+		LoginType: user.LoginType,
+		UserID:    user.ID,
+	})
+	require.NoError(t, err)
+	require.Len(t, keys, len(expiredTimes)+len(unexpiredTimes))
+
+	// Delete expired keys
+	// First verify the limit works by deleting one at a time
+	deletedCount, err := db.DeleteExpiredAPIKeys(ctx, database.DeleteExpiredAPIKeysParams{
+		Before:     expiredBefore,
+		LimitCount: 1,
+	})
+	require.NoError(t, err)
+	require.Equal(t, int64(1), deletedCount)
+
+	// Ensure it was deleted
+	remaining, err := db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{
+		LoginType: user.LoginType,
+		UserID:    user.ID,
+	})
+	require.NoError(t, err)
+	require.Len(t, remaining, len(expiredTimes)+len(unexpiredTimes)-1)
+
+	// Delete the rest of the expired keys
+	deletedCount, err = db.DeleteExpiredAPIKeys(ctx, database.DeleteExpiredAPIKeysParams{
+		Before:     expiredBefore,
+		LimitCount: 100,
+	})
+	require.NoError(t, err)
+	require.Equal(t, int64(len(expiredTimes)-1), deletedCount)
+
+	// Ensure only unexpired keys remain
+	remaining, err = db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{
+		LoginType: user.LoginType,
+		UserID:    user.ID,
+	})
+	require.NoError(t, err)
+	require.Len(t, remaining, len(unexpiredTimes))
+}

coderd/database/queries.sql.go

@@ -1103,6 +1103,38 @@ func (q *sqlQuerier) DeleteApplicationConnectAPIKeysByUserID(ctx context.Context
 	return err
 }
 
+const deleteExpiredAPIKeys = `-- name: DeleteExpiredAPIKeys :one
+WITH expired_keys AS (
+	SELECT id
+	FROM api_keys
+	-- expired keys only
+	WHERE expires_at < $1::timestamptz
+	LIMIT $2
+),
+deleted_rows AS (
+	 DELETE FROM
+		 api_keys
+	 USING
+		 expired_keys
+	 WHERE
+		 api_keys.id = expired_keys.id
+	 RETURNING api_keys.id
+ )
+SELECT COUNT(deleted_rows.id) AS deleted_count FROM deleted_rows
+`
+
+type DeleteExpiredAPIKeysParams struct {
+	Before     time.Time `db:"before" json:"before"`
+	LimitCount int32     `db:"limit_count" json:"limit_count"`
+}
+
+func (q *sqlQuerier) DeleteExpiredAPIKeys(ctx context.Context, arg DeleteExpiredAPIKeysParams) (int64, error) {
+	row := q.db.QueryRowContext(ctx, deleteExpiredAPIKeys, arg.Before, arg.LimitCount)
+	var deleted_count int64
+	err := row.Scan(&deleted_count)
+	return deleted_count, err
+}
+
 const expirePrebuildsAPIKeys = `-- name: ExpirePrebuildsAPIKeys :exec
 WITH unexpired_prebuilds_workspace_session_tokens AS (
 	SELECT id, SUBSTRING(token_name FROM 38 FOR 36)::uuid AS workspace_id

coderd/database/queries/apikeys.sql

@@ -85,6 +85,26 @@ DELETE FROM
 WHERE
 	user_id = $1;
 
+-- name: DeleteExpiredAPIKeys :one
+WITH expired_keys AS (
+	SELECT id
+	FROM api_keys
+	-- expired keys only
+	WHERE expires_at < @before::timestamptz
+	LIMIT @limit_count
+),
+deleted_rows AS (
+	 DELETE FROM
+		 api_keys
+	 USING
+		 expired_keys
+	 WHERE
+		 api_keys.id = expired_keys.id
+	 RETURNING api_keys.id
+ )
+SELECT COUNT(deleted_rows.id) AS deleted_count FROM deleted_rows;
+;
+
 -- name: ExpirePrebuildsAPIKeys :exec
 -- Firstly, collect api_keys owned by the prebuilds user that correlate
 -- to workspaces no longer owned by the prebuilds user.