@clickup/ent-framework
Version:
A PostgreSQL graph-database-alike library with microsharding and row-level security
130 lines • 6.74 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.PgQueryUpsert = void 0;
const mapValues_1 = __importDefault(require("lodash/mapValues"));
const omit_1 = __importDefault(require("lodash/omit"));
const pick_1 = __importDefault(require("lodash/pick"));
const pickBy_1 = __importDefault(require("lodash/pickBy"));
const union_1 = __importDefault(require("lodash/union"));
const misc_1 = require("../internal/misc");
const types_1 = require("../types");
const PgRunner_1 = require("./PgRunner");
class PgQueryUpsert {
constructor(schema, input) {
this.schema = schema;
this.input = input;
this.IS_WRITE = true;
}
async run(client, annotation) {
if (!this.schema.uniqueKey.length) {
throw Error(`Define unique key fields to use upsert for ${this.schema.name}`);
}
const fieldsWithExplicitValues = Object.keys(this.schema.table).filter((field) => this.input[field] !== undefined);
return client
.batcher(this.constructor, this.schema, fieldsWithExplicitValues.join(":"), false, () => new PgRunnerUpsert(this.schema, client, fieldsWithExplicitValues))
.run(this.input, annotation);
}
}
exports.PgQueryUpsert = PgQueryUpsert;
class PgRunnerUpsert extends PgRunner_1.PgRunner {
constructor(schema, client, fieldsWithExplicitValues) {
super(schema, client);
this.op = "UPSERT";
this.maxBatchSize = 100;
this.default = "never_happens"; // abstract property implementation
const table = this.schema.table;
const uniqueKey = this.schema.uniqueKey;
const allFields = this.addPK(Object.keys(table), "prepend");
// We must have at least some fields in the WITH CTE, because otherwise we
// won't be able to generate FROM rows WHERE ... clause for the top UPDATE.
fieldsWithExplicitValues = (0, union_1.default)(fieldsWithExplicitValues, uniqueKey);
const insertSelectClause = {
fields: allFields,
autos: (0, mapValues_1.default)((0, omit_1.default)((0, pick_1.default)(table, allFields), fieldsWithExplicitValues), ({ autoInsert, autoUpdate }) => autoInsert ?? autoUpdate),
};
const updateWhereClause = {
fields: uniqueKey,
autos: (0, mapValues_1.default)((0, omit_1.default)((0, pick_1.default)(table, uniqueKey), fieldsWithExplicitValues), ({ autoInsert, autoUpdate }) => autoInsert ?? autoUpdate),
};
const updateFields = (0, union_1.default)(fieldsWithExplicitValues, Object.keys((0, pickBy_1.default)(table, ({ autoUpdate }) => autoUpdate !== undefined)));
const updateSetClause = {
fields: updateFields,
autos: (0, mapValues_1.default)((0, omit_1.default)((0, pick_1.default)(table, updateFields), fieldsWithExplicitValues), ({ autoUpdate }) => autoUpdate),
};
this.builder = this.createWithBuilder({
fields: fieldsWithExplicitValues,
skipSorting: true, // THE ORDER MATTERS!!! See FRAGILE comment below.
suffix: ",\nupdates AS (\n" +
(0, misc_1.indent)(this.fmt("UPDATE %T ") +
this.fmt("SET %UPDATE_FIELD_VALUE_PAIRS(rows)\n", updateSetClause) +
this.fmt("FROM rows WHERE %WHERE_FIELD_VALUE_PAIRS(%T,rows)\n", updateWhereClause) +
this.fmt(`RETURNING rows._key, %PK(%T) AS ${types_1.ID})`)) +
",\ninserts AS (\n" +
(0, misc_1.indent)(this.fmt("INSERT INTO %T (%FIELDS)\n", { fields: allFields }) +
this.fmt("SELECT %FIELDS\n", insertSelectClause) +
"FROM rows WHERE _key NOT IN (SELECT _key FROM updates) OFFSET 1\n" +
this.fmt("ON CONFLICT (%FIELDS) DO UPDATE ", {
fields: uniqueKey,
}) +
this.fmt("SET %UPDATE_FIELD_VALUE_PAIRS(EXCLUDED)\n", updateSetClause) +
this.fmt(`RETURNING NULL AS _key, %PK AS ${types_1.ID})`)) +
`\nSELECT _key, ${types_1.ID} FROM updates UNION ALL SELECT _key, ${types_1.ID} FROM inserts`,
});
}
key(inputIn) {
const input = inputIn;
const key = [];
for (const field of this.schema.uniqueKey) {
key.push(input[field] === null || input[field] === undefined
? { guaranteed_unique_value: super.key(inputIn) }
: input[field]);
}
return JSON.stringify(key);
}
async runSingle(input, annotations) {
const sql = this.builder.prefix +
this.builder.func([["", input]]) +
this.builder.suffix;
const rows = await this.clientQuery(sql, annotations, 1);
return (0, misc_1.nullthrows)(rows[0], sql)[types_1.ID];
}
async runBatch(inputs, annotations) {
const sql = this.builder.prefix + this.builder.func(inputs) + this.builder.suffix;
const rows = await this.clientQuery(sql, annotations, inputs.size);
if (rows.length !== inputs.size) {
throw Error(`BUG: number of rows returned from upsert (${rows.length}) ` +
`is different from the number of input rows (${inputs.size}): ${sql}`);
}
const outputs = new Map();
// First, extract all top-level UPDATEd rows, we know their keys.
const inputsWithNullRowKeys = new Map(inputs);
const rowsWithNullKey = [];
for (const row of rows) {
if (row._key !== null) {
outputs.set(row._key, row[types_1.ID]);
inputsWithNullRowKeys.delete(row._key);
}
else {
rowsWithNullKey.push(row);
}
}
// FRAGILE! Then, extract INSERTed or on-conflict UPDATEd rows, we don't
// know their keys. In case insert didn't happen in "INSERT ... ON CONFLICT
// DO UPDATE ... RETURNING ..." clause, we can't match the updated row id
// with the key: one can only pull the fields of the updated table in
// RETURNING, where _key field just doesn't exist. Luckily, the order of
// rows returned is the same as the input rows order, and "ON CONFLICT DO
// UPDATE" update always succeeds entirely (or fails entirely).
let i = 0;
for (const key of inputsWithNullRowKeys.keys()) {
outputs.set(key, rowsWithNullKey[i][types_1.ID]);
i++;
}
return outputs;
}
}
PgRunnerUpsert.IS_WRITE = true;
//# sourceMappingURL=PgQueryUpsert.js.map