mirror of
https://github.com/hcengineering/platform.git
synced 2026-09-10 19:57:43 +02:00
Extract postgres base package (#9397)
Signed-off-by: Andrey Sobolev <haiodo@gmail.com>
This commit is contained in:
Andrey Sobolev
@@ -0,0 +1,434 @@
|
||||
//
|
||||
// Copyright © 2025 Hardcore Engineering Inc.
|
||||
//
|
||||
// Licensed under the Eclipse Public License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License. You may
|
||||
// obtain a copy of the License at https://www.eclipse.org/legal/epl-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
//
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
import postgres, { type Options, type ParameterOrJSON } from 'postgres'
|
||||
|
||||
const clientRefs = new Map<number, ClientRef>()
|
||||
|
||||
let clId = 0
|
||||
|
||||
export type DBResult = any[] & { count: number }
|
||||
export interface DBClient {
|
||||
execute: (query: string, parameters?: ParameterOrJSON<any>[] | undefined) => Promise<DBResult>
|
||||
|
||||
release: () => void
|
||||
|
||||
reserve: () => Promise<DBClient>
|
||||
|
||||
raw: () => postgres.Sql
|
||||
}
|
||||
|
||||
export function createDBClient (client: postgres.Sql, release: () => void = () => {}): DBClient {
|
||||
return {
|
||||
execute: (query, parameters) =>
|
||||
client.unsafe(query, doFetchTypes ? parameters : convertArrayParams(parameters), getPrepare()),
|
||||
release,
|
||||
reserve: async () => {
|
||||
const reserved = await client.reserve()
|
||||
return createDBClient(reserved, () => {
|
||||
reserved.release()
|
||||
})
|
||||
},
|
||||
raw: () => client
|
||||
}
|
||||
}
|
||||
|
||||
export function convertArrayParams (params?: unknown[]): any[] | undefined {
|
||||
if (params === undefined) return undefined
|
||||
|
||||
return params.map((param) => {
|
||||
if (!Array.isArray(param)) return param
|
||||
|
||||
if (param.length === 0) return '{}'
|
||||
|
||||
const sanitized = param.map((item) => {
|
||||
if (item === null || item === undefined) return 'NULL'
|
||||
|
||||
if (typeof item === 'number' || typeof item === 'boolean') {
|
||||
return String(item)
|
||||
}
|
||||
|
||||
if (typeof item === 'string') {
|
||||
const escaped = item.replace(/\\/g, '\\\\').replace(/"/g, '\\"')
|
||||
return `"${escaped}"`
|
||||
}
|
||||
|
||||
const json = JSON.stringify(item)
|
||||
const escapedJson = json.replace(/\\/g, '\\\\').replace(/"/g, '\\"')
|
||||
return `"${escapedJson}"`
|
||||
})
|
||||
|
||||
return `{${sanitized.join(',')}}`
|
||||
})
|
||||
}
|
||||
|
||||
export async function retryTxn (
|
||||
pool: postgres.Sql,
|
||||
operation: (client: postgres.TransactionSql) => Promise<any>
|
||||
): Promise<any> {
|
||||
await pool.begin(async (client) => {
|
||||
const result = await operation(client)
|
||||
return result
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* @public
|
||||
*/
|
||||
export async function shutdownPostgres (): Promise<void> {
|
||||
for (const c of connections.values()) {
|
||||
c.close(true)
|
||||
}
|
||||
connections.clear()
|
||||
}
|
||||
|
||||
export interface PostgresClientReference {
|
||||
getClient: () => Promise<postgres.Sql>
|
||||
|
||||
mgr: ConnectionMgr
|
||||
close: () => void
|
||||
url: () => string
|
||||
}
|
||||
|
||||
class PostgresClientReferenceImpl {
|
||||
count: number
|
||||
client: postgres.Sql
|
||||
|
||||
mgr: ConnectionMgr
|
||||
|
||||
constructor (
|
||||
readonly connectionString: string,
|
||||
client: postgres.Sql,
|
||||
readonly onclose: () => void
|
||||
) {
|
||||
this.count = 0
|
||||
this.client = client
|
||||
this.mgr = new ConnectionMgr(createDBClient(this.client))
|
||||
}
|
||||
|
||||
url (): string {
|
||||
return this.connectionString
|
||||
}
|
||||
|
||||
getClient (): postgres.Sql {
|
||||
return this.client
|
||||
}
|
||||
|
||||
close (force: boolean = false): void {
|
||||
this.count--
|
||||
if (this.count === 0 || force) {
|
||||
if (force) {
|
||||
this.count = 0
|
||||
}
|
||||
void (async () => {
|
||||
this.onclose()
|
||||
const cl = this.client
|
||||
await cl.end({ timeout: 1 })
|
||||
})()
|
||||
}
|
||||
}
|
||||
|
||||
addRef (): void {
|
||||
this.count++
|
||||
}
|
||||
}
|
||||
export class ClientRef implements PostgresClientReference {
|
||||
id = ++clId
|
||||
constructor (
|
||||
readonly client: PostgresClientReferenceImpl,
|
||||
readonly mgr: ConnectionMgr
|
||||
) {
|
||||
clientRefs.set(this.id, this)
|
||||
}
|
||||
|
||||
url (): string {
|
||||
return this.client.url()
|
||||
}
|
||||
|
||||
closed = false
|
||||
async getClient (): Promise<postgres.Sql> {
|
||||
if (!this.closed) {
|
||||
return this.client.getClient()
|
||||
} else {
|
||||
throw Error('DB client is already closed')
|
||||
}
|
||||
}
|
||||
|
||||
close (): void {
|
||||
// Do not allow double close of mongo connection client
|
||||
if (!this.closed) {
|
||||
clientRefs.delete(this.id)
|
||||
this.closed = true
|
||||
this.client.close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export let dbExtraOptions: Partial<Options<any>> = {}
|
||||
export function setDBExtraOptions (options: Partial<Options<any>>): void {
|
||||
dbExtraOptions = options
|
||||
}
|
||||
|
||||
export function getPrepare (): { prepare: boolean } {
|
||||
return { prepare: dbExtraOptions.prepare ?? false }
|
||||
}
|
||||
|
||||
export const doFetchTypes = true
|
||||
|
||||
const connections = new Map<string, PostgresClientReferenceImpl>()
|
||||
|
||||
/**
|
||||
* Initialize a connection to DB
|
||||
* @public
|
||||
*/
|
||||
export function getDBClient (
|
||||
connectionString: string,
|
||||
database?: string,
|
||||
serviceName: string = 'transactor'
|
||||
): PostgresClientReference {
|
||||
const extraOptions = JSON.parse(process.env.POSTGRES_OPTIONS ?? '{}')
|
||||
const key = `${connectionString}${extraOptions}`
|
||||
|
||||
let existing = connections.get(key)
|
||||
|
||||
if (existing === undefined) {
|
||||
const sql = postgres(connectionString, {
|
||||
connection: {
|
||||
application_name: serviceName
|
||||
},
|
||||
database,
|
||||
max: 10,
|
||||
min: 2,
|
||||
connect_timeout: 30,
|
||||
idle_timeout: 0,
|
||||
transform: {
|
||||
undefined: null
|
||||
},
|
||||
debug: false,
|
||||
notice: false,
|
||||
onnotice (notice) {},
|
||||
onparameter (key, value) {},
|
||||
...dbExtraOptions,
|
||||
...extraOptions,
|
||||
fetch_types: doFetchTypes
|
||||
})
|
||||
|
||||
existing = new PostgresClientReferenceImpl(connectionString, sql, () => {
|
||||
connections.delete(key)
|
||||
})
|
||||
connections.set(key, existing)
|
||||
}
|
||||
// Add reference and return once closable
|
||||
existing.addRef()
|
||||
return new ClientRef(existing, existing.mgr)
|
||||
}
|
||||
|
||||
class ConnectionInfo {
|
||||
// It should preserve at least one available connection in pool, other connection should be closed
|
||||
available: DBClient[] = []
|
||||
|
||||
released: boolean = false
|
||||
|
||||
constructor (
|
||||
readonly connectionId: string,
|
||||
protected readonly client: DBClient,
|
||||
readonly managed: boolean
|
||||
) {}
|
||||
|
||||
async withReserve (action: (reservedClient: DBClient) => Promise<any>, forced: boolean = false): Promise<any> {
|
||||
let reserved: DBClient | undefined
|
||||
|
||||
// Check if we have at least one available connection and reserve one more if required.
|
||||
if (this.available.length === 0) {
|
||||
if (this.managed || forced) {
|
||||
reserved = await this.client.reserve()
|
||||
}
|
||||
} else {
|
||||
reserved = this.available.shift() as DBClient
|
||||
}
|
||||
|
||||
try {
|
||||
// Use reserved or pool
|
||||
return await action(reserved ?? this.client)
|
||||
} catch (err: any) {
|
||||
console.error(err)
|
||||
throw err
|
||||
} finally {
|
||||
if (this.released) {
|
||||
try {
|
||||
reserved?.release()
|
||||
} catch (err: any) {
|
||||
console.error('failed to release', err)
|
||||
}
|
||||
} else if (reserved !== undefined) {
|
||||
if (this.available.length > 0) {
|
||||
reserved?.release()
|
||||
} else {
|
||||
this.available.push(reserved)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
release (): void {
|
||||
for (const c of [...this.available]) {
|
||||
c.release()
|
||||
}
|
||||
this.available = []
|
||||
}
|
||||
}
|
||||
|
||||
export class ConnectionMgr {
|
||||
private readonly connections = new Map<string, ConnectionInfo>()
|
||||
constructor (protected readonly client: DBClient) {}
|
||||
|
||||
async write (id: string | undefined, fn: (client: DBClient) => Promise<any>): Promise<void> {
|
||||
const backoffInterval = 25 // millis
|
||||
const maxTries = 5
|
||||
let tries = 0
|
||||
|
||||
const realId = id ?? `${++clId}`
|
||||
|
||||
const connection = this.getConnection(realId, false)
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const retry: boolean | Error = await connection.withReserve(async (client) => {
|
||||
tries++
|
||||
try {
|
||||
await client.execute('BEGIN;')
|
||||
await fn(client)
|
||||
await client.execute('COMMIT;')
|
||||
return true
|
||||
} catch (err: any) {
|
||||
await client.execute('ROLLBACK;')
|
||||
console.error({ message: 'failed to process tx', error: err.message, cause: err })
|
||||
|
||||
if (!this.isRetryableError(err) || tries === maxTries) {
|
||||
return err
|
||||
} else {
|
||||
console.log('Transaction failed. Retrying.')
|
||||
console.log(err.message)
|
||||
return false
|
||||
}
|
||||
}
|
||||
}, true)
|
||||
if (retry === true) {
|
||||
break
|
||||
}
|
||||
if (retry instanceof Error) {
|
||||
// Pass it to exit
|
||||
throw retry
|
||||
}
|
||||
// Retry for a timeout
|
||||
await new Promise((resolve) => setTimeout(resolve, backoffInterval))
|
||||
}
|
||||
} finally {
|
||||
if (!connection.managed) {
|
||||
// We need to relase in case it temporaty connection was used
|
||||
connection.release()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async retry (id: string | undefined, fn: (client: DBClient) => Promise<any>): Promise<any> {
|
||||
const backoffInterval = 25 // millis
|
||||
const maxTries = 5
|
||||
let tries = 0
|
||||
|
||||
const realId = id ?? `${++clId}`
|
||||
// Will reuse reserved if had and use new one if not
|
||||
const connection = this.getConnection(realId, false)
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const retry: false | { result: any } | Error = await connection.withReserve(async (client) => {
|
||||
tries++
|
||||
try {
|
||||
return { result: await fn(client) }
|
||||
} catch (err: any) {
|
||||
console.error({ message: 'failed to process sql', error: err.message, cause: err })
|
||||
if (!this.isRetryableError(err) || tries === maxTries) {
|
||||
return err
|
||||
} else {
|
||||
console.log('Read Transaction failed. Retrying.')
|
||||
console.log(err.message)
|
||||
return false
|
||||
}
|
||||
}
|
||||
})
|
||||
if (retry instanceof Error) {
|
||||
// Pass it to exit
|
||||
throw retry
|
||||
}
|
||||
if (retry === false) {
|
||||
// Retry for a timeout
|
||||
await new Promise((resolve) => setTimeout(resolve, backoffInterval))
|
||||
continue
|
||||
}
|
||||
return retry.result
|
||||
}
|
||||
} finally {
|
||||
if (!connection.managed) {
|
||||
// We need to relase in case it temporaty connection was used
|
||||
connection.release()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
release (id: string): void {
|
||||
const conn = this.connections.get(id)
|
||||
if (conn !== undefined) {
|
||||
conn.released = true
|
||||
this.connections.delete(id) // We need to delete first
|
||||
conn.release()
|
||||
}
|
||||
}
|
||||
|
||||
close (): void {
|
||||
const cnts = this.connections
|
||||
for (const [k, conn] of Array.from(cnts.entries())) {
|
||||
cnts.delete(k)
|
||||
try {
|
||||
conn.release()
|
||||
} catch (err: any) {
|
||||
console.error('failed to release connection')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
getConnection (id: string, managed: boolean = true): ConnectionInfo {
|
||||
let conn = this.connections.get(id)
|
||||
if (conn === undefined) {
|
||||
conn = new ConnectionInfo(id, this.client, managed)
|
||||
}
|
||||
if (managed) {
|
||||
this.connections.set(id, conn)
|
||||
}
|
||||
return conn
|
||||
}
|
||||
|
||||
private isRetryableError (err: any): boolean {
|
||||
const msg: string = err?.message ?? ''
|
||||
|
||||
return (
|
||||
err.code === '40001' || // Retry transaction
|
||||
err.code === '55P03' || // Lock not available
|
||||
err.code === 'CONNECTION_CLOSED' || // This error is thrown if the connection was closed without an error.
|
||||
err.code === 'CONNECTION_DESTROYED' || // This error is thrown for any queries that were pending when the timeout to sql.end({ timeout: X }) was reached. If the DB client is being closed completely retry will result in CONNECTION_ENDED which is not retried so should be fine.
|
||||
msg.includes('RETRY_SERIALIZABLE')
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user