@ydbjs/query
Version:
High-level, type-safe YQL query and transaction client for YDB. Supports tagged template syntax, parameter binding, transactions, and statistics.
360 lines • 13.3 kB
JavaScript
import { EventEmitter } from 'node:events';
import { create } from '@bufbuild/protobuf';
import { StatusIds_StatusCode } from '@ydbjs/api/operation';
import { ExecMode, QueryServiceDefinition, StatsMode, Syntax, TransactionControlSchema, } from '@ydbjs/api/query';
import { TypedValueSchema } from '@ydbjs/api/value';
import { loggers } from '@ydbjs/debug';
import { YDBError } from '@ydbjs/error';
import { defaultRetryConfig, retry } from '@ydbjs/retry';
import { fromYdb, toJs } from '@ydbjs/value';
import { typeToString } from '@ydbjs/value/print';
import { ctx } from './ctx.js';
let dbg = loggers.query;
export class Query extends EventEmitter {
#driver;
#promise = null;
#cleanup = [];
#text;
#parameters;
#idempotent = false;
#active = false;
#disposed = false;
#cancelled = false;
#controller = new AbortController();
#signal;
#timeout;
#syntax = Syntax.YQL_V1;
#poolId;
#stats;
#statsMode = StatsMode.UNSPECIFIED;
#isolation = 'implicit';
#isolationSettings = {};
#raw = false;
#values = false;
constructor(driver, text, params) {
super();
this.#text = text;
this.#driver = driver;
this.#parameters = {};
for (let key in params) {
key.startsWith('$') || (key = '$' + key);
this.#parameters[key] = params[key];
}
}
static get [Symbol.species]() {
return Promise;
}
/* oxlint-disable max-lines-per-function */
async #execute() {
let { nodeId, sessionId, transactionId, signal } = ctx.getStore() || {};
if (this.#disposed) {
dbg.log('query disposed, aborting execution');
throw new Error('Query has been disposed.');
}
// If we already have a promise, return it without executing the query again
if (this.#promise) {
dbg.log('query already has a promise, returning cached result');
return this.#promise;
}
if (this.#active) {
dbg.log('query is already executing, aborting');
throw new Error('Query is already executing.');
}
dbg.log('starting query execution: %s', this.text);
this.#active = true;
signal = signal ? AbortSignal.any([signal, this.#controller.signal]) : this.#controller.signal;
if (this.#signal) {
signal = AbortSignal.any([signal, this.#signal]);
}
if (this.#timeout) {
signal = AbortSignal.any([signal, AbortSignal.timeout(this.#timeout)]);
}
let retryConfig = {
...defaultRetryConfig,
signal,
idempotent: this.#idempotent,
onRetry: (retryCtx) => {
dbg.log('retrying query, attempt %d, error: %O', retryCtx.attempt, retryCtx.error);
this.emit('retry', retryCtx);
}
};
await this.#driver.ready(signal);
this.#promise = retry(retryConfig, async (signal) => {
dbg.log('creating query client for nodeId: %s', nodeId);
let client = this.#driver.createClient(QueryServiceDefinition, nodeId);
if (!sessionId) {
dbg.log('creating new session');
let sessionResponse = await client.createSession({}, { signal });
if (sessionResponse.status !== StatusIds_StatusCode.SUCCESS) {
dbg.log('failed to create session, status: %d', sessionResponse.status);
throw new YDBError(sessionResponse.status, sessionResponse.issues);
}
nodeId = sessionResponse.nodeId;
sessionId = sessionResponse.sessionId;
client = this.#driver.createClient(QueryServiceDefinition, nodeId);
this.#cleanup.push(async () => {
dbg.log('deleting session %s', sessionId);
await client.deleteSession({ sessionId: sessionId });
});
let attachSession = client.attachSession({ sessionId }, { signal })[Symbol.asyncIterator]();
let attachSessionResult = await attachSession.next();
if (attachSessionResult.value.status !== StatusIds_StatusCode.SUCCESS) {
dbg.log('failed to attach session, status: %d', attachSessionResult.value.status);
throw new YDBError(attachSessionResult.value.status, attachSessionResult.value.issues);
}
dbg.log('session %s created and attached', sessionId);
}
let parameters = {};
for (let key in this.#parameters) {
parameters[key] = create(TypedValueSchema, {
type: this.#parameters[key].type.encode(),
value: this.#parameters[key].encode(),
});
}
// If we have a transactionId, we should use it
// If we have an isolation level, we should use it
let txControl;
if (transactionId) {
txControl = create(TransactionControlSchema, {
txSelector: {
case: "txId",
value: transactionId,
},
});
}
else if (this.#isolation !== "implicit") {
txControl = create(TransactionControlSchema, {
commitTx: true,
txSelector: {
case: "beginTx",
value: {
txMode: {
case: this.#isolation,
value: this.#isolationSettings,
}
}
}
});
}
let stream = client.executeQuery({
sessionId,
execMode: ExecMode.EXECUTE,
query: {
case: 'queryContent',
value: {
syntax: this.#syntax,
text: this.text,
},
},
parameters,
statsMode: this.#statsMode,
...(this.#poolId && { poolId: this.#poolId }),
...(txControl && { txControl }),
}, {
signal,
onTrailer: (trailer) => {
this.emit('metadata', trailer);
},
});
let results = [];
for await (let part of stream) {
signal.throwIfAborted();
if (part.status !== StatusIds_StatusCode.SUCCESS) {
dbg.log('query part failed, status: %d', part.status);
throw new YDBError(part.status, part.issues);
}
if (part.execStats) {
dbg.log('received query stats');
this.#stats = part.execStats;
}
if (!part.resultSet) {
continue;
}
while (part.resultSetIndex >= results.length) {
results.push([]);
}
for (let i = 0; i < part.resultSet.rows.length; i++) {
let result = this.#values ? [] : {};
for (let j = 0; j < part.resultSet.columns.length; j++) {
let column = part.resultSet.columns[j];
let value = part.resultSet.rows[i].items[j];
if (this.#values) {
result.push(this.#raw ? value : toJs(fromYdb(value, column.type)));
continue;
}
result[column.name] = this.#raw ? value : toJs(fromYdb(value, column.type));
}
results[Number(part.resultSetIndex)].push(result);
}
}
dbg.log('query executed successfully');
return results;
})
.then((results) => {
if (this.#stats) {
this.emit('stats', this.#stats);
}
this.emit('done', results);
return results;
})
.catch((err) => {
this.emit('error', err);
throw err;
})
.finally(async () => {
this.#active = false;
this.#controller.abort('Query completed.');
this.#cleanup.forEach((fn) => void fn());
this.#cleanup = [];
});
return this.#promise;
}
/** Returns the result of the query */
/* oxlint-disable unicorn/no-thenable */
then(onfulfilled, onrejected) {
return this.#execute().then(onfulfilled, onrejected);
}
/** Indicates if the query is currently executing */
get active() {
return this.#active;
}
/** Indicates if the query has been cancelled */
get cancelled() {
return this.#cancelled;
}
get text() {
let queryText = this.#text;
if (this.#parameters) {
for (let [name, value] of Object.entries(this.#parameters)) {
name.startsWith('$') || (name = '$' + name);
queryText = `DECLARE ${name} AS ${typeToString(value.type)};\n` + queryText;
}
}
return queryText;
}
get parameters() {
return this.#parameters;
}
syntax(syntax) {
this.#syntax = syntax;
return this;
}
pool(poolId) {
this.#poolId = poolId;
return this;
}
/** Adds a parameter to the query */
parameter(name, parameter) {
name.startsWith('$') || (name = '$' + name);
if (parameter === undefined) {
delete this.#parameters[name];
return this;
}
this.#parameters[name] = parameter;
return this;
}
/** Adds a parameter to the query */
param(name, parameter) {
return this.parameter(name, parameter);
}
/**
* Sets the idempotent flag for the query.
*
* ONLY FOR SINGLE EXECUTE CALLS.
* DO NOTHING IN TRANSACTION CONTEXT (sql.begin or sql.transaction).
*
* Idempotent queries may be retried without side effects.
*/
idempotent(idempotent = true) {
this.#idempotent = idempotent;
return this;
}
/**
* Sets the transaction isolation level for a single execute call.
*
* ONLY FOR SINGLE EXECUTE CALLS.
* DO NOTHING IN TRANSACTION CONTEXT (sql.begin or sql.transaction).
*
* A transaction is always used. If `mode` is 'implicit', the database decides the isolation level.
* If a specific isolation `mode` is provided, the query will be executed within a single transaction (with inline begin and commit)
* using the specified isolation level.
*
* @param mode Transaction isolation level:
* - 'serializableReadWrite' — serializable read/write
* - 'snapshotReadOnly' — snapshot read-only
* - 'onlineReadOnly' — online read-only
* - 'staleReadOnly' — stale read-only
* - 'implicit' — isolation is not set, server decides
* - 'implicit' is the default value
* @param settings Additional options, e.g., allowInconsistentReads — allow inconsistent reads only with 'onlineReadOnly'
* @returns The current instance for chaining
*/
isolation(mode, settings = {}) {
this.#isolation = mode;
this.#isolationSettings = settings;
return this;
}
/** Returns the query execution statistics */
// TODO: Return user-friendly stats report
stats() {
return this.#stats;
}
/** Returns a query with statistics enabled */
withStats(mode) {
this.#statsMode = mode;
return this;
}
/** Sets the query timeout */
timeout(timeout) {
this.#timeout = timeout;
return this;
}
/** Cancels the executing query */
cancel() {
this.#controller.abort('Query cancelled by user.');
this.#cancelled = true;
this.emit('cancel');
return this;
}
signal(signal) {
this.#signal = signal;
return this;
}
/** Executes the query */
execute() {
void this.#execute();
return this;
}
/** Returns only the values from the query result */
values() {
this.#values = true;
return this;
}
/** Returns raw values */
raw() {
this.#raw = true;
return this;
}
/**
* Disposes the query and releases all resources.
* This method is called automatically when the query is done.
* It is recommended to call this method explicitly when the query is no longer needed.
*/
async dispose() {
if (this.#disposed) {
return;
}
this.#controller.abort('Query disposed.');
await Promise.all(this.#cleanup.map((fn) => fn()));
this.#cleanup = [];
this.#promise = null;
this.#disposed = true;
}
[Symbol.dispose]() {
this.dispose();
}
async [Symbol.asyncDispose]() {
await this.dispose();
}
}
//# sourceMappingURL=query.js.map