@clickup/ent-framework
Version:
A PostgreSQL graph-database-alike library with microsharding and row-level security
547 lines • 25.4 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
var _a, _b;
Object.defineProperty(exports, "__esModule", { value: true });
exports.PgClientPool = exports.PgClient = void 0;
const defaults_1 = __importDefault(require("lodash/defaults"));
const range_1 = __importDefault(require("lodash/range"));
const pg_1 = __importDefault(require("pg"));
const Client_1 = require("../abstract/Client");
const ClientError_1 = require("../abstract/ClientError");
const misc_1 = require("../abstract/internal/misc");
const TimelineManager_1 = require("../abstract/TimelineManager");
const misc_2 = require("../internal/misc");
const Ref_1 = require("../internal/Ref");
const escapeLiteral_1 = require("./helpers/escapeLiteral");
const buildHintQueries_1 = require("./internal/buildHintQueries");
const misc_3 = require("./internal/misc");
const parseLsn_1 = require("./internal/parseLsn");
const PgError_1 = require("./PgError");
/**
* An abstract PostgreSQL Client. Includes connection pooling logic.
*
* Since the class is cloneable internally (using the prototype substitution
* technique, see withShard()), the contract of this class is that ALL its
* derived classes may only have readonly immediate properties. Use Ref helper
* if you need some mutable properties.
*/
class PgClient extends (_b = Client_1.Client) {
/**
* Calls swallowedErrorLogger() doing some preliminary amendment.
*/
logSwallowedError(props) {
if (!this.ended.current) {
super.logSwallowedError(props);
}
}
/**
* Initializes an instance of PgClient.
*/
constructor(options) {
super(options);
/** PG named connection pools to use. The default pool has `null` key.*/
this.pools = new Map();
/** Prewarming periodic timer (if scheduled). */
this.prewarmTimeout = new Ref_1.Ref(null);
/** Whether the pool has been ended and is not usable anymore. */
this.ended = new Ref_1.Ref(false);
/** This value is non-null if there was an unsuccessful connection attempt
* (i.e. the PG is down), and there were no successful queries since then. */
this.reportedConnectionIssue = new Ref_1.Ref(null);
/** Name of the shard associated to this Client. */
this.shardName = "public";
this.options = (0, defaults_1.default)({}, options, this.options, {
maxReplicationLagMs: options.role !== "unknown"
? 2000 // e.g. AWS Aurora, assuming it always "catches up" fast
: undefined,
}, _a.DEFAULT_OPTIONS);
this.reportedRoleAfterLastQuery = new Ref_1.Ref(this.options.role);
this.timelineManager = new TimelineManager_1.TimelineManager(this.options.maxReplicationLagMs, this.options.replicaTimelinePosRefreshMs, async () => {
const startTime = performance.now();
try {
await this.query({
query: [`SELECT '${misc_1.OP_TIMELINE_POS_REFRESH}'`],
isWrite: false,
annotations: [],
op: misc_1.OP_TIMELINE_POS_REFRESH,
table: "pg_catalog",
});
}
catch (error) {
this.logSwallowedError({
where: misc_1.OP_TIMELINE_POS_REFRESH,
error,
elapsed: Math.round(performance.now() - startTime),
importance: "normal",
});
}
});
}
/**
* Represents the full destination address this Client is working with.
* Depending on the implementation, it may include hostname, port number,
* database name, shard name etc. It is required that the address is stable
* enough to be able to cache some destination database related metadata (e.g.
* shardNos) based on that address.
*/
address() {
const { host, port, database } = this.options.config;
return (host +
(port ? `:${port}` : "") +
(database ? `/${database}` : "") +
"#" +
this.shardName);
}
/**
* Gracefully closes all the connections of this Client to let the caller
* destroy it. The pending queries are awaited to finish before returning. The
* Client becomes unusable right after calling this method (even before the
* connections are drained): you should not send queries to it.
*/
async end() {
if (this.ended.current) {
return;
}
this.ended.current = true;
clearTimeout(this.prewarmTimeout.current ?? undefined);
this.prewarmTimeout.current = null;
await (0, misc_2.mapJoin)([...this.pools.values()], async (pool) => pool.end());
}
/**
* Returns true if the Client is ended and can't be used anymore.
*/
isEnded() {
return this.ended.current;
}
/**
* Returns all Shard numbers discoverable via the connection to the Client's
* database.
*/
async shardNos() {
const shardNamer = this.options.shardNamer;
// An installation without sharding enabled.
if (!shardNamer) {
return [0];
}
// e.g. sh0000, sh0123 and not e.g. sh1 or sh12345678
const rows = await this.query({
query: [(0, misc_2.maybeCall)(shardNamer.options.discoverQuery)],
isWrite: false,
annotations: [],
op: misc_1.OP_SHARD_NOS,
table: "pg_catalog",
});
return rows
.map((row) => Object.values(row)[0])
.map((name) => (name ? shardNamer.shardNoByName(name) : null))
.filter((no) => no !== null)
.sort((a, b) => a - b);
}
/**
* Sends a read or write test query to the server. Tells the server to sit and
* wait for at least the provided number of milliseconds.
*/
async ping({ execTimeMs, isWrite, annotation, }) {
await this.query({
query: [
"DO $$ BEGIN PERFORM pg_sleep(?); IF pg_is_in_recovery() AND ? THEN RAISE read_only_sql_transaction; END IF; END $$",
execTimeMs / 1000,
isWrite,
],
isWrite,
annotations: [annotation],
op: misc_1.OP_PING,
table: "pg_catalog",
});
}
/**
* Creates a new Client which is namespaced to the provided Shard number. The
* new Client will share the same connection pool with the parent's Client.
*/
withShard(no) {
return Object.assign(Object.create(this.constructor.prototype), {
...this,
shardName: this.options.shardNamer
? this.options.shardNamer.shardNameByNo(no)
: this.shardName,
// Notice that we can ONLY have readonly properties in this and all
// derived classes to make it work. If we need some mutable props shared
// across all of the clones, we need to wrap them in a Ref (and make the
// Ref object itself readonly). That's a pretty fragile contract though.
});
}
/**
* Returns the Client's role reported after the last successful query. Master
* and replica roles may switch online unpredictably, without reconnecting, so
* we only know the role after a query.
*/
role() {
return this.reportedRoleAfterLastQuery.current;
}
/**
* Returns a non-nullable value if the Client couldn't connect to the server
* (or it could, but the load balancer reported the remote server as not
* working), so it should ideally be removed from the list of active replicas
* until e.g. the next discovery query to it (or any query) succeeds.
*/
connectionIssue() {
return this.reportedConnectionIssue.current;
}
/**
* A convenience method to put connections prewarming logic to. The idea is to
* keep the needed number of open connections and also, in each connection,
* minimize the time which the very 1st query will take (e.g. pre-cache
* full-text dictionaries).
*/
prewarm() {
if (this.prewarmTimeout.current) {
// Already scheduled a prewarm, so skipping.
return;
}
const subPools = this.options.prewarmSubPools
? [...this.pools.entries()]
.filter(([name]) => name !== null)
.map(([_, pool]) => pool)
: [];
for (const pool of [this.pool(), ...subPools]) {
const config = pool.options;
if (!config.min) {
continue;
}
const min = Math.min(config.min, config.max ?? Infinity, pool.totalCount + ((0, misc_2.maybeCall)(this.options.prewarmIntervalStep) || 1));
const toPrewarm = min - pool.waitingCount;
if (toPrewarm > 0) {
const startTime = performance.now();
(0, range_1.default)(toPrewarm).forEach(() => (0, misc_2.runInVoid)(pool.query((0, misc_2.maybeCall)(this.options.prewarmQuery)).catch((error) => this.logSwallowedError({
where: `${this.constructor.name}.prewarm`,
error,
elapsed: Math.round(performance.now() - startTime),
importance: "normal",
}))));
}
}
this.prewarmTimeout.current = setTimeout(() => {
this.prewarmTimeout.current = null;
this.prewarm();
}, Math.round((0, misc_2.maybeCall)(this.options.prewarmIntervalMs) *
(0, misc_2.jitter)((0, misc_2.maybeCall)(this.options.prewarmIntervalJitter)))).unref();
}
/**
* Returns a default pool (when subPoolConfig is not passed), or a "sub-pool"
* (a named low-level Pool implementation compatible to node-postgres). The
* method is specific to the current class and is not a part of
* database-agnostic Client API.
* - Sub-pools are lazily created and memoized by the provided name. They may
* differ by config options (like statement_timeout or max connections).
* - Sub-pools inherit the properties from default PgClientOptions.config.
* - It is implied (but not enforced) that all sub-pools use the same physical
* database, because otherwise it makes not a lot of sense.
*/
pool(subPoolConfig) {
let pool = this.pools.get(subPoolConfig?.name ?? null);
if (pool) {
return pool;
}
pool = this.options
.createPool((0, defaults_1.default)({}, subPoolConfig, this.options.config, {
allowExitOnIdle: true,
}))
.on("connect", (poolClient) => {
// Called only once, after the connection is 1st created.
const client = poolClient;
// Initialize additional properties merged into the default PoolClient.
const maxConnLifetimeMs = (0, misc_2.maybeCall)(this.options.maxConnLifetimeMs) *
(0, misc_2.jitter)((0, misc_2.maybeCall)(this.options.maxConnLifetimeJitter));
client.pool = pool;
client.id = connId++;
client.queriesSent = 0;
client.closeAt =
maxConnLifetimeMs > 0
? Date.now() + Math.round(maxConnLifetimeMs)
: null;
// Sets a "default error" handler to not let errors leak to e.g. Jest
// and the outside world as "unhandled error". Appending an additional
// error handler to EventEmitter doesn't affect the existing error
// handlers anyhow, so should be safe.
client.on("error", () => { });
})
.on("error", (error) =>
// Having this hook prevents node from crashing.
this.logSwallowedError({
where: 'Pool.on("error")',
error,
elapsed: null,
importance: "low",
}));
this.pools.set(subPoolConfig?.name ?? null, pool);
return pool;
}
/**
* Called when the Client needs a connection in the default pool (when
* subPoolConfig is not passed), or in a sub-pool (see pool() method) to run a
* query against. Implies than the caller MUST call release() method on the
* returned object. The difference from pool().connect() is that when calling
* release() on a result of acquireConn(), it additionally closes the
* connection automatically if was OPENED (not queried!) more than
* maxConnLifetimeMs ago (node-postgres Pool doesn't have this feature) The
* method is specific to the current class and is not a part of
* database-agnostic Client API.
*/
async acquireConn(subPoolConfig) {
const pool = this.pool(subPoolConfig);
const conn = (await pool.connect());
const connReleaseOrig = conn.release.bind(conn);
conn.release = (arg) => {
// Manage maxConnLifetimeMs manually since it's not supported by the
// vanilla node-postgres.
const needClose = !!(conn.closeAt && Date.now() > conn.closeAt);
return connReleaseOrig(arg !== undefined ? arg : needClose);
};
return conn;
}
/**
* Sends a query (internally, a multi-query) through the default Pool (if
* subPoolConfig is not passed), or through a named sub-pool (see pool()
* method). After the query finishes, we should expect that role() returns the
* actual master/replica role. The method is specific to the current class and
* is not a part of database-agnostic Client API.
*/
async query({ query: queryLiteral, hints, isWrite, annotations, op, table, batchFactor, subPoolConfig, }) {
const { rawPrepend, queries, queriesRollback, debugQueryWithHints, resultPos, } = this.buildMultiQuery(hints, queryLiteral, this.options.role === "unknown"
? // For master, we read its WAL LSN (pg_current_wal_insert_lsn) after
// each query (notice that, when run on a replica,
// pg_current_wal_insert_lsn() throws, so we call it only if
// pg_is_in_recovery() returns false). For replica, we read its WAL
// LSN (pg_last_wal_replay_lsn).
"SELECT CASE WHEN pg_is_in_recovery() THEN NULL ELSE pg_current_wal_insert_lsn() END AS pg_current_wal_insert_lsn, pg_last_wal_replay_lsn()"
: undefined, isWrite);
const startTime = performance.now();
let queryTime = undefined;
let conn = undefined;
let res = undefined;
let e = undefined;
let postAction = "fail";
try {
if (this.isEnded()) {
throw new ClientError_1.ClientError(Error(`Cannot use ${this.constructor.name} since it's ended`), this.options.name, "choose-another-client", "data-on-server-is-unchanged", "client_is_ended");
}
conn = await this.acquireConn(subPoolConfig);
conn.queriesSent++;
queryTime = Math.round(performance.now() - startTime);
const resMulti = await this.sendMultiQuery(conn, rawPrepend, queries, queriesRollback);
this.reportedConnectionIssue.current = null;
res = resMulti[resultPos].rows;
if (this.options.role === "unknown") {
const lsns = resMulti[resMulti.length - 1].rows[0];
if (lsns.pg_current_wal_insert_lsn !== null) {
this.reportedRoleAfterLastQuery.current = "master";
this.timelineManager.setCurrentPos((0, parseLsn_1.parseLsn)(lsns.pg_current_wal_insert_lsn));
}
else if (lsns.pg_last_wal_replay_lsn !== null) {
this.reportedRoleAfterLastQuery.current = "replica";
this.timelineManager.setCurrentPos((0, parseLsn_1.parseLsn)(lsns.pg_last_wal_replay_lsn));
}
else {
throw Error("BUG: both pg_current_wal_insert_lsn() and pg_last_wal_replay_lsn() returned null");
}
}
else if (this.options.role === "master") {
this.reportedRoleAfterLastQuery.current = "master";
// In this mode, master pos is always =1 constant.
this.timelineManager.setCurrentPos(BigInt(1), true);
}
else {
this.reportedRoleAfterLastQuery.current = "replica";
// In this mode, replica pos is always =0 constant (i.e. always behind
// the master), and we solely rely on maxReplicationLagMs timeline data
// expiration in Timeline object.
this.timelineManager.setCurrentPos(BigInt(0), true);
}
return res;
}
catch (cause) {
e = cause;
if (e instanceof ClientError_1.ClientError) {
throw e;
}
// Infer ClientError which affects Client choosing logic.
for (const predicate of misc_3.CLIENT_ERROR_PREDICATES) {
const res = predicate({
code: "" + e?.code,
message: "" + e?.message,
});
if (res) {
if (!isWrite) {
// For read queries, we know for sure that the data wasn't changed.
res.kind = "data-on-server-is-unchanged";
}
postAction =
this.role() === "master"
? res.postAction.ifMaster
: res.postAction.ifReplica;
if (res.postAction.reportConnectionIssue) {
// Mark the current Client as non-healthy, so the retry logic will
// likely choose another one if available.
this.reportedConnectionIssue.current = {
timestamp: new Date(),
cause,
postAction,
kind: res.kind,
comment: res.comment,
};
}
throw new ClientError_1.ClientError(e, this.options.name, postAction, res.kind, res.abbreviation, res.comment +
(res.kind === "unknown-server-state"
? " The write might have been committed on the PG server though."
: ""));
}
}
// Only wrap the errors which PG sent to us explicitly. Those errors mean
// that there was some aborted transaction, so it's safe to retry.
if (e?.severity) {
throw new PgError_1.PgError(e, this.options.name, debugQueryWithHints);
}
// Some other error which should not trigger query retries or
// Shards/Islands rediscovery.
throw e;
}
finally {
conn?.release();
const pool = conn?.pool ?? this.pool(subPoolConfig);
const now = performance.now();
this.options.loggers?.clientQueryLogger?.({
annotations,
op,
shard: this.shardName,
table,
batchFactor: batchFactor ?? 1,
msg: debugQueryWithHints,
output: res ? res : undefined,
elapsed: {
total: Math.round(now - startTime),
acquire: queryTime !== undefined ? queryTime : Math.round(now - startTime),
},
connStats: {
id: conn ? "" + (conn.id ?? 0) : "?",
queriesSent: conn?.queriesSent ?? 0,
},
poolStats: {
totalConns: pool.totalCount,
idleConns: pool.idleCount,
queuedReqs: pool.waitingCount,
},
error: e === undefined
? undefined
: (0, misc_2.addSentenceSuffixes)(`${e}`, e?.code ? ` (${e.code})` : undefined, ` [${postAction}]`),
role: this.role(),
backend: this.options.name,
address: this.address(),
});
}
}
/**
* Prepares a PG Client multi-query from the query literal and hints.
*/
buildMultiQuery(hints, literal, epilogue, isWrite) {
const query = (0, escapeLiteral_1.escapeLiteral)(literal).trimEnd();
if (query === "") {
throw Error("Empty query passed to query()");
}
const queriesPrologue = [];
const queriesEpilogue = [];
const queriesRollback = [];
const [rawPrepend, hintQueriesDefault, hintQueries] = (0, buildHintQueries_1.buildHintQueries)(this.options.hints ? (0, misc_2.maybeCall)(this.options.hints) : undefined, hints);
// Prepend per-query hints to the prologue (if any); they will be logged.
queriesPrologue.unshift(...hintQueries);
// The query which is logged to the logging infra. For more brief messages,
// we don't log internal hints (this.hints) and search_path; see below.
const debugQueryWithHints = `${rawPrepend}/*${this.shardName}*/` +
[...queriesPrologue, query].join("; ").trim();
// Prepend internal per-Client hints to the prologue.
queriesPrologue.unshift(...hintQueriesDefault);
// We must always have "public" in search_path, because extensions are by
// default installed in "public" schema. Some extensions may expose
// operators (e.g. "citext" exposes comparison operators) which must be
// available in all Shards by default, so they should live in "public".
// (There is a way to install an extension to a particular schema, but a)
// there can be only one such schema, and b) there are problems running
// pg_dump when migrating this Shard to another machine since pg_dump
// doesn't emit CREATE EXTENSION statement when filtering by schema name).
queriesPrologue.unshift(`SET LOCAL search_path TO ${this.shardName}, public`);
if (epilogue) {
queriesEpilogue.push(epilogue);
}
// Why wrapping with BEGIN...COMMIT for write queries? See here:
// https://www.postgresql.org/message-id/20220803.163217.1789690807623885906.horikyota.ntt%40gmail.com
if (isWrite && queriesEpilogue.length > 0) {
queriesPrologue.unshift("BEGIN");
queriesRollback.unshift("ROLLBACK");
queriesEpilogue.unshift("COMMIT");
}
return {
rawPrepend,
queries: [...queriesPrologue, query, ...queriesEpilogue],
queriesRollback,
debugQueryWithHints,
resultPos: queriesPrologue.length,
};
}
/**
* Sends a multi-query to PG Client.
*
* A good and simple explanation of the protocol is here:
* https://www.postgresql.org/docs/13/protocol-flow.html. In short, we can't
* use prepared-statement-based operations even theoretically, because this
* mode doesn't support multi-queries. Also notice that TS typing is doomed
* for multi-queries:
* https://github.com/DefinitelyTyped/DefinitelyTyped/pull/33297
*/
async sendMultiQuery(conn, rawPrepend, queries, queriesRollback) {
const queriesStr = `${rawPrepend}/*${this.shardName}*/${queries.join("; ")}`;
// For multi-query, query() actually returns an array of pg.QueryResult, but
// it's not reflected in its TS typing, so patching this.
const resMulti = (await conn.query(queriesStr).catch(async (e) => {
// We must run a ROLLBACK if we used BEGIN in the queries, because
// otherwise the connection is released to the pool in "aborted
// transaction" state (see the protocol link above).
queriesRollback.length > 0 &&
(await conn.query(queriesRollback.join("; ")).catch(() => { }));
throw e;
}));
if (resMulti.length !== queries.length) {
throw Error(`Multi-query (with semicolons) is not allowed as an input to query(); got ${queriesStr}`);
}
return resMulti;
}
}
exports.PgClient = PgClient;
_a = PgClient;
/** Default values for the constructor options. */
PgClient.DEFAULT_OPTIONS = {
...Reflect.get(_b, "DEFAULT_OPTIONS", _a),
createPool: (config) => new pg_1.default.Pool(config),
maxConnLifetimeMs: 0,
maxConnLifetimeJitter: 0.5,
prewarmIntervalStep: 1,
prewarmIntervalMs: 5000,
prewarmIntervalJitter: 0.5,
prewarmQuery: 'SELECT 1 AS "prewarmQuery"',
prewarmSubPools: false,
hints: null,
role: "unknown",
maxReplicationLagMs: 60000,
replicaTimelinePosRefreshMs: 1000,
};
/**
* For backward compatibility, exposing the old name as well.
* @deprecated Use PgClient instead.
* @ignore
*/
exports.PgClientPool = PgClient;
/**
* Auto-incrementing connection number (for debugging purposes).
*/
let connId = 1;
//# sourceMappingURL=PgClient.js.map