Skip to content

Commit cefe07d

Browse files
EmyrkEmyrk
authored
feat: purge expired api keys in dbpurge (#20863)
closes #19889 This is in response to a migration in v2.27 that takes very long on deployments with large `api_key` tables.
1 parent c6631e1 commit cefe07d

10 files changed

Lines changed: 197 additions & 4 deletions

File tree

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1641,6 +1641,15 @@ func (q *querier) DeleteCustomRole(ctx context.Context, arg database.DeleteCusto
16411641
return q.db.DeleteCustomRole(ctx, arg)
16421642
}
16431643

1644+
func (q *querier) DeleteExpiredAPIKeys(ctx context.Context, arg database.DeleteExpiredAPIKeysParams) (int64, error) {
1645+
// Requires DELETE across all API keys.
1646+
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceApiKey); err != nil {
1647+
return 0, err
1648+
}
1649+
1650+
return q.db.DeleteExpiredAPIKeys(ctx, arg)
1651+
}
1652+
16441653
func (q *querier) DeleteExternalAuthLink(ctx context.Context, arg database.DeleteExternalAuthLinkParams) error {
16451654
return fetchAndExec(q.log, q.auth, policy.ActionUpdatePersonal, func(ctx context.Context, arg database.DeleteExternalAuthLinkParams) (database.ExternalAuthLink, error) {
16461655
//nolint:gosimple

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,14 @@ func (s *MethodTestSuite) TestAPIKey() {
216216
dbm.EXPECT().DeleteAPIKeyByID(gomock.Any(), key.ID).Return(nil).AnyTimes()
217217
check.Args(key.ID).Asserts(key, policy.ActionDelete).Returns()
218218
}))
219+
s.Run("DeleteExpiredAPIKeys", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
220+
args := database.DeleteExpiredAPIKeysParams{
221+
Before: time.Date(2025, 11, 21, 0, 0, 0, 0, time.UTC),
222+
LimitCount: 1000,
223+
}
224+
dbm.EXPECT().DeleteExpiredAPIKeys(gomock.Any(), args).Return(int64(0), nil).AnyTimes()
225+
check.Args(args).Asserts(rbac.ResourceApiKey, policy.ActionDelete).Returns(int64(0))
226+
}))
219227
s.Run("GetAPIKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
220228
key := testutil.Fake(s.T(), faker, database.APIKey{})
221229
dbm.EXPECT().GetAPIKeyByID(gomock.Any(), key.ID).Return(key, nil).AnyTimes()

‎coderd/database/dbgen/dbgen.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,13 @@ func APIKey(t testing.TB, db database.Store, seed database.APIKey, munge ...func
175175
}
176176
}
177177

178+
// It does not make sense for the created_at to be after the expires_at.
179+
// So if expires is set, change the default created_at to be 24 hours before.
180+
var createdAt time.Time
181+
if !seed.ExpiresAt.IsZero() && seed.CreatedAt.IsZero() {
182+
createdAt = seed.ExpiresAt.Add(-24 * time.Hour)
183+
}
184+
178185
params := database.InsertAPIKeyParams{
179186
ID: takeFirst(seed.ID, id),
180187
// 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
184191
UserID: takeFirst(seed.UserID, uuid.New()),
185192
LastUsed: takeFirst(seed.LastUsed, dbtime.Now()),
186193
ExpiresAt: takeFirst(seed.ExpiresAt, dbtime.Now().Add(time.Hour)),
187-
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
194+
CreatedAt: takeFirst(seed.CreatedAt, createdAt, dbtime.Now()),
188195
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
189196
LoginType: takeFirst(seed.LoginType, database.LoginTypePassword),
190197
Scopes: takeFirstSlice([]database.APIKeyScope(seed.Scopes), []database.APIKeyScope{database.ApiKeyScopeCoderAll}),

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbpurge/dbpurge.go‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,19 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
7878
if err := tx.ExpirePrebuildsAPIKeys(ctx, dbtime.Time(start)); err != nil {
7979
return xerrors.Errorf("failed to expire prebuilds user api keys: %w", err)
8080
}
81+
expiredAPIKeys, err := tx.DeleteExpiredAPIKeys(ctx, database.DeleteExpiredAPIKeysParams{
82+
// Leave expired keys for a week to allow the backend to know the difference
83+
// between a 404 and an expired key. This purge code is just to bound the size of
84+
// the table to something more reasonable.
85+
Before: dbtime.Time(start.Add(time.Hour * 24 * 7 * -1)),
86+
// There could be a lot of expired keys here, so set a limit to prevent this
87+
// taking too long.
88+
// This runs every 10 minutes, so it deletes ~1.5m keys per day at most.
89+
LimitCount: 10000,
90+
})
91+
if err != nil {
92+
return xerrors.Errorf("failed to delete expired api keys: %w", err)
93+
}
8194
deleteOldTelemetryLocksBefore := start.Add(-maxTelemetryHeartbeatAge)
8295
if err := tx.DeleteOldTelemetryLocks(ctx, deleteOldTelemetryLocksBefore); err != nil {
8396
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
93106

94107
deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value())
95108
// nolint:gocritic // Needs to run as aibridge context.
96-
count, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
109+
purgedAIBridgeRecords, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore)
97110
if err != nil {
98111
return xerrors.Errorf("failed to delete old aibridge records: %w", err)
99112
}
100-
logger.Debug(ctx, "purged aibridge entries", slog.F("count", count), slog.F("since", deleteAIBridgeRecordsBefore.Format(time.RFC3339)))
101113

