Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 | 23x 23x 23x 23x 23x 23x 1x 114x 564x 564x 564x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 27x 4x 4x 4x 27x 27x 22x 22x 2x 2x 2x 22x 22x 22x 22x 22x 22x 22x 22x 5x 5x 5x 5x 2x 2x 2x 5x 5x 5x 5x 5x 5x 27x 27x 4x 4x 4x 27x 5x 5x 5x 27x 114x 31x | import type { Abortable } from 'node:events'
import { tracingChannel } from 'node:diagnostics_channel'
import { linkSignals } from '@ydbjs/abortable'
import { StatusIds_StatusCode } from '@ydbjs/api/operation'
import { QueryServiceDefinition } from '@ydbjs/api/query'
import type { Driver, DriverIdentity } from '@ydbjs/core'
import { CommitError, YDBError } from '@ydbjs/error'
import {
type RetryConfig,
type RetryContext,
defaultRetryConfig,
isRetryableError,
retry,
} from '@ydbjs/retry'
import { loggers } from '@ydbjs/debug'
import { Query } from './query.js'
import { ctx } from './ctx.js'
import { Fragment, UnsafeString, fragment, identifier, join, unsafe, yql } from './yql.js'
import { SessionPool, type SessionPoolOptions, sessionAcquireCh } from './session-pool.js'
type TransactionContext = {
driver: DriverIdentity
isolation: string
idempotent: boolean
}
type TxControlContext = {
driver: DriverIdentity
sessionId: string
nodeId: bigint
txId: string
}
type TxBeginContext = {
driver: DriverIdentity
nodeId: bigint
sessionId: string
isolation: string
}
let txBeginCh = tracingChannel<TxBeginContext>('tracing:ydb:query.begin')
let txCommitCh = tracingChannel<TxControlContext>('tracing:ydb:query.commit')
let txRollbackCh = tracingChannel<TxControlContext>('tracing:ydb:query.rollback')
let transactionCh = tracingChannel<TransactionContext>('tracing:ydb:query.transaction')
let dbg = loggers.query
export type QueryOptions = {
poolOptions?: SessionPoolOptions
}
export type SQL = <T extends any[] = unknown[], P extends any[] = unknown[]>(
strings: string | TemplateStringsArray,
...values: P
) => Query<T>
export type RegisterPrecommitHook = (fn: () => Promise<void> | void) => void
export type TX = SQL & {
nodeId: bigint
sessionId: string
transactionId: string
onRollback: (fn: (error: unknown, signal?: AbortSignal) => Promise<void> | void) => void
onCommit: (fn: (signal?: AbortSignal) => Promise<void> | void) => void
onClose: (fn: (committed: boolean, signal?: AbortSignal) => Promise<void> | void) => void
}
interface SessionContextCallback<T> {
(signal: AbortSignal): Promise<T>
}
interface TransactionExecuteOptions extends Abortable {
isolation?: 'serializableReadWrite' | 'snapshotReadOnly' | 'snapshotReadWrite'
idempotent?: boolean
}
interface TransactionContextCallback<T> {
(tx: TX, signal: AbortSignal): Promise<T>
}
export interface QueryClient extends SQL, AsyncDisposable {
// unsafe<T extends any[] = unknown[], P extends { toString(): string }[] = []>(
// strings: string | TemplateStringsArray,
// ...values: P
// ): Query<T>
do<T = unknown>(fn: SessionContextCallback<T>): Promise<T>
do<T = unknown>(options: any, fn: SessionContextCallback<T>): Promise<T>
begin<T = unknown>(fn: TransactionContextCallback<T>): Promise<T>
begin<T = unknown>(
options: TransactionExecuteOptions,
fn: TransactionContextCallback<T>
): Promise<T>
transaction<T = unknown>(fn: TransactionContextCallback<T>): Promise<T>
transaction<T = unknown>(
options: TransactionExecuteOptions,
fn: TransactionContextCallback<T>
): Promise<T>
/**
* Create an UnsafeString that represents a DB identifier (table, column).
* When used in a query, the identifier will be escaped.
*
* **WARNING: This function does not offer any protection against SQL injections,
* so you must validate any user input beforehand.**
*
* @example ```ts
* const query = sql`SELECT * FROM ${sql.identifier('my-table')}`;
* // 'SELECT * FROM `my-table`'
* ```
*/
identifier(value: string | { toString(): string }): UnsafeString
/**
* Create a composable {@link Fragment} from a tagged template. A fragment is
* not executable — splice it into another query or {@link QueryClient.join}
* to build dynamic SQL while keeping values bound as parameters.
*
* @example ```ts
* const cond = sql.fragment`${sql.identifier('age')} > ${18}`
* sql`SELECT * FROM users WHERE ${cond}`
* ```
*/
fragment(strings: TemplateStringsArray, ...values: unknown[]): Fragment
/**
* Create an UnsafeString that will be injected into the query as-is.
*
* WARNING: Use only with trusted SQL fragments (e.g., migrations) and never
* with user-provided input. Prefer parameters or {@link identifier} for
* dynamic values.
*/
unsafe(value: string | { toString(): string }): UnsafeString
/**
* Combine fragments into one, interleaving a separator. Empty list → empty
* fragment; single fragment → no separator. A string separator is inserted
* as raw SQL.
*
* @example ```ts
* sql`... WHERE ${sql.join(conditions, ' AND ')}`
* ```
*/
join(fragments: readonly Fragment[], separator?: Fragment | string): Fragment
}
let doImpl = function <T = unknown>(): Promise<T> {
throw new Error('Not implemented')
}
/**
* Creates a query client for executing YQL queries and managing transactions.
*
* @param driver - The YDB driver instance used to communicate with the database.
* @returns A `QueryClient` object that provides methods for executing queries and managing transactions.
*
* @remarks
* The returned client provides a tagged template function for YQL queries, as well as transactional helpers.
*
* @example
* ```typescript
* const client = query(driver);
* const result = await client`SELECT 1;`;
* ```
*
* @example
* ```typescript
* await client.transaction(async (yql, signal) => {
* const res = await yql`SELECT * FROM users WHERE id = ${userId}`;
* // ...
* });
* ```
*
* @see {@link QueryClient}
*/
export function query(driver: Driver, options?: QueryOptions): QueryClient {
let sessionPool = new SessionPool(driver, options?.poolOptions)
function yqlQuery<P extends any[] = unknown[], T extends any[] = unknown[]>(
strings: string | TemplateStringsArray,
...values: P
): Query<T> {
let { text, params } = yql(strings, ...values)
dbg.log('creating query instance for text: %s', text)
return ctx.run(ctx.getStore() ?? {}, () => new Query<T>(driver, text, params, sessionPool))
}
function txIml<T = unknown>(fn: TransactionContextCallback<T>): Promise<T>
function txIml<T = unknown>(
options: TransactionExecuteOptions,
fn: TransactionContextCallback<T>
): Promise<T>
/**
* Executes a transactional operation with automatic session and transaction management,
* including retries on retryable errors.
*
* This function handles the lifecycle of a YDB session and transaction, including session creation,
* transaction begin, commit, and rollback. It also manages retries for retryable errors and ensures
* proper cleanup of resources. The transaction isolation level and idempotency can be configured.
*
* @template T The return type of the transactional operation.
* @param optOrFn - Either the transaction execution options or the transactional callback function.
* If a function is provided as the first argument, it is used as the transactional callback and
* default options are applied.
* @param fn - The transactional callback function, if options are provided as the first argument.
* The callback receives a YQL query executor and an AbortSignal.
* @returns A promise that resolves with the result of the transactional operation.
* @throws {YDBError} If session creation, transaction begin, or commit fails.
* @throws {CommitError} If the transaction commit fails.
* @throws {Error} If a non-retryable error occurs during the transaction.
*
* @remarks
* - The function automatically retries the transaction on retryable errors if the operation is idempotent.
* - The session and transaction are automatically cleaned up after execution.
* - The transaction isolation level defaults to "serializableReadWrite" if not specified.
* - The function uses the driver's QueryServiceDefinition to interact with YDB.
*/
async function txIml<T = unknown>(
optOrFn: TransactionExecuteOptions | TransactionContextCallback<T>,
fn?: TransactionContextCallback<T>
): Promise<T> {
dbg.log('starting transaction')
await driver.ready()
// Snapshot the parent context once — we'll build a fresh attempt-scoped
// store inside each retry so failed attempts don't leave stale
// sessionId/transactionId in the shared store for the next one.
let parentStore = ctx.getStore() || {}
let caller = typeof optOrFn === 'function' ? optOrFn : fn
let options = typeof optOrFn === 'function' ? ({} as TransactionExecuteOptions) : optOrFn
options.isolation ??= 'serializableReadWrite'
let idempotent = options.idempotent ?? false
let sessionDiedLastAttempt = false
let acquireSession = (signal: AbortSignal) =>
sessionAcquireCh.tracePromise(() => sessionPool.acquire(signal), {
driver: driver.identity,
})
let runAttempt = async (retrySignal: AbortSignal) => {
dbg.log('acquiring session from pool for transaction')
using sessionLease = await acquireSession(retrySignal)
// `sessionLease` exists only for lifecycle (auto-release on scope
// exit). All RPC-level work — begin/commit/rollback, ctx wiring,
// nested queries inside the body — talks to the underlying Session
// directly. One source of truth, no parallel id/nodeId fields.
let session = sessionLease.session
dbg.log('session %s acquired for transaction', session.id)
// linkSignals is disposed on scope exit — no listener buildup
// on the long-lived session signal across tx retries.
using linked = linkSignals(retrySignal, sessionLease.signal)
let signal = linked.signal
let attemptStore: typeof parentStore = {
...parentStore,
signal,
session,
}
let client = driver.createClient(QueryServiceDefinition, session.nodeId)
let beginTransactionResult = await txBeginCh.tracePromise(
() =>
client.beginTransaction(
{
sessionId: session.id,
txSettings: {
txMode: { case: options.isolation!, value: {} },
},
},
{ signal }
),
{
driver: driver.identity,
nodeId: session.nodeId,
sessionId: session.id,
isolation: options.isolation!,
}
)
Iif (beginTransactionResult.status !== StatusIds_StatusCode.SUCCESS) {
dbg.log('failed to begin transaction, status: %d', beginTransactionResult.status)
throw new YDBError(beginTransactionResult.status, beginTransactionResult.issues)
}
attemptStore.transactionId = beginTransactionResult.txMeta!.id
let commitHooks: Array<(signal?: AbortSignal) => Promise<void> | void> = []
let rollbackHooks: Array<
(error: unknown, signal?: AbortSignal) => Promise<void> | void
> = []
let closeHooks: Array<
(committed: boolean, signal?: AbortSignal) => Promise<void> | void
> = []
let committed = false
try {
let tx = Object.assign(yqlQuery, {
nodeId: session.nodeId,
sessionId: session.id,
transactionId: attemptStore.transactionId,
onRollback: (fn: () => Promise<void> | void) => {
rollbackHooks.push(fn)
},
onCommit: (fn: () => Promise<void> | void) => {
commitHooks.push(fn)
},
onClose: (fn: () => Promise<void> | void) => {
closeHooks.push(fn)
},
}) as TX
dbg.log('executing transaction body')
let result = await ctx.run(attemptStore, () => caller!(tx, signal))
dbg.log('executing %d commit hooks', commitHooks.length)
await Promise.all(
commitHooks.map(async (hook, i) => {
dbg.log('executing commit hook #%d', i + 1)
await hook(signal)
dbg.log('commit hook #%d completed', i + 1)
})
)
dbg.log('committing transaction')
let txId = attemptStore.transactionId!
let commitResult = await txCommitCh.tracePromise(
() =>
client.commitTransaction(
{
sessionId: session.id,
txId,
},
{ signal }
),
{
driver: driver.identity,
sessionId: session.id,
nodeId: session.nodeId,
txId,
}
)
Iif (commitResult.status !== StatusIds_StatusCode.SUCCESS) {
dbg.log('failed to commit transaction, status: %d', commitResult.status)
throw new CommitError(
'Transaction commit failed.',
new YDBError(commitResult.status, commitResult.issues)
)
}
committed = true
dbg.log('transaction committed successfully')
return result
} catch (error) {
dbg.log('transaction error: %O', error)
// Signal up to the retry callback that this attempt tore down
// because the session died — lets it retry with a fresh one
// if the caller opted into idempotent retries.
Iif (session.signal.aborted) {
sessionDiedLastAttempt = true
}
dbg.log('executing %d rollback hooks', rollbackHooks.length)
await Promise.all(
rollbackHooks.map(async (hook, i) => {
dbg.log('executing rollback hook #%d', i + 1)
await hook(error, signal)
dbg.log('rollback hook #%d completed', i + 1)
})
)
// Fire-and-forget rollback. tracePromise sees the underlying
// promise (so `error` sub-channel fires on failure); the outer
// `.catch` keeps the rejection from becoming unhandled here.
let rollbackTxId = attemptStore.transactionId!
txRollbackCh
.tracePromise(
() =>
client.rollbackTransaction({
sessionId: session.id,
txId: rollbackTxId,
}),
{
driver: driver.identity,
sessionId: session.id,
nodeId: session.nodeId,
txId: rollbackTxId,
}
)
.catch(() => {})
Eif (!isRetryableError(error, idempotent)) {
dbg.log('transaction not retryable, aborting')
throw new Error('Transaction failed.', { cause: error })
}
throw error
} finally {
dbg.log('executing %d close hooks', closeHooks.length)
await Promise.all(
closeHooks.map(async (hook, i) => {
dbg.log('executing close hook #%d', i + 1)
await hook(committed, signal)
dbg.log('close hook #%d completed', i + 1)
})
)
}
}
let retryConfig: RetryConfig = {
...defaultRetryConfig,
signal: options.signal,
// Caller's flag flows through the retry layer unchanged and
// lands back in the callback as the second arg — single source
// of truth for whether the body is safe to replay.
idempotent,
retry: (error: unknown, idempotent: boolean) => {
let sessionDied = sessionDiedLastAttempt
sessionDiedLastAttempt = false
// A session abort mid-tx means we don't know whether the
// server applied any side effects — only re-open if the
// caller opted into idempotent retries.
return isRetryableError(error, idempotent) || (sessionDied && idempotent)
},
onRetry: (retryCtx: RetryContext) => {
dbg.log(
'retrying transaction, attempt %d, error: %O',
retryCtx.attempt,
retryCtx.error
)
},
}
return transactionCh.tracePromise(() => retry(retryConfig, runAttempt), {
driver: driver.identity,
isolation: options.isolation!,
idempotent,
})
}
return Object.assign(yqlQuery, {
do: doImpl,
begin: txIml,
transaction: txIml,
identifier: identifier,
unsafe: unsafe,
fragment: fragment,
join: join,
async [Symbol.asyncDispose]() {
await sessionPool.close()
},
})
}
export { type Query } from './query.ts'
export { identifier, unsafe, UnsafeString, fragment, join, Fragment } from './yql.js'
export { type SessionPoolOptions } from './session-pool.js'
|