use one select for cleanup and abandened resource cleanup

Author: TheodoreSpeaks

apps/sim/background/cleanup-soft-deletes.ts

@@ -15,31 +15,41 @@ import {
 import { createLogger } from '@sim/logger'
 import { task } from '@trigger.dev/sdk'
 import { and, inArray, isNotNull, lt } from 'drizzle-orm'
+import type { PgColumn, PgTable } from 'drizzle-orm/pg-core'
 import { type CleanupJobPayload, resolveCleanupScope } from '@/lib/billing/cleanup-dispatcher'
 import {
   batchDeleteByWorkspaceAndTimestamp,
   DEFAULT_BATCH_SIZE,
   DEFAULT_MAX_BATCHES_PER_TABLE,
+  type TableCleanupResult,
 } from '@/lib/cleanup/batch-delete'
 import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup'
 import type { StorageContext } from '@/lib/uploads'
 import { isUsingCloudStorage, StorageService } from '@/lib/uploads'
 
 const logger = createLogger('CleanupSoftDeletes')
 
-async function cleanupWorkspaceFileStorage(
+interface WorkspaceFileScope {
+  /** Rows from `workspace_file` (singular, legacy workspace-context only). */
+  legacyRows: Array<{ id: string; key: string }>
+  /** Rows from `workspace_files` (plural, multi-context). */
+  multiContextRows: Array<{ id: string; key: string; context: StorageContext }>
+}
+
+/**
+ * Select every soft-deleted file row that's eligible for permanent removal.
+ * Returned once and reused for both S3 deletion and DB deletion so the external
+ * cleanup cannot drift from the row-level cleanup.
+ */
+async function selectExpiredWorkspaceFiles(
   workspaceIds: string[],
   retentionDate: Date
-): Promise<{ filesDeleted: number; filesFailed: number }> {
-  const stats = { filesDeleted: 0, filesFailed: 0 }
-
-  if (!isUsingCloudStorage() || workspaceIds.length === 0) return stats
-
+): Promise<WorkspaceFileScope> {
   const limit = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE
 
-  const [legacyFiles, multiContextFiles] = await Promise.all([
+  const [legacyRows, multiContextRows] = await Promise.all([
     db
-      .select({ key: workspaceFile.key })
+      .select({ id: workspaceFile.id, key: workspaceFile.key })
       .from(workspaceFile)
       .where(
         and(
@@ -50,7 +60,11 @@ async function cleanupWorkspaceFileStorage(
       )
       .limit(limit),
     db
-      .select({ key: workspaceFiles.key, context: workspaceFiles.context })
+      .select({
+        id: workspaceFiles.id,
+        key: workspaceFiles.key,
+        context: workspaceFiles.context,
+      })
       .from(workspaceFiles)
       .where(
         and(
@@ -62,9 +76,25 @@ async function cleanupWorkspaceFileStorage(
       .limit(limit),
   ])
 
+  return {
+    legacyRows,
+    multiContextRows: multiContextRows.map((r) => ({
+      id: r.id,
+      key: r.key,
+      context: r.context as StorageContext,
+    })),
+  }
+}
+
+async function cleanupWorkspaceFileStorage(
+  scope: WorkspaceFileScope
+): Promise<{ filesDeleted: number; filesFailed: number }> {
+  const stats = { filesDeleted: 0, filesFailed: 0 }
+  if (!isUsingCloudStorage()) return stats
+
   const toDelete: Array<{ key: string; context: StorageContext }> = [
-    ...legacyFiles.map((f) => ({ key: f.key, context: 'workspace' as StorageContext })),
-    ...multiContextFiles.map((f) => ({ key: f.key, context: f.context as StorageContext })),
+    ...scope.legacyRows.map((r) => ({ key: r.key, context: 'workspace' as StorageContext })),
+    ...scope.multiContextRows.map((r) => ({ key: r.key, context: r.context })),
   ]
 
   await Promise.all(
@@ -82,13 +112,32 @@ async function cleanupWorkspaceFileStorage(
   return stats
 }
 
+async function deleteRowsById(
+  tableDef: PgTable,
+  idCol: PgColumn,
+  ids: string[],
+  tableName: string
+): Promise<TableCleanupResult> {
+  const result: TableCleanupResult = { table: tableName, deleted: 0, failed: 0 }
+  if (ids.length === 0) return result
+  try {
+    const deleted = await db.delete(tableDef).where(inArray(idCol, ids)).returning({ id: idCol })
+    result.deleted = deleted.length
+    logger.info(`[${tableName}] Deleted ${deleted.length} rows`)
+  } catch (error) {
+    result.failed++
+    logger.error(`[${tableName}] Delete failed:`, { error })
+  }
+  return result
+}
+
+/**
+ * Tables cleaned by the generic workspace-scoped batched DELETE. Tables whose
+ * hard-delete triggers external side effects (workflow → copilot chats cascade,
+ * workspace files → S3 storage) are handled explicitly so the SELECT that drives
+ * the external cleanup and the SELECT that drives the DB delete see the same rows.
+ */
 const CLEANUP_TARGETS = [
-  {
-    table: workflow,
-    softDeleteCol: workflow.archivedAt,
-    wsCol: workflow.workspaceId,
-    name: 'workflow',
-  },
   {
     table: workflowFolder,
     softDeleteCol: workflowFolder.archivedAt,
@@ -107,18 +156,6 @@ const CLEANUP_TARGETS = [
     wsCol: userTableDefinitions.workspaceId,
     name: 'userTableDefinitions',
   },
-  {
-    table: workspaceFile,
-    softDeleteCol: workspaceFile.deletedAt,
-    wsCol: workspaceFile.workspaceId,
-    name: 'workspaceFile',
-  },
-  {
-    table: workspaceFiles,
-    softDeleteCol: workspaceFiles.deletedAt,
-    wsCol: workspaceFiles.workspaceId,
-    name: 'workspaceFiles',
-  },
   { table: memory, softDeleteCol: memory.deletedAt, wsCol: memory.workspaceId, name: 'memory' },
   {
     table: mcpServers,
@@ -161,18 +198,23 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
     `[${label}] Processing ${workspaceIds.length} workspaces, cutoff: ${retentionDate.toISOString()}`
   )
 
-  // Find chats linked to workflows that are about to be cascade-deleted
-  const doomedWorkflows = await db
-    .select({ id: workflow.id })
-    .from(workflow)
-    .where(
-      and(
-        inArray(workflow.workspaceId, workspaceIds),
-        isNotNull(workflow.archivedAt),
-        lt(workflow.archivedAt, retentionDate)
+  // Select workflows + files once. These sets drive BOTH external cleanup
+  // (chats + S3) AND the DB deletes below — selecting twice could return
+  // different subsets above the LIMIT cap and orphan or prematurely purge data.
+  const [doomedWorkflows, fileScope] = await Promise.all([
+    db
+      .select({ id: workflow.id })
+      .from(workflow)
+      .where(
+        and(
+          inArray(workflow.workspaceId, workspaceIds),
+          isNotNull(workflow.archivedAt),
+          lt(workflow.archivedAt, retentionDate)
+        )
       )
-    )
-    .limit(DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE)
+      .limit(DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE),
+    selectExpiredWorkspaceFiles(workspaceIds, retentionDate),
+  ])
 
   const doomedWorkflowIds = doomedWorkflows.map((w) => w.id)
   let chatCleanup: { execute: () => Promise<void> } | null = null
@@ -190,9 +232,35 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
     }
   }
 
-  const fileStats = await cleanupWorkspaceFileStorage(workspaceIds, retentionDate)
+  const fileStats = await cleanupWorkspaceFileStorage(fileScope)
 
   let totalDeleted = 0
+
+  // Delete the workflow + file rows using the exact IDs we already selected.
+  const workflowResult = await deleteRowsById(
+    workflow,
+    workflow.id,
+    doomedWorkflowIds,
+    `${label}/workflow`
+  )
+  totalDeleted += workflowResult.deleted
+
+  const legacyFileResult = await deleteRowsById(
+    workspaceFile,
+    workspaceFile.id,
+    fileScope.legacyRows.map((r) => r.id),
+    `${label}/workspaceFile`
+  )
+  totalDeleted += legacyFileResult.deleted
+
+  const multiContextFileResult = await deleteRowsById(
+    workspaceFiles,
+    workspaceFiles.id,
+    fileScope.multiContextRows.map((r) => r.id),
+    `${label}/workspaceFiles`
+  )
+  totalDeleted += multiContextFileResult.deleted
+
   for (const target of CLEANUP_TARGETS) {
     const result = await batchDeleteByWorkspaceAndTimestamp({
       tableDef: target.table,