From 6baec938601da61cb709e4ec61af2fcf058ba71f Mon Sep 17 00:00:00 2001
From: Ong Chung Yau <33013947+chungyau97@users.noreply.github.com>
Date: Wed, 9 Jul 2025 18:36:47 +0800
Subject: [PATCH] Optimize export import (#4795)

* feat: add saveBatch and optimize duplication id handling

* feat: improve lookup performance by using Set
---
 .../src/controllers/export-import/index.ts    |  2 +-
 .../src/services/export-import/index.ts       | 41 +++++++++++++------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/packages/server/src/controllers/export-import/index.ts b/packages/server/src/controllers/export-import/index.ts
index 91c6d485..b066df0c 100644
--- a/packages/server/src/controllers/export-import/index.ts
+++ b/packages/server/src/controllers/export-import/index.ts
@@ -39,7 +39,7 @@ const importData = async (req: Request, res: Response, next: NextFunction) => {
         }
 
         await exportImportService.importData(importData, orgId, workspaceId, subscriptionId)
-        return res.json({ message: 'success' })
+        return res.status(StatusCodes.OK).json({ message: 'success' })
     } catch (error) {
         next(error)
     }
diff --git a/packages/server/src/services/export-import/index.ts b/packages/server/src/services/export-import/index.ts
index 87ddc6ad..696981a0 100644
--- a/packages/server/src/services/export-import/index.ts
+++ b/packages/server/src/services/export-import/index.ts
@@ -1,5 +1,5 @@
 import { StatusCodes } from 'http-status-codes'
-import { In, QueryRunner } from 'typeorm'
+import { EntityManager, In, QueryRunner } from 'typeorm'
 import { v4 as uuidv4 } from 'uuid'
 import { Assistant } from '../../database/entities/Assistant'
 import { ChatFlow } from '../../database/entities/ChatFlow'
@@ -255,11 +255,15 @@ async function replaceDuplicateIdsForChatMessage(
             where: { id: In(ids) }
         })
         if (records.length < 0) return originalData
-        for (let record of records) {
-            const oldId = record.id
-            const newId = uuidv4()
-            originalData = JSON.parse(JSON.stringify(originalData).replaceAll(oldId, newId))
-        }
+
+        // replace duplicate ids found in db to new id
+        const dbExistingIds = new Set(records.map((record) => record.id))
+        originalData.ChatMessage = originalData.ChatMessage.map((item) => {
+            if (dbExistingIds.has(item.id)) {
+                return { ...item, id: uuidv4() }
+            }
+            return item
+        })
         return originalData
     } catch (error) {
         throw new InternalFlowiseError(
@@ -459,11 +463,15 @@ async function replaceDuplicateIdsForDocumentStoreFileChunk(
             where: { id: In(ids) }
         })
         if (records.length < 0) return originalData
-        for (let record of records) {
-            const oldId = record.id
-            const newId = uuidv4()
-            originalData = JSON.parse(JSON.stringify(originalData).replaceAll(oldId, newId))
-        }
+
+        // replace duplicate ids found in db to new id
+        const dbExistingIds = new Set(records.map((record) => record.id))
+        originalData.DocumentStoreFileChunk = originalData.DocumentStoreFileChunk.map((item) => {
+            if (dbExistingIds.has(item.id)) {
+                return { ...item, id: uuidv4() }
+            }
+            return item
+        })
         return originalData
     } catch (error) {
         throw new InternalFlowiseError(
@@ -550,6 +558,13 @@ function insertWorkspaceId(importedData: any, activeWorkspaceId?: string) {
     return importedData
 }
 
+async function saveBatch(manager: EntityManager, entity: any, items: any[], batchSize = 900) {
+    for (let i = 0; i < items.length; i += batchSize) {
+        const batch = items.slice(i, i + batchSize)
+        await manager.save(entity, batch)
+    }
+}
+
 const importData = async (importData: ExportData, orgId: string, activeWorkspaceId: string, subscriptionId: string) => {
     // Initialize missing properties with empty arrays to avoid "undefined" errors
     importData.AgentFlow = importData.AgentFlow || []
@@ -705,13 +720,13 @@ const importData = async (importData: ExportData, orgId: string, activeWorkspace
             if (importData.AssistantOpenAI.length > 0) await queryRunner.manager.save(Assistant, importData.AssistantOpenAI)
             if (importData.AssistantAzure.length > 0) await queryRunner.manager.save(Assistant, importData.AssistantAzure)
             if (importData.ChatFlow.length > 0) await queryRunner.manager.save(ChatFlow, importData.ChatFlow)
-            if (importData.ChatMessage.length > 0) await queryRunner.manager.save(ChatMessage, importData.ChatMessage)
+            if (importData.ChatMessage.length > 0) await saveBatch(queryRunner.manager, ChatMessage, importData.ChatMessage)
             if (importData.ChatMessageFeedback.length > 0)
                 await queryRunner.manager.save(ChatMessageFeedback, importData.ChatMessageFeedback)
             if (importData.CustomTemplate.length > 0) await queryRunner.manager.save(CustomTemplate, importData.CustomTemplate)
             if (importData.DocumentStore.length > 0) await queryRunner.manager.save(DocumentStore, importData.DocumentStore)
             if (importData.DocumentStoreFileChunk.length > 0)
-                await queryRunner.manager.save(DocumentStoreFileChunk, importData.DocumentStoreFileChunk)
+                await saveBatch(queryRunner.manager, DocumentStoreFileChunk, importData.DocumentStoreFileChunk)
             if (importData.Tool.length > 0) await queryRunner.manager.save(Tool, importData.Tool)
             if (importData.Execution.length > 0) await queryRunner.manager.save(Execution, importData.Execution)
             if (importData.Variable.length > 0) await queryRunner.manager.save(Variable, importData.Variable)