Bugfix/Query Runner already released (#3525)

Merge branch 'feature/add-couchbase-vectore-store' into feature/Couchbase
FlowiseAI · Nov 16, 2024 · 47e723b · 47e723b
1 parent 1ccd976
commit 47e723b
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 37 deletions.
diff --git a/packages/server/src/services/assistants/index.ts b/packages/server/src/services/assistants/index.ts
@@ -7,7 +7,7 @@ import { Credential } from '../../database/entities/Credential'
 import { decryptCredentialData, getAppVersion } from '../../utils'
 import { InternalFlowiseError } from '../../errors/internalFlowiseError'
 import { getErrorMessage } from '../../errors/utils'
-import { DeleteResult } from 'typeorm'
+import { DeleteResult, QueryRunner } from 'typeorm'
 import { FLOWISE_METRIC_COUNTERS, FLOWISE_COUNTER_STATUS } from '../../Interface.Metrics'
 
 const createAssistant = async (requestBody: any): Promise<Assistant> => {
@@ -291,9 +291,10 @@ const updateAssistant = async (assistantId: string, requestBody: any): Promise<A
     }
 }
 
-const importAssistants = async (newAssistants: Partial<Assistant>[]): Promise<any> => {
+const importAssistants = async (newAssistants: Partial<Assistant>[], queryRunner?: QueryRunner): Promise<any> => {
     try {
         const appServer = getRunningExpressApp()
+        const repository = queryRunner ? queryRunner.manager.getRepository(Assistant) : appServer.AppDataSource.getRepository(Assistant)
 
         // step 1 - check whether array is zero
         if (newAssistants.length == 0) return
@@ -309,7 +310,7 @@ const importAssistants = async (newAssistants: Partial<Assistant>[]): Promise<an
             count += 1
         })
 
-        const selectResponse = await appServer.AppDataSource.getRepository(Assistant)
+        const selectResponse = await repository
             .createQueryBuilder('assistant')
             .select('assistant.id')
             .where(`assistant.id IN ${ids}`)
@@ -329,7 +330,7 @@ const importAssistants = async (newAssistants: Partial<Assistant>[]): Promise<an
         })
 
         // step 4 - transactional insert array of entities
-        const insertResponse = await appServer.AppDataSource.getRepository(Assistant).insert(prepVariables)
+        const insertResponse = await repository.insert(prepVariables)
 
         return insertResponse
     } catch (error) {

diff --git a/packages/server/src/services/chatflows/index.ts b/packages/server/src/services/chatflows/index.ts
@@ -14,6 +14,7 @@ import { getRunningExpressApp } from '../../utils/getRunningExpressApp'
 import { utilGetUploadsConfig } from '../../utils/getUploadsConfig'
 import logger from '../../utils/logger'
 import { FLOWISE_METRIC_COUNTERS, FLOWISE_COUNTER_STATUS } from '../../Interface.Metrics'
+import { QueryRunner } from 'typeorm'
 
 // Check if chatflow valid for streaming
 const checkIfChatflowIsValidForStreaming = async (chatflowId: string): Promise<any> => {
@@ -206,9 +207,10 @@ const saveChatflow = async (newChatFlow: ChatFlow): Promise<any> => {
     }
 }
 
-const importChatflows = async (newChatflows: Partial<ChatFlow>[]): Promise<any> => {
+const importChatflows = async (newChatflows: Partial<ChatFlow>[], queryRunner?: QueryRunner): Promise<any> => {
     try {
         const appServer = getRunningExpressApp()
+        const repository = queryRunner ? queryRunner.manager.getRepository(ChatFlow) : appServer.AppDataSource.getRepository(ChatFlow)
 
         // step 1 - check whether file chatflows array is zero
         if (newChatflows.length == 0) return
@@ -224,11 +226,7 @@ const importChatflows = async (newChatflows: Partial<ChatFlow>[]): Promise<any>
             count += 1
         })
 
-        const selectResponse = await appServer.AppDataSource.getRepository(ChatFlow)
-            .createQueryBuilder('cf')
-            .select('cf.id')
-            .where(`cf.id IN ${ids}`)
-            .getMany()
+        const selectResponse = await repository.createQueryBuilder('cf').select('cf.id').where(`cf.id IN ${ids}`).getMany()
         const foundIds = selectResponse.map((response) => {
             return response.id
         })
@@ -248,7 +246,7 @@ const importChatflows = async (newChatflows: Partial<ChatFlow>[]): Promise<any>
         })
 
         // step 4 - transactional insert array of entities