102-
logger.Debug(ctx, "purged old database entries", slog.F("duration", clk.Since(start)))
114+
logger.Debug(ctx, "purged old database entries",
115+
slog.F("expired_api_keys", expiredAPIKeys),
116+
slog.F("aibridge_records", purgedAIBridgeRecords),
117+
slog.F("duration", clk.Since(start)),
118+
)
103119

104120
return nil
105121
}, database.DefaultTXOptions().WithID("db_purge")); err != nil {

‎coderd/database/querier.go‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier_test.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7835,3 +7835,81 @@ func TestUpdateAIBridgeInterceptionEnded(t *testing.T) {
78357835
}
78367836
})
78377837
}
7838+
7839+
func TestDeleteExpiredAPIKeys(t *testing.T) {
7840+
t.Parallel()
7841+
db, _ := dbtestutil.NewDB(t)
7842+
7843+
// Constant time for testing
7844+
now := time.Date(2025, 11, 20, 12, 0, 0, 0, time.UTC)
7845+
expiredBefore := now.Add(-time.Hour) // Anything before this is expired
7846+
7847+
ctx := testutil.Context(t, testutil.WaitLong)
7848+
7849+
user := dbgen.User(t, db, database.User{})
7850+
7851+
expiredTimes := []time.Time{
7852+
expiredBefore.Add(-time.Hour * 24 * 365),
7853+
expiredBefore.Add(-time.Hour * 24),
7854+
expiredBefore.Add(-time.Hour),
7855+
expiredBefore.Add(-time.Minute),
7856+
expiredBefore.Add(-time.Second),
7857+
}
7858+
for _, exp := range expiredTimes {
7859+
// Expired api keys
7860+
dbgen.APIKey(t, db, database.APIKey{UserID: user.ID, ExpiresAt: exp})
7861+
}
7862+
7863+
unexpiredTimes := []time.Time{
7864+
expiredBefore.Add(time.Hour * 24 * 365),
7865+
expiredBefore.Add(time.Hour * 24),
7866+
expiredBefore.Add(time.Hour),
7867+
expiredBefore.Add(time.Minute),
7868+
expiredBefore.Add(time.Second),
7869+
}
7870+
for _, unexp := range unexpiredTimes {
7871+
// Unexpired api keys
7872+
dbgen.APIKey(t, db, database.APIKey{UserID: user.ID, ExpiresAt: unexp})
7873+
}
7874+
7875+
// All keys are present before deletion
7876+
keys, err := db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{
7877+
LoginType: user.LoginType,
7878+
UserID: user.ID,
7879+
})
7880+
require.NoError(t, err)
7881+
require.Len(t, keys, len(expiredTimes)+len(unexpiredTimes))
7882+
7883+
// Delete expired keys
7884+
// First verify the limit works by deleting one at a time
7885+
deletedCount, err := db.DeleteExpiredAPIKeys(ctx, database.DeleteExpiredAPIKeysParams{
7886+
Before: expiredBefore,
7887+
LimitCount: 1,
7888+
})
7889+
require.NoError(t, err)
7890+
require.Equal(t, int64(1), deletedCount)
7891+
7892+
// Ensure it was deleted
7893+
remaining, err := db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{
7894+
LoginType: user.LoginType,
7895+
UserID: user.ID,
7896+
})
7897+
require.NoError(t, err)
7898+
require.Len(t, remaining, len(expiredTimes)+len(unexpiredTimes)-1)
7899+
7900+
// Delete the rest of the expired keys
7901+
deletedCount, err = db.DeleteExpiredAPIKeys(ctx, database.DeleteExpiredAPIKeysParams{
7902+
Before: expiredBefore,
7903+
LimitCount: 100,
7904+
})
7905+
require.NoError(t, err)
7906+
require.Equal(t, int64(len(expiredTimes)-1), deletedCount)
7907+
7908+
// Ensure only unexpired keys remain
7909+
remaining, err = db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{
7910+
LoginType: user.LoginType,
7911+
UserID: user.ID,
7912+
})
7913+
require.NoError(t, err)
7914+
require.Len(t, remaining, len(unexpiredTimes))
7915+
}

‎coderd/database/queries.sql.go‎

Lines changed: 32 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/apikeys.sql‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,26 @@ DELETE FROM
8585
WHERE
8686
user_id = $1;
8787

88+
-- name: DeleteExpiredAPIKeys :one
89+
WITH expired_keys AS (
90+
SELECT id
91+
FROM api_keys
92+
-- expired keys only
93+
WHERE expires_at < @before::timestamptz
94+
LIMIT @limit_count
95+
),
96+
deleted_rows AS (
97+
DELETE FROM
98+
api_keys
99+
USING
100+
expired_keys
101+
WHERE
102+
api_keys.id = expired_keys.id
103+
RETURNING api_keys.id
104+
)
105+
SELECT COUNT(deleted_rows.id) AS deleted_count FROM deleted_rows;
106+
;
107+
88108
-- name: ExpirePrebuildsAPIKeys :exec
89109
-- Firstly, collect api_keys owned by the prebuilds user that correlate
90110
-- to workspaces no longer owned by the prebuilds user.

0 commit comments

Comments
 (0)