Fix MySQLRecordManager DataSource connection churn

packages/components/nodes/recordmanager/MySQLRecordManager/MySQLrecordManager.test.ts

@@ -0,0 +1,138 @@
+const mockDataSourceInstances: Array<{
+    isInitialized: boolean
+    initialize: jest.Mock
+    destroy: jest.Mock
+    createQueryRunner: jest.Mock
+}> = []
+
+const mockQueryRunners: Array<{
+    isReleased: boolean
+    manager: {
+        query: jest.Mock
+    }
+    release: jest.Mock
+}> = []
+
+jest.mock('typeorm', () => {
+    class MockDataSource {
+        options: unknown
+        isInitialized = false
+        initialize = jest.fn(async () => {
+            this.isInitialized = true
+            return this
+        })
+        destroy = jest.fn(async () => {
+            this.isInitialized = false
+        })
+        createQueryRunner = jest.fn(() => {
+            const queryRunner = {
+                isReleased: false,
+                manager: {
+                    query: jest.fn(async (query: string) => {
+                        if (query.includes('UNIX_TIMESTAMP')) {
+                            return [{ epoch: '123' }]
+                        }
+
+                        if (query.includes('SELECT `key`')) {
+                            return [{ key: 'a' }]
+                        }
+
+                        return []
+                    })
+                },
+                release: jest.fn(async () => {
+                    queryRunner.isReleased = true
+                })
+            }
+
+            mockQueryRunners.push(queryRunner)
+            return queryRunner
+        })
+
+        constructor(options: unknown) {
+            this.options = options
+            mockDataSourceInstances.push(this)
+        }
+    }
+
+    return {
+        DataSource: MockDataSource
+    }
+})
+
+const { MySQLRecordManager } = require('./MySQLrecordManager')
+
+const createManager = () =>
+    new MySQLRecordManager('test_namespace', {
+        mysqlOptions: {
+            type: 'mysql',
+            host: 'localhost',
+            port: 3306,
+            username: 'user',
+            password: 'password',
+            database: 'flowise'
+        },
+        tableName: 'upsertion_records'
+    })
+
+describe('MySQLRecordManager connection lifecycle', () => {
+    beforeEach(() => {
+        jest.clearAllMocks()
+        mockDataSourceInstances.length = 0
+        mockQueryRunners.length = 0
+    })
+
+    it('reuses one initialized DataSource across operations and destroys it only on close', async () => {
+        const manager = createManager()
+
+        await expect(manager.exists(['a'])).resolves.toEqual([true])
+        await expect(manager.listKeys()).resolves.toEqual(['a'])
+        await expect(manager.getTime()).resolves.toBe(123)
+
+        expect(mockDataSourceInstances).toHaveLength(1)
+        expect(mockDataSourceInstances[0].initialize).toHaveBeenCalledTimes(1)
+        expect(mockDataSourceInstances[0].destroy).not.toHaveBeenCalled()
+        expect(mockDataSourceInstances[0].createQueryRunner).toHaveBeenCalledTimes(3)
+        expect(mockQueryRunners).toHaveLength(3)
+        mockQueryRunners.forEach((queryRunner) => {
+            expect(queryRunner.release).toHaveBeenCalledTimes(1)
+            expect(queryRunner.isReleased).toBe(true)
+        })
+
+        await manager.close()
+
+        expect(mockDataSourceInstances[0].destroy).toHaveBeenCalledTimes(1)
+        expect(mockDataSourceInstances[0].isInitialized).toBe(false)
+    })
+
+    it('shares the pending DataSource initialization across concurrent operations', async () => {
+        const manager = createManager()
+
+        await expect(Promise.all([manager.exists(['a']), manager.listKeys()])).resolves.toEqual([[true], ['a']])
+
+        expect(mockDataSourceInstances).toHaveLength(1)
+        expect(mockDataSourceInstances[0].initialize).toHaveBeenCalledTimes(1)
+        expect(mockDataSourceInstances[0].destroy).not.toHaveBeenCalled()
+        expect(mockDataSourceInstances[0].createQueryRunner).toHaveBeenCalledTimes(2)
+
+        await manager.close()
+    })
+
+    it('releases query runners when update validation fails', async () => {
+        const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => undefined)
+        const manager = createManager()
+
+        await expect(manager.update(['a'], { timeAtLeast: 999 })).rejects.toThrow('Time sync issue with database 123 < 999')
+
+        expect(mockDataSourceInstances).toHaveLength(1)
+        expect(mockDataSourceInstances[0].destroy).not.toHaveBeenCalled()
+        expect(mockQueryRunners).toHaveLength(2)
+        mockQueryRunners.forEach((queryRunner) => {
+            expect(queryRunner.release).toHaveBeenCalledTimes(1)
+            expect(queryRunner.isReleased).toBe(true)
+        })
+
+        await manager.close()
+        consoleErrorSpy.mockRestore()
+    })
+})

packages/components/nodes/recordmanager/MySQLRecordManager/MySQLrecordManager.ts