-        const insertResponse = await appServer.AppDataSource.getRepository(ChatFlow).insert(prepChatflows)
+        const insertResponse = await repository.insert(prepChatflows)
 
         return insertResponse
     } catch (error) {

diff --git a/packages/server/src/services/export-import/index.ts b/packages/server/src/services/export-import/index.ts
@@ -87,22 +87,22 @@ const importData = async (importData: ExportData) => {
         const queryRunner = appServer.AppDataSource.createQueryRunner()
 
         try {
-            queryRunner.startTransaction()
+            await queryRunner.startTransaction()
 
-            // step 1 - importTools
-            if (importData.Tool.length > 0) await toolsService.importTools(importData.Tool)
-            // step 2 - importChatflows
-            if (importData.ChatFlow.length > 0) await chatflowService.importChatflows(importData.ChatFlow)
-            // step 3 - importAgentlows
-            if (importData.AgentFlow.length > 0) await chatflowService.importChatflows(importData.AgentFlow)
-            if (importData.Variable.length > 0) await variableService.importVariables(importData.Variable)
-            if (importData.Assistant.length > 0) await assistantService.importAssistants(importData.Assistant)
-            queryRunner.commitTransaction()
+            if (importData.Tool.length > 0) await toolsService.importTools(importData.Tool, queryRunner)
+            if (importData.ChatFlow.length > 0) await chatflowService.importChatflows(importData.ChatFlow, queryRunner)
+            if (importData.AgentFlow.length > 0) await chatflowService.importChatflows(importData.AgentFlow, queryRunner)
+            if (importData.Variable.length > 0) await variableService.importVariables(importData.Variable, queryRunner)
+            if (importData.Assistant.length > 0) await assistantService.importAssistants(importData.Assistant, queryRunner)
+
+            await queryRunner.commitTransaction()
         } catch (error) {
-            queryRunner.rollbackTransaction()
+            await queryRunner.rollbackTransaction()
             throw error
         } finally {
-            queryRunner.release()
+            if (!queryRunner.isReleased) {
+                await queryRunner.release()
+            }
         }
     } catch (error) {
         throw new InternalFlowiseError(

diff --git a/packages/server/src/services/tools/index.ts b/packages/server/src/services/tools/index.ts
@@ -5,6 +5,7 @@ import { getErrorMessage } from '../../errors/utils'
 import { getAppVersion } from '../../utils'
 import { getRunningExpressApp } from '../../utils/getRunningExpressApp'
 import { FLOWISE_METRIC_COUNTERS, FLOWISE_COUNTER_STATUS } from '../../Interface.Metrics'
+import { QueryRunner } from 'typeorm'
 
 const createTool = async (requestBody: any): Promise<any> => {
     try {
@@ -81,9 +82,10 @@ const updateTool = async (toolId: string, toolBody: any): Promise<any> => {
     }
 }
 
-const importTools = async (newTools: Partial<Tool>[]) => {
+const importTools = async (newTools: Partial<Tool>[], queryRunner?: QueryRunner) => {
     try {
         const appServer = getRunningExpressApp()
+        const repository = queryRunner ? queryRunner.manager.getRepository(Tool) : appServer.AppDataSource.getRepository(Tool)
 
         // step 1 - check whether file tools array is zero
         if (newTools.length == 0) return
@@ -99,11 +101,7 @@ const importTools = async (newTools: Partial<Tool>[]) => {
             count += 1
         })
 
-        const selectResponse = await appServer.AppDataSource.getRepository(Tool)
-            .createQueryBuilder('t')
-            .select('t.id')
-            .where(`t.id IN ${ids}`)
-            .getMany()
+        const selectResponse = await repository.createQueryBuilder('t').select('t.id').where(`t.id IN ${ids}`).getMany()
         const foundIds = selectResponse.map((response) => {
             return response.id
         })
@@ -120,7 +118,7 @@ const importTools = async (newTools: Partial<Tool>[]) => {
         })
 
         // step 4 - transactional insert array of entities
-        const insertResponse = await appServer.AppDataSource.getRepository(Tool).insert(prepTools)
+        const insertResponse = await repository.insert(prepTools)
 
         return insertResponse
     } catch (error) {

diff --git a/packages/server/src/services/variables/index.ts b/packages/server/src/services/variables/index.ts
@@ -3,6 +3,7 @@ import { getRunningExpressApp } from '../../utils/getRunningExpressApp'
 import { Variable } from '../../database/entities/Variable'
 import { InternalFlowiseError } from '../../errors/internalFlowiseError'
 import { getErrorMessage } from '../../errors/utils'
+import { QueryRunner } from 'typeorm'
 
 const createVariable = async (newVariable: Variable) => {
     try {
@@ -73,9 +74,10 @@ const updateVariable = async (variable: Variable, updatedVariable: Variable) =>
     }
 }
 
-const importVariables = async (newVariables: Partial<Variable>[]): Promise<any> => {
+const importVariables = async (newVariables: Partial<Variable>[], queryRunner?: QueryRunner): Promise<any> => {
     try {
         const appServer = getRunningExpressApp()
+        const repository = queryRunner ? queryRunner.manager.getRepository(Variable) : appServer.AppDataSource.getRepository(Variable)
 
         // step 1 - check whether array is zero
         if (newVariables.length == 0) return
@@ -91,11 +93,7 @@ const importVariables = async (newVariables: Partial<Variable>[]): Promise<any>
             count += 1
         })
 
-        const selectResponse = await appServer.AppDataSource.getRepository(Variable)
-            .createQueryBuilder('v')
-            .select('v.id')
-            .where(`v.id IN ${ids}`)
-            .getMany()
+        const selectResponse = await repository.createQueryBuilder('v').select('v.id').where(`v.id IN ${ids}`).getMany()
         const foundIds = selectResponse.map((response) => {
             return response.id
         })
@@ -112,7 +110,7 @@ const importVariables = async (newVariables: Partial<Variable>[]): Promise<any>
         })
 
         // step 4 - transactional insert array of entities
-        const insertResponse = await appServer.AppDataSource.getRepository(Variable).insert(prepVariables)
+        const insertResponse = await repository.insert(prepVariables)
 
         return insertResponse
     } catch (error) {