From ad93262d072dc470c3e3210f0efe42fbd1dc678e Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 3 Dec 2025 11:37:18 +0200 Subject: [PATCH] fix(coderd/database/dbpurge): allow disabling AI Bridge retention with 0 (#21062) Previously setting AI Bridge retention to 0 would cause records to be deleted immediately since we didn't check for the zero value before calculating the deletion threshold. This adds a check for aibridgeRetention > 0 to skip deletion when retention is disabled, matching the pattern used for other retention settings (connection logs, audit logs, etc.). Also fixes the return type of DeleteOldAIBridgeRecords from int32 to int64 since COUNT(*) returns bigint in PostgreSQL. Refs #21055 --- coderd/database/dbauthz/dbauthz.go | 2 +- coderd/database/dbauthz/dbauthz_test.go | 2 +- coderd/database/dbmetrics/querymetrics.go | 2 +- coderd/database/dbmock/dbmock.go | 4 +- coderd/database/dbpurge/dbpurge.go | 14 +- coderd/database/dbpurge/dbpurge_test.go | 365 +++++++++++++--------- coderd/database/querier.go | 2 +- coderd/database/queries.sql.go | 9 +- coderd/database/queries/aibridge.sql | 5 +- 9 files changed, 239 insertions(+), 166 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 94473c4899..ec56e244a0 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1732,7 +1732,7 @@ func (q *querier) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Contex return q.db.DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx, arg) } -func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { +func (q *querier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAibridgeInterception); err != nil { return -1, err } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 53a51edaab..0e1e684587 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4705,7 +4705,7 @@ func (s *MethodTestSuite) TestAIBridge() { s.Run("DeleteOldAIBridgeRecords", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { t := dbtime.Now() - db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int32(0), nil).AnyTimes() + db.EXPECT().DeleteOldAIBridgeRecords(gomock.Any(), t).Return(int64(0), nil).AnyTimes() check.Args(t).Asserts(rbac.ResourceAibridgeInterception, policy.ActionDelete) })) } diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index bffa55a94d..6a018f4190 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -396,7 +396,7 @@ func (m queryMetricsStore) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx conte return r0 } -func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { +func (m queryMetricsStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) { start := time.Now() r0, r1 := m.s.DeleteOldAIBridgeRecords(ctx, beforeTime) m.queryLatencies.WithLabelValues("DeleteOldAIBridgeRecords").Observe(time.Since(start).Seconds()) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 263ea59b79..f25e91e90c 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -725,10 +725,10 @@ func (mr *MockStoreMockRecorder) DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx } // DeleteOldAIBridgeRecords mocks base method. -func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { +func (m *MockStore) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "DeleteOldAIBridgeRecords", ctx, beforeTime) - ret0, _ := ret[0].(int32) + ret0, _ := ret[0].(int64) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/coderd/database/dbpurge/dbpurge.go b/coderd/database/dbpurge/dbpurge.go index 8e7357d8ba..8646fb6d02 100644 --- a/coderd/database/dbpurge/dbpurge.go +++ b/coderd/database/dbpurge/dbpurge.go @@ -118,11 +118,15 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder return xerrors.Errorf("failed to delete old audit log connection events: %w", err) } - deleteAIBridgeRecordsBefore := start.Add(-vals.AI.BridgeConfig.Retention.Value()) - // nolint:gocritic // Needs to run as aibridge context. - purgedAIBridgeRecords, err := tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore) - if err != nil { - return xerrors.Errorf("failed to delete old aibridge records: %w", err) + var purgedAIBridgeRecords int64 + aibridgeRetention := vals.AI.BridgeConfig.Retention.Value() + if aibridgeRetention > 0 { + deleteAIBridgeRecordsBefore := start.Add(-aibridgeRetention) + // nolint:gocritic // Needs to run as aibridge context. + purgedAIBridgeRecords, err = tx.DeleteOldAIBridgeRecords(dbauthz.AsAIBridged(ctx), deleteAIBridgeRecordsBefore) + if err != nil { + return xerrors.Errorf("failed to delete old aibridge records: %w", err) + } } var purgedConnectionLogs int64 diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index 5b8ab92a0d..05092dd3a3 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -973,170 +973,237 @@ func TestDeleteOldConnectionLogs(t *testing.T) { func TestDeleteOldAIBridgeRecords(t *testing.T) { t.Parallel() - ctx := testutil.Context(t, testutil.WaitShort) - - clk := quartz.NewMock(t) now := time.Date(2025, 1, 15, 7, 30, 0, 0, time.UTC) retentionPeriod := 30 * 24 * time.Hour // 30 days afterThreshold := now.Add(-retentionPeriod).Add(-24 * time.Hour) // 31 days ago (older than threshold) beforeThreshold := now.Add(-15 * 24 * time.Hour) // 15 days ago (newer than threshold) closeBeforeThreshold := now.Add(-retentionPeriod).Add(24 * time.Hour) // 29 days ago - clk.Set(now).MustWait(ctx) - db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure()) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - user := dbgen.User(t, db, database.User{}) - - // Create old AI Bridge interception (should be deleted) - oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ - ID: uuid.New(), - APIKeyID: sql.NullString{}, - InitiatorID: user.ID, - Provider: "anthropic", - Model: "claude-3-5-sonnet", - StartedAt: afterThreshold, - }, &afterThreshold) - - // Create old interception with related records (should all be deleted) - oldInterceptionWithRelated := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ - ID: uuid.New(), - APIKeyID: sql.NullString{}, - InitiatorID: user.ID, - Provider: "openai", - Model: "gpt-4", - StartedAt: afterThreshold, - }, &afterThreshold) - - _ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ - ID: uuid.New(), - InterceptionID: oldInterceptionWithRelated.ID, - ProviderResponseID: "resp-1", - InputTokens: 100, - OutputTokens: 50, - CreatedAt: afterThreshold, - }) - - _ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ - ID: uuid.New(), - InterceptionID: oldInterceptionWithRelated.ID, - ProviderResponseID: "resp-1", - Prompt: "test prompt", - CreatedAt: afterThreshold, - }) - - _ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ - ID: uuid.New(), - InterceptionID: oldInterceptionWithRelated.ID, - ProviderResponseID: "resp-1", - Tool: "test-tool", - ServerUrl: sql.NullString{String: "http://test", Valid: true}, - Input: "{}", - Injected: true, - CreatedAt: afterThreshold, - }) - - // Create recent AI Bridge interception (should be kept) - recentInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ - ID: uuid.New(), - APIKeyID: sql.NullString{}, - InitiatorID: user.ID, - Provider: "anthropic", - Model: "claude-3-5-sonnet", - StartedAt: beforeThreshold, - }, &beforeThreshold) - - // Create interception close to threshold (should be kept) - nearThresholdInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ - ID: uuid.New(), - APIKeyID: sql.NullString{}, - InitiatorID: user.ID, - Provider: "anthropic", - Model: "claude-3-5-sonnet", - StartedAt: closeBeforeThreshold, - }, &closeBeforeThreshold) - - _ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ - ID: uuid.New(), - InterceptionID: nearThresholdInterception.ID, - ProviderResponseID: "resp-1", - InputTokens: 100, - OutputTokens: 50, - CreatedAt: closeBeforeThreshold, - }) - - _ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ - ID: uuid.New(), - InterceptionID: nearThresholdInterception.ID, - ProviderResponseID: "resp-1", - Prompt: "test prompt", - CreatedAt: closeBeforeThreshold, - }) - - _ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ - ID: uuid.New(), - InterceptionID: nearThresholdInterception.ID, - ProviderResponseID: "resp-1", - Tool: "test-tool", - ServerUrl: sql.NullString{String: "http://test", Valid: true}, - Input: "{}", - Injected: true, - CreatedAt: closeBeforeThreshold, - }) - - // Run the purge with configured retention period - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{ - AI: codersdk.AIConfig{ - BridgeConfig: codersdk.AIBridgeConfig{ - Retention: serpent.Duration(retentionPeriod), - }, - }, - }, clk) - defer closer.Close() - // Wait for tick - testutil.TryReceive(ctx, t, done) - - // Verify results by querying all AI Bridge records - interceptions, err := db.GetAIBridgeInterceptions(ctx) - require.NoError(t, err) - - // Extract interception IDs for comparison - interceptionIDs := make([]uuid.UUID, len(interceptions)) - for i, interception := range interceptions { - interceptionIDs[i] = interception.ID + type testFixtures struct { + oldInterception database.AIBridgeInterception + oldInterceptionWithRelated database.AIBridgeInterception + recentInterception database.AIBridgeInterception + nearThresholdInterception database.AIBridgeInterception } - require.NotContains(t, interceptionIDs, oldInterception.ID, "old interception should be deleted") - require.NotContains(t, interceptionIDs, oldInterceptionWithRelated.ID, "old interception with related records should be deleted") + testCases := []struct { + name string + retention time.Duration + verify func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) + }{ + { + name: "RetentionEnabled", + retention: retentionPeriod, + verify: func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) { + t.Helper() - // Verify related records were also deleted - oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) - require.NoError(t, err) - require.Empty(t, oldTokenUsages, "old token usages should be deleted") + interceptions, err := db.GetAIBridgeInterceptions(ctx) + require.NoError(t, err) + require.Len(t, interceptions, 2, "expected 2 interceptions remaining") - oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, oldInterceptionWithRelated.ID) - require.NoError(t, err) - require.Empty(t, oldUserPrompts, "old user prompts should be deleted") + interceptionIDs := make([]uuid.UUID, len(interceptions)) + for i, interception := range interceptions { + interceptionIDs[i] = interception.ID + } - oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, oldInterceptionWithRelated.ID) - require.NoError(t, err) - require.Empty(t, oldToolUsages, "old tool usages should be deleted") + require.NotContains(t, interceptionIDs, fixtures.oldInterception.ID, "old interception should be deleted") + require.NotContains(t, interceptionIDs, fixtures.oldInterceptionWithRelated.ID, "old interception with related records should be deleted") + require.Contains(t, interceptionIDs, fixtures.recentInterception.ID, "recent interception should be kept") + require.Contains(t, interceptionIDs, fixtures.nearThresholdInterception.ID, "near threshold interception should be kept") - require.Contains(t, interceptionIDs, recentInterception.ID, "recent interception should be kept") - require.Contains(t, interceptionIDs, nearThresholdInterception.ID, "near threshold interception should be kept") + // Verify related records were deleted for old interception. + oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Empty(t, oldTokenUsages, "old token usages should be deleted") - // Verify related records were NOT deleted - newTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, nearThresholdInterception.ID) - require.NoError(t, err) - require.Len(t, newTokenUsages, 1, "near threshold token usages should not be deleted") + oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Empty(t, oldUserPrompts, "old user prompts should be deleted") - newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, nearThresholdInterception.ID) - require.NoError(t, err) - require.Len(t, newUserPrompts, 1, "near threshold user prompts should not be deleted") + oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Empty(t, oldToolUsages, "old tool usages should be deleted") - newToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, nearThresholdInterception.ID) - require.NoError(t, err) - require.Len(t, newToolUsages, 1, "near threshold tool usages should not be deleted") + // Verify related records were NOT deleted for near-threshold interception. + newTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.nearThresholdInterception.ID) + require.NoError(t, err) + require.Len(t, newTokenUsages, 1, "near threshold token usages should not be deleted") + + newUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.nearThresholdInterception.ID) + require.NoError(t, err) + require.Len(t, newUserPrompts, 1, "near threshold user prompts should not be deleted") + + newToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.nearThresholdInterception.ID) + require.NoError(t, err) + require.Len(t, newToolUsages, 1, "near threshold tool usages should not be deleted") + }, + }, + { + name: "RetentionDisabled", + retention: 0, + verify: func(t *testing.T, ctx context.Context, db database.Store, fixtures testFixtures) { + t.Helper() + + interceptions, err := db.GetAIBridgeInterceptions(ctx) + require.NoError(t, err) + require.Len(t, interceptions, 4, "expected all 4 interceptions to be retained") + + interceptionIDs := make([]uuid.UUID, len(interceptions)) + for i, interception := range interceptions { + interceptionIDs[i] = interception.ID + } + + require.Contains(t, interceptionIDs, fixtures.oldInterception.ID, "old interception should be kept") + require.Contains(t, interceptionIDs, fixtures.oldInterceptionWithRelated.ID, "old interception with related records should be kept") + require.Contains(t, interceptionIDs, fixtures.recentInterception.ID, "recent interception should be kept") + require.Contains(t, interceptionIDs, fixtures.nearThresholdInterception.ID, "near threshold interception should be kept") + + // Verify all related records were kept. + oldTokenUsages, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Len(t, oldTokenUsages, 1, "old token usages should be kept") + + oldUserPrompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Len(t, oldUserPrompts, 1, "old user prompts should be kept") + + oldToolUsages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, fixtures.oldInterceptionWithRelated.ID) + require.NoError(t, err) + require.Len(t, oldToolUsages, 1, "old tool usages should be kept") + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clk := quartz.NewMock(t) + clk.Set(now).MustWait(ctx) + + db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure()) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + user := dbgen.User(t, db, database.User{}) + + // Create old AI Bridge interception (should be deleted when retention enabled). + oldInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + APIKeyID: sql.NullString{}, + InitiatorID: user.ID, + Provider: "anthropic", + Model: "claude-3-5-sonnet", + StartedAt: afterThreshold, + }, &afterThreshold) + + // Create old interception with related records (should all be deleted when retention enabled). + oldInterceptionWithRelated := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + APIKeyID: sql.NullString{}, + InitiatorID: user.ID, + Provider: "openai", + Model: "gpt-4", + StartedAt: afterThreshold, + }, &afterThreshold) + + _ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + ID: uuid.New(), + InterceptionID: oldInterceptionWithRelated.ID, + ProviderResponseID: "resp-1", + InputTokens: 100, + OutputTokens: 50, + CreatedAt: afterThreshold, + }) + + _ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + ID: uuid.New(), + InterceptionID: oldInterceptionWithRelated.ID, + ProviderResponseID: "resp-1", + Prompt: "test prompt", + CreatedAt: afterThreshold, + }) + + _ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ + ID: uuid.New(), + InterceptionID: oldInterceptionWithRelated.ID, + ProviderResponseID: "resp-1", + Tool: "test-tool", + ServerUrl: sql.NullString{String: "http://test", Valid: true}, + Input: "{}", + Injected: true, + CreatedAt: afterThreshold, + }) + + // Create recent AI Bridge interception (should be kept). + recentInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + APIKeyID: sql.NullString{}, + InitiatorID: user.ID, + Provider: "anthropic", + Model: "claude-3-5-sonnet", + StartedAt: beforeThreshold, + }, &beforeThreshold) + + // Create interception close to threshold (should be kept). + nearThresholdInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + APIKeyID: sql.NullString{}, + InitiatorID: user.ID, + Provider: "anthropic", + Model: "claude-3-5-sonnet", + StartedAt: closeBeforeThreshold, + }, &closeBeforeThreshold) + + _ = dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + ID: uuid.New(), + InterceptionID: nearThresholdInterception.ID, + ProviderResponseID: "resp-1", + InputTokens: 100, + OutputTokens: 50, + CreatedAt: closeBeforeThreshold, + }) + + _ = dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + ID: uuid.New(), + InterceptionID: nearThresholdInterception.ID, + ProviderResponseID: "resp-1", + Prompt: "test prompt", + CreatedAt: closeBeforeThreshold, + }) + + _ = dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ + ID: uuid.New(), + InterceptionID: nearThresholdInterception.ID, + ProviderResponseID: "resp-1", + Tool: "test-tool", + ServerUrl: sql.NullString{String: "http://test", Valid: true}, + Input: "{}", + Injected: true, + CreatedAt: closeBeforeThreshold, + }) + + fixtures := testFixtures{ + oldInterception: oldInterception, + oldInterceptionWithRelated: oldInterceptionWithRelated, + recentInterception: recentInterception, + nearThresholdInterception: nearThresholdInterception, + } + + // Run the purge with configured retention period. + done := awaitDoTick(ctx, t, clk) + closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{ + AI: codersdk.AIConfig{ + BridgeConfig: codersdk.AIBridgeConfig{ + Retention: serpent.Duration(tc.retention), + }, + }, + }, clk) + defer closer.Close() + testutil.TryReceive(ctx, t, done) + + tc.verify(t, ctx, db, fixtures) + }) + } } func TestDeleteOldAuditLogs(t *testing.T) { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 22abe565a1..7202d22f3d 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -104,7 +104,7 @@ type sqlcQuerier interface { DeleteOAuth2ProviderAppSecretByID(ctx context.Context, id uuid.UUID) error DeleteOAuth2ProviderAppTokensByAppAndUserID(ctx context.Context, arg DeleteOAuth2ProviderAppTokensByAppAndUserIDParams) error // Cumulative count. - DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) + DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) DeleteOldAuditLogConnectionEvents(ctx context.Context, arg DeleteOldAuditLogConnectionEventsParams) error // Deletes old audit logs based on retention policy, excluding deprecated // connection events (connect, disconnect, open, close) which are handled diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 69184be52a..d1277ff945 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -354,17 +354,18 @@ WITH WHERE id IN (SELECT id FROM to_delete) RETURNING 1 ) -SELECT +SELECT ( (SELECT COUNT(*) FROM tool_usages) + (SELECT COUNT(*) FROM token_usages) + (SELECT COUNT(*) FROM user_prompts) + - (SELECT COUNT(*) FROM interceptions) as total_deleted + (SELECT COUNT(*) FROM interceptions) +)::bigint as total_deleted ` // Cumulative count. -func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int32, error) { +func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime time.Time) (int64, error) { row := q.db.QueryRowContext(ctx, deleteOldAIBridgeRecords, beforeTime) - var total_deleted int32 + var total_deleted int64 err := row.Scan(&total_deleted) return total_deleted, err } diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index cf87598115..960fe18ec0 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -360,8 +360,9 @@ WITH RETURNING 1 ) -- Cumulative count. -SELECT +SELECT ( (SELECT COUNT(*) FROM tool_usages) + (SELECT COUNT(*) FROM token_usages) + (SELECT COUNT(*) FROM user_prompts) + - (SELECT COUNT(*) FROM interceptions) as total_deleted; + (SELECT COUNT(*) FROM interceptions) +)::bigint as total_deleted;