fix(knowledge): take the image restore locks in sync's order

Author: waleedlatif1

apps/sim/lib/knowledge/application/connectors.ts

@@ -11,7 +11,7 @@ import {
 } from '@sim/db/schema'
 import { toError } from '@sim/utils/errors'
 import { escapeLikePattern } from '@sim/utils/string'
-import { and, asc, desc, eq, inArray, isNull, lt, or, type SQL, sql } from 'drizzle-orm'
+import { and, asc, desc, eq, inArray, isNull, lt, or, sql } from 'drizzle-orm'
 import type { ConnectorDocumentFilter } from '@/lib/api/contracts/knowledge/connectors'
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import { requireCurrentHumanRole } from '@/lib/core/application'
@@ -1278,34 +1278,31 @@ export const listKnowledgeConnectorDocuments = defineAuthorizedKnowledgeUseCase(
 /**
  * An excluded image keeps its indexed text, and turning image attachments off
  * skips excluded rows so the exclusion survives turning them back on. Restoring
- * one while images are off would re-enable text the source no longer lists. The
- * connector row stays locked until the restore commits, so a concurrent change of
- * the setting either sees the restored image or is seen by this check.
+ * one while images are off would re-enable text the source no longer lists.
+ *
+ * Runs after the restore's document update in the same transaction, matching a
+ * sync's lock order (documents, then a share lock on the connector). The share
+ * lock holds off a concurrent change of the setting until the restore commits,
+ * so that change queues removal of the restored image or this check sees it.
  */
 async function assertImageRestoreAllowed(
   tx: DbTransaction,
   connector: { id: string; connectorType: string },
-  restorable: SQL | undefined
+  restored: { mimeType: string }[]
 ) {
   const imageExtraction = getConnectorMeta(connector.connectorType)?.imageExtraction
-  if (!imageExtraction) return
+  if (!imageExtraction || !restored.some((row) => SEMANTIC_IMAGE_MIME_TYPES.has(row.mimeType))) {
+    return
+  }
   const [stored] = await tx
     .select({ sourceConfig: knowledgeConnector.sourceConfig })
     .from(knowledgeConnector)
     .where(eq(knowledgeConnector.id, connector.id))
-    .for('update')
+    .for('share')
   if (
-    !stored ||
-    isImageExtractionEnabled(imageExtraction, stored.sourceConfig as Record<string, unknown>)
+    stored &&
+    !isImageExtractionEnabled(imageExtraction, stored.sourceConfig as Record<string, unknown>)
   ) {
-    return
-  }
-  const [image] = await tx
-    .select({ id: document.id })
-    .from(document)
-    .where(and(restorable, inArray(document.mimeType, [...SEMANTIC_IMAGE_MIME_TYPES])))
-    .limit(1)
-  if (image) {
     throw new OrchestrationError(
       'validation',
       'Turn on image attachments for this connector before restoring an image.'
@@ -1337,21 +1334,24 @@ export const updateKnowledgeConnectorDocuments = defineAuthorizedKnowledgeUseCas
     if (restoring && !requiresConnectorIndexing(context.knowledgeBase.isSearchIndex)) {
       throw new OrchestrationError('validation', 'This search index is inactive; use Sim Search.')
     }
-    const target = and(
-      eq(document.connectorId, context.connectorId),
-      inArray(document.id, documentIds),
-      eq(document.userExcluded, restoring),
-      isNull(document.archivedAt),
-      isNull(document.deletedAt),
-      knowledgeAccessCondition(await context.access.get())
-    )
+    const access = await context.access.get()
     const updated = await db.transaction(async (tx) => {
-      if (restoring) await assertImageRestoreAllowed(tx, context.connector, target)
-      return tx
+      const rows = await tx
         .update(document)
         .set({ userExcluded: !restoring, enabled: restoring })
-        .where(target)
-        .returning({ id: document.id })
+        .where(
+          and(
+            eq(document.connectorId, context.connectorId),
+            inArray(document.id, documentIds),
+            eq(document.userExcluded, restoring),
+            isNull(document.archivedAt),
+            isNull(document.deletedAt),
+            knowledgeAccessCondition(access)
+          )
+        )
+        .returning({ id: document.id, mimeType: document.mimeType })
+      if (restoring) await assertImageRestoreAllowed(tx, context.connector, rows)
+      return rows
     })
     return {
       operation: input.operation,