|
1 | | -import { types, Client, ClientConfig, Pool } from 'pg' |
| 1 | +import { types, Pool, PoolConfig } from 'pg' |
2 | 2 | import { PostgresMetaResult } from './types' |
3 | 3 |
|
4 | 4 | types.setTypeParser(20, parseInt) |
5 | 5 |
|
6 | | -export const init = ( |
7 | | - config: ClientConfig, |
8 | | - { pooled = true } = {} |
9 | | -): ((sql: string) => Promise<PostgresMetaResult<any>>) => { |
10 | | - const client = pooled ? new Pool(config) : new Client(config) |
11 | | - return async (sql: string) => { |
12 | | - try { |
13 | | - const { rows } = await client.query(sql) |
14 | | - return { data: rows, error: null } |
15 | | - } catch (e) { |
16 | | - return { data: null, error: { message: e.message } } |
17 | | - } |
| 6 | +export const init: ( |
| 7 | + config: PoolConfig |
| 8 | +) => { |
| 9 | + query: (sql: string) => Promise<PostgresMetaResult<any>> |
| 10 | + end: () => Promise<void> |
| 11 | +} = (config) => { |
| 12 | + // XXX: Race condition could happen here: one async task may be doing |
| 13 | + // `pool.end()` which invalidates the pool and subsequently all existing |
| 14 | + // handles to `query`. Normally you might only deal with one DB so you don't |
| 15 | + // need to call `pool.end()`, but since the server needs this, we make a |
| 16 | + // compromise: if we run `query` after `pool.end()` is called (i.e. pool is |
| 17 | + // `null`), we temporarily create a pool and close is right after. |
| 18 | + let pool: Pool | null = new Pool(config) |
| 19 | + return { |
| 20 | + async query(sql) { |
| 21 | + try { |
| 22 | + if (!pool) { |
| 23 | + const pool = new Pool(config) |
| 24 | + const { rows } = await pool.query(sql) |
| 25 | + await pool.end() |
| 26 | + return { data: rows, error: null } |
| 27 | + } |
| 28 | + |
| 29 | + const { rows } = await pool.query(sql) |
| 30 | + return { data: rows, error: null } |
| 31 | + } catch (e) { |
| 32 | + return { data: null, error: { message: e.message } } |
| 33 | + } |
| 34 | + }, |
| 35 | + |
| 36 | + async end() { |
| 37 | + const _pool = pool |
| 38 | + pool = null |
| 39 | + // Gracefully wait for active connections to be idle, then close all |
| 40 | + // connections in the pool. |
| 41 | + if (_pool) await _pool.end() |
| 42 | + }, |
18 | 43 | } |
19 | 44 | } |
0 commit comments