mirror of
https://github.com/farcasclaudiu/Flowise.git
synced 2026-06-28 07:00:49 +03:00
Bugfix/Prevent open connections on typeorm datasource (#3652)
prevent open connections on typeorm datasource
This commit is contained in:
+86
-29
@@ -1,7 +1,7 @@
|
||||
import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
|
||||
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
|
||||
import { ListKeyOptions, RecordManagerInterface, UpdateOptions } from '@langchain/community/indexes/base'
|
||||
import { DataSource, QueryRunner } from 'typeorm'
|
||||
import { DataSource } from 'typeorm'
|
||||
import { getHost } from '../../vectorstores/Postgres/utils'
|
||||
import { getDatabase, getPort, getTableName } from './utils'
|
||||
|
||||
@@ -175,29 +175,37 @@ type PostgresRecordManagerOptions = {
|
||||
|
||||
class PostgresRecordManager implements RecordManagerInterface {
|
||||
lc_namespace = ['langchain', 'recordmanagers', 'postgres']
|
||||
|
||||
datasource: DataSource
|
||||
|
||||
queryRunner: QueryRunner
|
||||
|
||||
config: PostgresRecordManagerOptions
|
||||
tableName: string
|
||||
|
||||
namespace: string
|
||||
|
||||
constructor(namespace: string, config: PostgresRecordManagerOptions) {
|
||||
const { postgresConnectionOptions, tableName } = config
|
||||
const { tableName } = config
|
||||
this.namespace = namespace
|
||||
this.datasource = new DataSource(postgresConnectionOptions)
|
||||
this.tableName = tableName
|
||||
this.config = config
|
||||
}
|
||||
|
||||
private async getDataSource(): Promise<DataSource> {
|
||||
const { postgresConnectionOptions } = this.config
|
||||
if (!postgresConnectionOptions) {
|
||||
throw new Error('No datasource options provided')
|
||||
}
|
||||
// Prevent using default MySQL port, otherwise will throw uncaught error and crashing the app
|
||||
if (postgresConnectionOptions.port === 3006) {
|
||||
throw new Error('Invalid port number')
|
||||
}
|
||||
const dataSource = new DataSource(postgresConnectionOptions)
|
||||
await dataSource.initialize()
|
||||
return dataSource
|
||||
}
|
||||
|
||||
async createSchema(): Promise<void> {
|
||||
try {
|
||||
const appDataSource = await this.datasource.initialize()
|
||||
const dataSource = await this.getDataSource()
|
||||
const queryRunner = dataSource.createQueryRunner()
|
||||
|
||||
this.queryRunner = appDataSource.createQueryRunner()
|
||||
|
||||
await this.queryRunner.manager.query(`
|
||||
await queryRunner.manager.query(`
|
||||
CREATE TABLE IF NOT EXISTS "${this.tableName}" (
|
||||
uuid UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
key TEXT NOT NULL,
|
||||
@@ -210,6 +218,8 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
CREATE INDEX IF NOT EXISTS key_index ON "${this.tableName}" (key);
|
||||
CREATE INDEX IF NOT EXISTS namespace_index ON "${this.tableName}" (namespace);
|
||||
CREATE INDEX IF NOT EXISTS group_id_index ON "${this.tableName}" (group_id);`)
|
||||
|
||||
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
|
||||
@@ -223,8 +233,18 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
}
|
||||
|
||||
async getTime(): Promise<number> {
|
||||
const res = await this.queryRunner.manager.query('SELECT EXTRACT(EPOCH FROM CURRENT_TIMESTAMP)')
|
||||
return Number.parseFloat(res[0].extract)
|
||||
const dataSource = await this.getDataSource()
|
||||
try {
|
||||
const queryRunner = dataSource.createQueryRunner()
|
||||
const res = await queryRunner.manager.query('SELECT EXTRACT(EPOCH FROM CURRENT_TIMESTAMP)')
|
||||
await queryRunner.release()
|
||||
return Number.parseFloat(res[0].extract)
|
||||
} catch (error) {
|
||||
console.error('Error getting time in PostgresRecordManager:')
|
||||
throw error
|
||||
} finally {
|
||||
await dataSource.destroy()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -247,6 +267,9 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
return
|
||||
}
|
||||
|
||||
const dataSource = await this.getDataSource()
|
||||
const queryRunner = dataSource.createQueryRunner()
|
||||
|
||||
const updatedAt = await this.getTime()
|
||||
const { timeAtLeast, groupIds: _groupIds } = updateOptions ?? {}
|
||||
|
||||
@@ -265,7 +288,15 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
const valuesPlaceholders = recordsToUpsert.map((_, j) => this.generatePlaceholderForRowAt(j, recordsToUpsert[0].length)).join(', ')
|
||||
|
||||
const query = `INSERT INTO "${this.tableName}" (key, namespace, updated_at, group_id) VALUES ${valuesPlaceholders} ON CONFLICT (key, namespace) DO UPDATE SET updated_at = EXCLUDED.updated_at;`
|
||||
await this.queryRunner.manager.query(query, recordsToUpsert.flat())
|
||||
try {
|
||||
await queryRunner.manager.query(query, recordsToUpsert.flat())
|
||||
await queryRunner.release()
|
||||
} catch (error) {
|
||||
console.error('Error updating in PostgresRecordManager:')
|
||||
throw error
|
||||
} finally {
|
||||
await dataSource.destroy()
|
||||
}
|
||||
}
|
||||
|
||||
async exists(keys: string[]): Promise<boolean[]> {
|
||||
@@ -273,14 +304,25 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
return []
|
||||
}
|
||||
|
||||
const dataSource = await this.getDataSource()
|
||||
const queryRunner = dataSource.createQueryRunner()
|
||||
|
||||
const startIndex = 2
|
||||
const arrayPlaceholders = keys.map((_, i) => `$${i + startIndex}`).join(', ')
|
||||
|
||||
const query = `
|
||||
SELECT k, (key is not null) ex from unnest(ARRAY[${arrayPlaceholders}]) k left join "${this.tableName}" on k=key and namespace = $1;
|
||||
`
|
||||
const res = await this.queryRunner.manager.query(query, [this.namespace, ...keys.flat()])
|
||||
return res.map((row: { ex: boolean }) => row.ex)
|
||||
try {
|
||||
const res = await queryRunner.manager.query(query, [this.namespace, ...keys.flat()])
|
||||
await queryRunner.release()
|
||||
return res.map((row: { ex: boolean }) => row.ex)
|
||||
} catch (error) {
|
||||
console.error('Error checking existence of keys in PostgresRecordManager:')
|
||||
throw error
|
||||
} finally {
|
||||
await dataSource.destroy()
|
||||
}
|
||||
}
|
||||
|
||||
async listKeys(options?: ListKeyOptions): Promise<string[]> {
|
||||
@@ -314,8 +356,20 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
}
|
||||
|
||||
query += ';'
|
||||
const res = await this.queryRunner.manager.query(query, values)
|
||||
return res.map((row: { key: string }) => row.key)
|
||||
|
||||
const dataSource = await this.getDataSource()
|
||||
const queryRunner = dataSource.createQueryRunner()
|
||||
|
||||
try {
|
||||
const res = await queryRunner.manager.query(query, values)
|
||||
await queryRunner.release()
|
||||
return res.map((row: { key: string }) => row.key)
|
||||
} catch (error) {
|
||||
console.error('Error listing keys in PostgresRecordManager:')
|
||||
throw error
|
||||
} finally {
|
||||
await dataSource.destroy()
|
||||
}
|
||||
}
|
||||
|
||||
async deleteKeys(keys: string[]): Promise<void> {
|
||||
@@ -323,16 +377,19 @@ class PostgresRecordManager implements RecordManagerInterface {
|
||||
return
|
||||
}
|
||||
|
||||
const query = `DELETE FROM "${this.tableName}" WHERE namespace = $1 AND key = ANY($2);`
|
||||
await this.queryRunner.manager.query(query, [this.namespace, keys])
|
||||
}
|
||||
const dataSource = await this.getDataSource()
|
||||
const queryRunner = dataSource.createQueryRunner()
|
||||
|
||||
/**
|
||||
* Terminates the connection pool.
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async end(): Promise<void> {
|
||||
if (this.datasource && this.datasource.isInitialized) await this.datasource.destroy()
|
||||
try {
|
||||
const query = `DELETE FROM "${this.tableName}" WHERE namespace = $1 AND key = ANY($2);`
|
||||
await queryRunner.manager.query(query, [this.namespace, keys])
|
||||
await queryRunner.release()
|
||||
} catch (error) {
|
||||
console.error('Error deleting keys')
|
||||
throw error
|
||||
} finally {
|
||||
await dataSource.destroy()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user