@@ -2,7 +2,7 @@ import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Inter
 import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
 import { sanitizeDataSourceOptions } from '../../../src/sanitizeDataSourceOptions'
 import { ListKeyOptions, RecordManagerInterface, UpdateOptions } from '@langchain/community/indexes/base'
-import { DataSource } from 'typeorm'
+import { DataSource, QueryRunner } from 'typeorm'
 
 class MySQLRecordManager_RecordManager implements INode {
     label: string
@@ -173,6 +173,8 @@ class MySQLRecordManager implements RecordManagerInterface {
     config: MySQLRecordManagerOptions
     tableName: string
     namespace: string
+    private dataSource?: DataSource
+    private dataSourcePromise?: Promise<DataSource>
 
     constructor(namespace: string, config: MySQLRecordManagerOptions) {
         const { tableName } = config
@@ -202,15 +204,56 @@ class MySQLRecordManager implements RecordManagerInterface {
         if (mysqlOptions.port === 5432) {
             throw new Error('Invalid port number')
         }
-        const dataSource = new DataSource(mysqlOptions)
-        await dataSource.initialize()
-        return dataSource
+
+        if (this.dataSource?.isInitialized) {
+            return this.dataSource
+        }
+
+        if (this.dataSourcePromise) {
+            return this.dataSourcePromise
+        }
+
+        this.dataSourcePromise = (async () => {
+            const dataSource = new DataSource(mysqlOptions)
+            await dataSource.initialize()
+            this.dataSource = dataSource
+            return dataSource
+        })()
+
+        try {
+            return await this.dataSourcePromise
+        } catch (error) {
+            this.dataSource = undefined
+            throw error
+        } finally {
+            this.dataSourcePromise = undefined
+        }
     }
 
-    async createSchema(): Promise<void> {
+    private async getQueryRunner(): Promise<QueryRunner> {
         const dataSource = await this.getDataSource()
+        return dataSource.createQueryRunner()
+    }
+
+    async close(): Promise<void> {
+        const dataSource = this.dataSourcePromise ? await this.dataSourcePromise.catch(() => undefined) : this.dataSource
+        this.dataSourcePromise = undefined
+        this.dataSource = undefined
+
+        if (dataSource?.isInitialized) {
+            await dataSource.destroy()
+        }
+    }
+
+    private async releaseQueryRunner(queryRunner: QueryRunner): Promise<void> {
+        if (!queryRunner.isReleased) {
+            await queryRunner.release()
+        }
+    }
+
+    async createSchema(): Promise<void> {
+        const queryRunner = await this.getQueryRunner()
         try {
-            const queryRunner = dataSource.createQueryRunner()
             const tableName = this.sanitizeTableName(this.tableName)
 
             await queryRunner.manager.query(`create table if not exists \`${this.sanitizeTableName(tableName)}\` (
@@ -259,8 +302,6 @@ class MySQLRecordManager implements RecordManagerInterface {
                     }
                 }
             }
-
-            await queryRunner.release()
         } catch (e: any) {
             // This error indicates that the table already exists
             // Due to asynchronous nature of the code, it is possible that
@@ -271,22 +312,20 @@ class MySQLRecordManager implements RecordManagerInterface {
             }
             throw e
         } finally {
-            await dataSource.destroy()
+            await this.releaseQueryRunner(queryRunner)
         }
     }
 
     async getTime(): Promise<number> {
-        const dataSource = await this.getDataSource()
+        const queryRunner = await this.getQueryRunner()
         try {
-            const queryRunner = dataSource.createQueryRunner()
             const res = await queryRunner.manager.query(`SELECT UNIX_TIMESTAMP(NOW()) AS epoch`)
-            await queryRunner.release()
             return Number.parseFloat(res[0].epoch)
         } catch (error) {
             console.error('Error getting time in MySQLRecordManager:')
             throw error
         } finally {
-            await dataSource.destroy()
+            await this.releaseQueryRunner(queryRunner)
         }
     }
 
@@ -295,48 +334,45 @@ class MySQLRecordManager implements RecordManagerInterface {
             return
         }
 
-        const dataSource = await this.getDataSource()
-        const queryRunner = dataSource.createQueryRunner()
-        const tableName = this.sanitizeTableName(this.tableName)
+        const queryRunner = await this.getQueryRunner()
+        try {
+            const tableName = this.sanitizeTableName(this.tableName)
 
-        const updatedAt = await this.getTime()
-        const { timeAtLeast, groupIds: _groupIds } = updateOptions ?? {}
+            const updatedAt = await this.getTime()
+            const { timeAtLeast, groupIds: _groupIds } = updateOptions ?? {}
 
-        if (timeAtLeast && updatedAt < timeAtLeast) {
-            throw new Error(`Time sync issue with database ${updatedAt} < ${timeAtLeast}`)
-        }
+            if (timeAtLeast && updatedAt < timeAtLeast) {
+                throw new Error(`Time sync issue with database ${updatedAt} < ${timeAtLeast}`)
+            }
 
-        // Handle both new format (objects with uid and docId) and old format (strings)
-        const isNewFormat = keys.length > 0 && typeof keys[0] === 'object' && 'uid' in keys[0]
-        const keyStrings = isNewFormat ? (keys as Array<{ uid: string; docId: string }>).map((k) => k.uid) : (keys as string[])
-        const docIds = isNewFormat ? (keys as Array<{ uid: string; docId: string }>).map((k) => k.docId) : keys.map(() => null)
+            // Handle both new format (objects with uid and docId) and old format (strings)
+            const isNewFormat = keys.length > 0 && typeof keys[0] === 'object' && 'uid' in keys[0]
+            const keyStrings = isNewFormat ? (keys as Array<{ uid: string; docId: string }>).map((k) => k.uid) : (keys as string[])
+            const docIds = isNewFormat ? (keys as Array<{ uid: string; docId: string }>).map((k) => k.docId) : keys.map(() => null)
 
-        const groupIds = _groupIds ?? keyStrings.map(() => null)
+            const groupIds = _groupIds ?? keyStrings.map(() => null)
 
-        if (groupIds.length !== keyStrings.length) {
-            throw new Error(`Number of keys (${keyStrings.length}) does not match number of group_ids (${groupIds.length})`)
-        }
+            if (groupIds.length !== keyStrings.length) {
+                throw new Error(`Number of keys (${keyStrings.length}) does not match number of group_ids (${groupIds.length})`)
+            }
 
-        const recordsToUpsert = keyStrings.map((key, i) => [key, this.namespace, updatedAt, groupIds[i] ?? null, docIds[i] ?? null])
+            const recordsToUpsert = keyStrings.map((key, i) => [key, this.namespace, updatedAt, groupIds[i] ?? null, docIds[i] ?? null])
 
-        const query = `
+            const query = `
             INSERT INTO \`${tableName}\` (\`key\`, \`namespace\`, \`updated_at\`, \`group_id\`, \`doc_id\`)
             VALUES (?, ?, ?, ?, ?)
             ON DUPLICATE KEY UPDATE \`updated_at\` = VALUES(\`updated_at\`), \`doc_id\` = VALUES(\`doc_id\`)`
 
-        // To handle multiple files upsert
-        try {
+            // To handle multiple files upsert
             for (const record of recordsToUpsert) {
                 // Consider using a transaction for batch operations
                 await queryRunner.manager.query(query, record.flat())
             }
-
-            await queryRunner.release()
         } catch (error) {
             console.error('Error updating in MySQLRecordManager:')
             throw error
         } finally {
-            await dataSource.destroy()
+            await this.releaseQueryRunner(queryRunner)
         }
     }
 
@@ -345,8 +381,7 @@ class MySQLRecordManager implements RecordManagerInterface {
             return []
         }
 
-        const dataSource = await this.getDataSource()
-        const queryRunner = dataSource.createQueryRunner()
+        const queryRunner = await this.getQueryRunner()
         const tableName = this.sanitizeTableName(this.tableName)
 
         // Prepare the placeholders and the query
@@ -368,19 +403,17 @@ class MySQLRecordManager implements RecordManagerInterface {
             keys.forEach((key, index) => {
                 existsArray[index] = existingKeysSet.has(key)
             })
-            await queryRunner.release()
             return existsArray
         } catch (error) {
             console.error('Error checking existence of keys')
             throw error
         } finally {
-            await dataSource.destroy()
+            await this.releaseQueryRunner(queryRunner)
         }
     }
 
     async listKeys(options?: ListKeyOptions & { docId?: string }): Promise<string[]> {
-        const dataSource = await this.getDataSource()
-        const queryRunner = dataSource.createQueryRunner()
+        const queryRunner = await this.getQueryRunner()
         const tableName = this.sanitizeTableName(this.tableName)
 
         try {
@@ -420,13 +453,12 @@ class MySQLRecordManager implements RecordManagerInterface {
 
             // Directly using try/catch with async/await for cleaner flow
             const result = await queryRunner.manager.query(query, values)
-            await queryRunner.release()
             return result.map((row: { key: string }) => row.key)
         } catch (error) {
             console.error('MySQLRecordManager listKeys Error: ')
             throw error
         } finally {
-            await dataSource.destroy()
+            await this.releaseQueryRunner(queryRunner)
         }
     }
 
@@ -435,8 +467,7 @@ class MySQLRecordManager implements RecordManagerInterface {
             return
         }
 
-        const dataSource = await this.getDataSource()
-        const queryRunner = dataSource.createQueryRunner()
+        const queryRunner = await this.getQueryRunner()
         const tableName = this.sanitizeTableName(this.tableName)
 
         const placeholders = keys.map(() => '?').join(', ')
@@ -446,14 +477,13 @@ class MySQLRecordManager implements RecordManagerInterface {
         // Directly using try/catch with async/await for cleaner flow
         try {
             await queryRunner.manager.query(query, values)
-            await queryRunner.release()
         } catch (error) {
             console.error('Error deleting keys')
             throw error
         } finally {
-            await dataSource.destroy()
+            await this.releaseQueryRunner(queryRunner)
         }
     }
 }
 
-module.exports = { nodeClass: MySQLRecordManager_RecordManager }
+module.exports = { nodeClass: MySQLRecordManager_RecordManager, MySQLRecordManager }