@clickup/ent-framework
Version:
A PostgreSQL graph-database-alike library with microsharding and row-level security
411 lines • 19.1 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.PgClient = void 0;
const defaults_1 = __importDefault(require("lodash/defaults"));
const Client_1 = require("../abstract/Client");
const ClientError_1 = require("../abstract/ClientError");
const misc_1 = require("../abstract/internal/misc");
const ShardError_1 = require("../abstract/ShardError");
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 parseCompositeRow_1 = require("./internal/parseCompositeRow");
const parseLsn_1 = require("./internal/parseLsn");
const PgError_1 = require("./PgError");
/**
* An abstract PostgreSQL Client which doesn't know how to acquire an actual
* connection and send queries; these things are up to the derived classes to
* implement.
*
* The idea is that in each particular project, people may have they own classes
* derived from PgClient, in case the codebase already has some existing
* connection pooling solution. They don't have to use PgClientPool.
*
* Since the class is cloneable internally (using the prototype substitution
* technique), the contract of this class is that ALL its derived classes may
* only have readonly immediate properties.
*/
class PgClient extends (_b = Client_1.Client) {
/**
* Initializes an instance of PgClient.
*/
constructor(options) {
super(options);
/** Number of decimal digits in an ID allocated for shard number. Calculated
* dynamically based on shards.nameFormat (e.g. for "sh%04d", it will be 4
* since it expands to "sh0012"). */
this.shardNoPadLen = 0;
/** 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 === "master" || options.role === "replica"
? 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",
});
}
});
if (this.options.shards) {
this.shardNoPadLen = this.buildShardName(0).match(/(\d+)/)
? RegExp.$1.length
: 0;
if (!this.shardNoPadLen) {
throw Error("Invalid shards.nameFormat value");
}
}
}
/**
* Sends a query (internally, a multi-query). After the query finishes, we
* should expect that role() returns the actual master/replica role.
*/
async query({ query: queryLiteral, hints, isWrite, annotations, op, table, batchFactor, }) {
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();
conn.id ??= connNo++;
conn.queriesSent = (conn.queriesSent ?? 0) + 1;
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 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: this.poolStats(),
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(),
});
}
}
async shardNos() {
// An installation without sharding enabled.
if (!this.options.shards) {
return [0];
}
// e.g. sh0000, sh0123 and not e.g. sh1 or sh12345678
const rows = await this.query({
query: [(0, misc_2.maybeCall)(this.options.shards.discoverQuery)],
isWrite: false,
annotations: [],
op: misc_1.OP_SHARD_NOS,
table: "pg_catalog",
});
return rows
.map((row) => Object.values(row)[0])
.map((name) => {
const no = name?.match(/(\d+)/) ? parseInt(RegExp.$1) : null;
return no !== null && name === this.buildShardName(no) ? no : null;
})
.filter((no) => no !== null)
.sort((a, b) => a - b);
}
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",
});
}
shardNoByID(id) {
// An installation without sharding enabled.
if (!this.options.shards) {
return 0;
}
// Just a historical exception for id="1".
if (id === "1") {
return 1;
}
// Composite ID: `(100008888888,1023499999999)` - try extracting non-zero
// Shard from parts (left to right) first, and if there is none, allow shard
// zero too.
if (typeof id === "string" && id.startsWith("(") && id.endsWith(")")) {
let no = NaN;
for (const subID of (0, parseCompositeRow_1.parseCompositeRow)(id)) {
const tryNo = subID && subID.length >= this.shardNoPadLen + 1
? parseInt(subID.substring(1, this.shardNoPadLen + 1))
: NaN;
if (!isNaN(tryNo)) {
if (tryNo > 0) {
return tryNo;
}
else if (isNaN(no)) {
no = tryNo;
}
}
}
if (isNaN(no)) {
const idSafe = (0, misc_2.sanitizeIDForDebugPrinting)(id);
throw Error(`Cannot extract shard number from the composite ID ${idSafe}`);
}
return no;
}
// Plain ID.
const no = typeof id === "string" && id.length >= this.shardNoPadLen + 1
? parseInt(id.substring(1, this.shardNoPadLen + 1))
: NaN;
if (isNaN(no)) {
const idSafe = (0, misc_2.sanitizeIDForDebugPrinting)(id);
throw new ShardError_1.ShardError(`Cannot parse ID ${idSafe} to detect shard number`, this.options.name);
}
return no;
}
withShard(no) {
return Object.assign(Object.create(this.constructor.prototype), {
...this,
shardName: this.buildShardName(no),
// 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.
});
}
role() {
return this.reportedRoleAfterLastQuery.current;
}
connectionIssue() {
return this.reportedConnectionIssue.current;
}
/**
* 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;
}
/**
* Builds the schema name (aka "Shard name") by Shard number using
* `options#shards#nameFormat`.
*
* E.g. nameFormat="sh%04d" generates names like "sh0042".
*/
buildShardName(no) {
//
return this.options.shards
? this.options.shards.nameFormat.replace(/%(0?)(\d+)[sd]/, (_, zero, d) => no.toString().padStart(zero ? parseInt(d) : 0, "0"))
: this.shardName;
}
}
exports.PgClient = PgClient;
_a = PgClient;
/** Default values for the constructor options. */
PgClient.DEFAULT_OPTIONS = {
...Reflect.get(_b, "DEFAULT_OPTIONS", _a),
shards: null,
hints: null,
role: "unknown",
maxReplicationLagMs: 60000,
replicaTimelinePosRefreshMs: 1000,
};
let connNo = 1;
//# sourceMappingURL=PgClient.js.map