fix(tables): probe for leftover rows on budget exhaustion to avoid needless deferral

Author: TheodoreSpeaks

apps/sim/lib/cleanup/batch-delete.test.ts

@@ -2,7 +2,7 @@
  * @vitest-environment node
  */
 
-import { dbChainMock, dbChainMockFns } from '@sim/testing'
+import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 vi.mock('@sim/db', () => dbChainMock)
@@ -24,6 +24,7 @@ function returnRows(count: number) {
 describe('drainRowsByColumn', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
   })
 
   it('drains in batches until a short batch and reports the set fully drained', async () => {
@@ -37,17 +38,31 @@ describe('drainRowsByColumn', () => {
     expect(dbChainMockFns.returning).toHaveBeenCalledTimes(2)
   })
 
-  it('stops at the row budget and reports the set not fully drained', async () => {
+  it('stops at the budget and reports not fully drained when rows remain', async () => {
     dbChainMockFns.returning
       .mockResolvedValueOnce(returnRows(2))
       .mockResolvedValueOnce(returnRows(2))
+    // Existence probe after the budget is spent finds a leftover row.
+    dbChainMockFns.limit.mockResolvedValue([{ id: 'leftover' }])
 
     const result = await drainRowsByColumn({ ...baseOpts, batchSize: 2, rowBudget: 4 })
 
     expect(result).toEqual({ deleted: 4, fullyDrained: false })
     expect(dbChainMockFns.returning).toHaveBeenCalledTimes(2)
   })
 
+  it('reports fully drained when the budget is hit but the set emptied exactly', async () => {
+    dbChainMockFns.returning
+      .mockResolvedValueOnce(returnRows(2))
+      .mockResolvedValueOnce(returnRows(2))
+    // Existence probe finds nothing remaining.
+    dbChainMockFns.limit.mockResolvedValue([])
+
+    const result = await drainRowsByColumn({ ...baseOpts, batchSize: 2, rowBudget: 4 })
+
+    expect(result).toEqual({ deleted: 4, fullyDrained: true })
+  })
+
   it('reports fully drained immediately when the match set is already empty', async () => {
     dbChainMockFns.returning.mockResolvedValueOnce([])
 

apps/sim/lib/cleanup/batch-delete.ts

@@ -279,10 +279,11 @@ export interface DrainByColumnOptions {
 export interface DrainResult {
   deleted: number
   /**
-   * True only when the match set was confirmed empty (a short final batch).
-   * Budget exhaustion and batch errors both yield `false` — callers must treat
-   * `false` as "rows may remain" and defer any dependent parent-delete (whose
-   * ON DELETE CASCADE would otherwise fire on the leftovers) to a later run.
+   * True only when the match set was confirmed empty — via a short final batch
+   * or an existence probe after the budget was spent. Batch errors yield
+   * `false`; callers must treat `false` as "rows may remain" and defer any
+   * dependent parent-delete (whose ON DELETE CASCADE would otherwise fire on the
+   * leftovers) to a later run.
    */
   fullyDrained: boolean
 }
@@ -331,6 +332,18 @@ export async function drainRowsByColumn({
     if (batchDeleted.length < limit) return { deleted, fullyDrained: true }
   }
 
-  // Hit the per-call budget on a full batch — rows may remain.
-  return { deleted, fullyDrained: false }
+  // Budget hit on a full final batch — rows may or may not remain. A cheap
+  // indexed existence probe disambiguates so a set whose size divides the budget
+  // exactly isn't needlessly deferred to a later run.
+  try {
+    const [leftover] = await db
+      .select({ id: idCol })
+      .from(tableDef)
+      .where(eq(matchCol, matchValue))
+      .limit(1)
+    return { deleted, fullyDrained: !leftover }
+  } catch (error) {
+    logger.error(`[${tableName}] Drain remainder probe failed for ${matchValue}:`, { error })
+    return { deleted, fullyDrained: false }
+  }
 }