From ce31c2f6e889c8ef4ce39545de2d6d9f5ce3c56e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 10 Sep 2026 10:00:53 +0200 Subject: [PATCH] perf(spanner): bypass stream pipeline for single-chunk query results Historically, database.run and transaction.run route all queries through a multi-stage Transform stream pipeline (Readable -> CheckpointStream -> PartialResultStream -> Transform). For small, single-chunk queries, stream state buffering, event emitter dispatch, and microtask scheduling introduce significant CPU and latency overhead. This change introduces an internal fast path for single-chunk queries: - For Database.run (with multiplexed sessions) and Snapshot.run, executes the gRPC request directly without wrapping it in a Transform pipeline. - If the query completes in a single response chunk (chunk.last is true and chunkedValue is unset), rows are decoded directly in a tight synchronous loop (decodeRowsDirect) into pre-allocated row arrays or plain JSON objects. - If the result spans multiple chunks, it seamlessly falls back to the full partialResultStream pipeline by replaying the first chunk through a pass-through stream with zero data loss or token mismatch. - PartialResultStream also utilizes decodeRowsDirect when streaming queries arrive in a single chunk. - Preserves full OpenTelemetry trace span hierarchy and Cloud Spanner transaction retry semantics. --- handwritten/spanner/src/database.ts | 160 ++- .../spanner/src/partial-result-stream.ts | 286 +++-- handwritten/spanner/src/transaction.ts | 457 ++++++- handwritten/spanner/test/database.ts | 208 ++++ .../spanner/test/partial-result-stream.ts | 250 +++- handwritten/spanner/test/transaction.ts | 1104 ++++++++++++++++- 6 files changed, 2361 insertions(+), 104 deletions(-) diff --git a/handwritten/spanner/src/database.ts b/handwritten/spanner/src/database.ts index 23904f4abab..b6ba8347e4c 100644 --- a/handwritten/spanner/src/database.ts +++ b/handwritten/spanner/src/database.ts @@ -73,6 +73,7 @@ import { RunCallback, RunResponse, RunUpdateCallback, + Rows, Snapshot, TimestampBounds, Transaction, @@ -2891,9 +2892,6 @@ class Database extends common.GrpcServiceObject { optionsOrCallback?: TimestampBounds | RunCallback, cb?: RunCallback, ): void | Promise { - let stats: ResultSetStats; - let metadata: ResultSetMetadata; - const rows: Row[] = []; const callback = typeof optionsOrCallback === 'function' ? (optionsOrCallback as RunCallback) @@ -2903,7 +2901,33 @@ class Database extends common.GrpcServiceObject { ? (optionsOrCallback as TimestampBounds) : {}; - return startTrace( + if ( + this.runStream !== Database.prototype.runStream || + !this.sessionFactory_.isMultiplexedEnabled() + ) { + this._runLegacy(query, options, callback!); + return; + } + this._run(query, options, callback!); + } + + /** + * Always runs the query through the full streaming pipeline (Database.prototype.runStream). + * Used when runStream has been overridden on Database, or when multiplexed + * sessions are disabled (requiring standard session pool checkout and session-not-found retries). + * + * @private + */ + private _runLegacy( + query: string | ExecuteSqlRequest, + options: TimestampBounds, + callback: RunCallback, + ): void { + const rows: Row[] = []; + let stats: ResultSetStats; + let metadata: ResultSetMetadata; + + startTrace( 'Database.run', { ...(query as ExecuteSqlRequest), @@ -2932,6 +2956,134 @@ class Database extends common.GrpcServiceObject { }, ); } + + /** + * Executes a query using an optimized streaming model: + * 1. If all results are returned in a single PartialResultSet, the internal streaming + * pipeline is simplified and rows are decoded directly, bypassing Stream overhead. + * 2. If there are more than one PartialResultSets, it seamlessly falls back to + * the standard streaming model. + * + * @private + */ + private _run( + query: string | ExecuteSqlRequest, + options: TimestampBounds, + callback: RunCallback, + ): void { + const traceConfig = { + ...(query as ExecuteSqlRequest), + ...this._traceConfig, + }; + + startTrace('Database.run', traceConfig, runSpan => { + startTrace('Database.runStream', traceConfig, streamSpan => { + this._executeRunOnSession( + query, + options, + runSpan, + streamSpan, + callback, + ); + }); + }); + } + + /** + * Acquires a session and executes the query on a snapshot, managing span + * lifetimes and session release. + * + * @private + */ + private _executeRunOnSession( + query: string | ExecuteSqlRequest, + options: TimestampBounds, + runSpan: Span, + streamSpan: Span, + callback: RunCallback, + ): void { + let snapshot: Snapshot | undefined; + let completed = false; + + const complete = ( + error: grpc.ServiceError | null, + rows: Rows = [], + stats?: ResultSetStats, + metadata?: ResultSetMetadata, + ) => { + if (completed) { + return; + } + completed = true; + if (error) { + setSpanError(streamSpan, error as Error); + setSpanError(runSpan, error as Error); + } + snapshot?.end(); + streamSpan.end(); + runSpan.end(); + callback!(error, rows, stats!, metadata!); + }; + + this.sessionFactory_.getSession((error, session) => { + if (error) { + complete(error as grpc.ServiceError); + return; + } + + streamSpan.addEvent('Using Session', {'session.id': session?.id}); + snapshot = session!.snapshot(options, this.queryOptions_); + this._runOnSnapshot(snapshot, session!, query, complete); + }); + } + + /** + * Executes the query on the snapshot and binds session release to snapshot end. + * + * @private + */ + private _runOnSnapshot( + snapshot: Snapshot, + session: Session, + query: string | ExecuteSqlRequest, + callback: ( + error: grpc.ServiceError | null, + rows?: Rows, + stats?: ResultSetStats, + metadata?: ResultSetMetadata, + ) => void, + ): void { + snapshot.once('end', () => { + try { + this.sessionFactory_.release(session); + } catch (releaseError) { + this.emit('error', releaseError); + } + }); + + const snapshotWithRun = snapshot as Snapshot & { + _run?: ( + query: string | ExecuteSqlRequest, + callback: RunCallback, + options?: {startRunSpan?: boolean}, + ) => void; + }; + + try { + if ( + typeof snapshotWithRun._run === 'function' && + snapshot.runStream === Snapshot.prototype.runStream + ) { + snapshotWithRun._run(query, callback as RunCallback, { + startRunSpan: false, + }); + } else { + snapshot.run(query, callback as RunCallback); + } + } catch (syncError) { + callback(syncError as grpc.ServiceError); + } + } /** * Partitioned DML transactions are used to execute DML statements with a * different execution strategy that provides different, and often better, diff --git a/handwritten/spanner/src/partial-result-stream.ts b/handwritten/spanner/src/partial-result-stream.ts index 970afd84c65..8b5b03e4052 100644 --- a/handwritten/spanner/src/partial-result-stream.ts +++ b/handwritten/spanner/src/partial-result-stream.ts @@ -137,6 +137,170 @@ Object.defineProperty(RowImpl.prototype, 'constructor', { enumerable: false, }); +/** + * Creates an array of decoder functions for the specified struct fields. + * + * @private + */ +export function createFieldDecoders( + fields: google.spanner.v1.StructType.Field[], + options?: RowOptions, +): Function[] { + const jsonMode = Boolean(options?.json); + const jsonOptions = options?.jsonOptions; + const columnsMetadata = options?.columnsMetadata; + + return fields.map(({name, type}) => { + const columnMetadata = + columnsMetadata && + name !== null && + name !== undefined && + Object.prototype.hasOwnProperty.call(columnsMetadata, name) + ? (columnsMetadata as Record)[name] + : undefined; + if (codec.decode !== originalDecode) { + return (val: Value) => + codec.decode(val, type as google.spanner.v1.Type, columnMetadata); + } + return codec.getDecoder( + type as google.spanner.v1.Type, + columnMetadata, + jsonMode ? jsonOptions || {} : undefined, + ); + }); +} + +/** + * Directly creates a plain JSON object from row cell values, bypassing + * Struct, Row, and WrappedNumber class wrappers when possible. + * + * @private + */ +export function createJsonRow( + fields: google.spanner.v1.StructType.Field[], + decoders: Function[], + values: Value[], + includeNameless?: boolean, +): Json { + const json: Json = {}; + const len = fields.length; + + for (let i = 0; i < len; i++) { + const {name} = fields[i]; + if (!name && !includeNameless) { + continue; + } + const fieldName = name || `_${i}`; + try { + json[fieldName] = decoders[i](values[i]); + } catch (e) { + (e as Error).message = [ + `Serializing column "${fieldName}" encountered an error: ${ + (e as Error).message + }`, + 'Call row.toJSON({ wrapNumbers: true }) to receive a custom type.', + ].join(' '); + throw e; + } + } + return json; +} + +/** + * Converts an array of decoded cell values into a Row. + * + * @private + */ +export function createRow( + fields: google.spanner.v1.StructType.Field[], + decoders: Function[], + values: Value[], +): Row { + const len = fields.length; + const row = new RowImpl(len); + for (let i = 0; i < len; i++) { + row[i] = { + name: fields[i].name, + value: decoders[i](values[i]), + }; + } + return row; +} + +/** + * Formats raw row values into either a plain JSON object or a Row instance + * according to the provided RowOptions. + * + * @private + */ +export function formatRow( + fields: google.spanner.v1.StructType.Field[], + decoders: Function[], + values: Value[], + options?: RowOptions, +): Row { + const jsonMode = Boolean(options?.json); + const jsonOptions = options?.jsonOptions; + const isJsonStubbed = + codec.convertFieldsToJson !== originalConvertFieldsToJson; + + if (jsonMode && !isJsonStubbed) { + return createJsonRow( + fields, + decoders, + values, + Boolean(jsonOptions?.includeNameless), + ) as unknown as Row; + } + + const row = createRow(fields, decoders, values); + return jsonMode ? (row.toJSON(jsonOptions) as unknown as Row) : row; +} + +/** + * Directly decodes rows from a PartialResultSet without going through the stream + * pipeline. Used by the fast-path for queries returning small results in a single chunk. + * + * @private + */ +export function decodeRowsDirect( + chunk: google.spanner.v1.PartialResultSet, + options?: RowOptions, + existingFields?: google.spanner.v1.StructType.Field[], + existingDecoders?: Function[], +): Row[] { + const fields = + existingFields || + ((chunk.metadata?.rowType?.fields || + []) as google.spanner.v1.StructType.Field[]); + const numFields = fields.length; + const chunkValues = chunk.values || []; + const numValues = chunkValues.length; + if (numFields === 0 || numValues === 0) { + return []; + } + + const rowCount = Math.floor(numValues / numFields); + const rows: Row[] = new Array(rowCount); + + const decoders: Function[] = + existingDecoders || createFieldDecoders(fields, options); + + const rowValues: Value[] = new Array(numFields); + + for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) { + const offset = rowIndex * numFields; + for (let columnIndex = 0; columnIndex < numFields; columnIndex++) { + rowValues[columnIndex] = GrpcService.decodeValue_( + chunkValues[offset + columnIndex], + ); + } + rows[rowIndex] = formatRow(fields, decoders, rowValues, options); + } + + return rows; +} + /** * @callback PartialResultStream~rowCallback * @param {Row|object} row The row data. @@ -205,12 +369,14 @@ export class PartialResultStream extends Transform implements ResultEvents { private _pendingValueForResume?: p.IValue; private _values: p.IValue[]; private _numPushFailed = 0; + private _isFirstChunk = true; constructor(options = {}) { super({objectMode: true}); this._destroyed = false; this._options = Object.assign({maxResumeRetries: 20}, options); this._values = []; + this._isFirstChunk = true; } /** * Destroys the stream. @@ -255,28 +421,7 @@ export class PartialResultStream extends Transform implements ResultEvents { if (!this._fields && chunk.metadata) { this._fields = chunk.metadata.rowType! .fields as google.spanner.v1.StructType.Field[]; - - this._decoders = this._fields.map(({name, type}) => { - const columnMetadata = - this._options.columnsMetadata && - name !== null && - name !== undefined && - Object.prototype.hasOwnProperty.call( - this._options.columnsMetadata, - name, - ) - ? (this._options.columnsMetadata as any)[name] - : undefined; - if (codec.decode !== originalDecode) { - return val => - codec.decode(val, type as google.spanner.v1.Type, columnMetadata); - } - return codec.getDecoder( - type as google.spanner.v1.Type, - columnMetadata, - this._options.json ? this._options.jsonOptions || {} : undefined, - ); - }); + this._decoders = createFieldDecoders(this._fields, this._options); } let res = true; @@ -356,6 +501,14 @@ export class PartialResultStream extends Transform implements ResultEvents { * @param {object} chunk The partial result set. */ private _addChunk(chunk: google.spanner.v1.PartialResultSet): boolean { + const isFirstChunk = this._isFirstChunk; + this._isFirstChunk = false; + + // Fast path for single-chunk stream responses: + if (isFirstChunk && chunk.last && !chunk.chunkedValue) { + return this._addSingleChunk(chunk); + } + const chunkValues = chunk.values; const numValues = chunkValues.length; const values: Value[] = new Array(numValues); @@ -398,6 +551,33 @@ export class PartialResultStream extends Transform implements ResultEvents { } return res; } + + /** + * Fast-path handler for single-chunk stream responses. Decodes rows directly + * and pushes them into the stream without incremental chunk buffering. + * + * @private + * @param {google.spanner.v1.PartialResultSet} chunk The partial result set. + * @returns {boolean} Whether the stream can accept more data. + */ + private _addSingleChunk(chunk: google.spanner.v1.PartialResultSet): boolean { + const rows = decodeRowsDirect( + chunk, + this._options, + this._fields, + this._decoders, + ); + let canAcceptMore = true; + for (let i = 0; i < rows.length; i++) { + const accepted = this.push(rows[i]); + if (!accepted && canAcceptMore) { + canAcceptMore = false; + this.emit('paused'); + } + } + return canAcceptMore; + } + /** * Manages complete values, pushing a completed row into the stream once all * values have been received. @@ -417,20 +597,9 @@ export class PartialResultStream extends Transform implements ResultEvents { this._values = []; - const isJsonStubbed = - codec.convertFieldsToJson !== originalConvertFieldsToJson; - - if (this._options.json && !isJsonStubbed) { - return this.push(this._createJsonRow(values)); - } - - const row: Row = this._createRow(values); - - if (this._options.json) { - return this.push(row.toJSON(this._options.jsonOptions)); - } - - return this.push(row); + return this.push( + formatRow(this._fields, this._decoders, values, this._options), + ); } /** @@ -443,31 +612,12 @@ export class PartialResultStream extends Transform implements ResultEvents { * @returns {Json} The plain JavaScript object representing the row. */ private _createJsonRow(values: Value[]): Json { - const json: Json = {}; - const fields = this._fields; - const decoders = this._decoders; - const len = fields.length; - const includeNameless = !!this._options.jsonOptions?.includeNameless; - - for (let i = 0; i < len; i++) { - const {name} = fields[i]; - if (!name && !includeNameless) { - continue; - } - const fieldName = name ? name : `_${i}`; - try { - json[fieldName] = decoders[i](values[i]); - } catch (e) { - (e as Error).message = [ - `Serializing column "${fieldName}" encountered an error: ${ - (e as Error).message - }`, - 'Call row.toJSON({ wrapNumbers: true }) to receive a custom type.', - ].join(' '); - throw e; - } - } - return json; + return createJsonRow( + this._fields, + this._decoders, + values, + Boolean(this._options.jsonOptions?.includeNameless), + ); } /** * Converts an array of values into a row. @@ -478,19 +628,7 @@ export class PartialResultStream extends Transform implements ResultEvents { * @returns {Row} */ private _createRow(values: Value[]): Row { - const len = values.length; - const fields = new RowImpl(len); - const decoders = this._decoders; - const classFields = this._fields; - - for (let i = 0; i < len; i++) { - fields[i] = { - name: classFields[i].name, - value: decoders[i](values[i]), - }; - } - - return fields; + return createRow(this._fields, this._decoders, values); } /** * Attempts to merge chunked values together. diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index bbd87d49fa9..490406d9fa9 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -29,6 +29,7 @@ import { partialResultStream, ResumeToken, Row, + decodeRowsDirect, } from './partial-result-stream'; import {Session} from './session'; import {Key} from './table'; @@ -48,12 +49,27 @@ import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel; import IAny = google.protobuf.IAny; import IQueryOptions = google.spanner.v1.ExecuteSqlRequest.IQueryOptions; import IRequestOptions = google.spanner.v1.IRequestOptions; +import IResultSetStats = google.spanner.v1.IResultSetStats; +import ResultSetStats = google.spanner.v1.ResultSetStats; +import IResultSetMetadata = google.spanner.v1.IResultSetMetadata; +import ResultSetMetadata = google.spanner.v1.ResultSetMetadata; import {Database, Spanner} from '.'; import ReadLockMode = google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode; -import {RunTransactionOptions} from './transaction-runner'; -import {injectRequestIDIntoHeaders, nextNthRequest} from './request_id_header'; +import { + DeadlineError, + RunTransactionOptions, + isRetryableInternalError, +} from './transaction-runner'; +import { + injectRequestIDIntoHeaders, + nextNthRequest, + X_GOOG_SPANNER_REQUEST_ID_HEADER, + X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, +} from './request_id_header'; export type Rows = Array; +type ResultStats = ResultSetStats | IResultSetStats; +type ResultMetadata = ResultSetMetadata | IResultSetMetadata; const RETRY_INFO_TYPE = 'type.googleapis.com/google.rpc.retryinfo'; const RETRY_INFO_BIN = 'google.rpc.retryinfo-bin'; @@ -73,6 +89,18 @@ function injectGaxOpt(existingOpts: any, key: string, value: any): any { }); } +function normalizeError(err: Error): ServiceError { + const errorWithCode = err as Partial; + if (errorWithCode.code === undefined) { + Object.assign(errorWithCode, { + code: grpc.status.UNKNOWN, + details: err.message, + metadata: new grpc.Metadata(), + }); + } + return errorWithCode as ServiceError; +} + export interface TimestampBounds { strong?: boolean; minReadTimestamp?: PreciseDate | spannerClient.protobuf.ITimestamp; @@ -224,10 +252,6 @@ export interface BatchUpdateCallback { response?: spannerClient.spanner.v1.ExecuteBatchDmlResponse, ): void; } -export interface BatchUpdateOptions { - requestOptions?: Omit; - gaxOptions?: CallOptions; -} export type ReadCallback = NormalCallback; @@ -1107,6 +1131,7 @@ export class Snapshot extends EventEmitter { } this.ended = true; + this._releaseWaitingRequests(); process.nextTick(() => this.emit('end')); if (this._affinityKey) { @@ -1117,7 +1142,7 @@ export class Snapshot extends EventEmitter { if (client?.spannerStub) { Promise.resolve(client.spannerStub) .then((stub: any) => { - stub?.getChannel?.()?.unbind?.(this._affinityKey); + return stub?.getChannel?.()?.unbind?.(this._affinityKey); }) .catch(() => {}); } @@ -1368,6 +1393,23 @@ export class Snapshot extends EventEmitter { query: string | ExecuteSqlRequest, callback?: RunCallback, ): void | Promise { + if (this.runStream !== Snapshot.prototype.runStream) { + return this._runLegacy(query, callback!); + } + return this._run(query, callback!); + } + + /** + * Always runs the query through the full streaming pipeline (Snapshot.prototype.runStream). + * Used when runStream has been overridden on Snapshot to preserve backward compatibility + * with custom stream implementations. + * + * @private + */ + private _runLegacy( + query: string | ExecuteSqlRequest, + callback: RunCallback, + ): void { const rows: Rows = []; let stats: google.spanner.v1.ResultSetStats; let metadata: google.spanner.v1.ResultSetMetadata; @@ -1379,18 +1421,11 @@ export class Snapshot extends EventEmitter { ...this._traceConfig, }, span => { - return this.runStream(query) + this.runStream(query) .on('error', err => { setSpanError(span, err); span.end(); - if (!('code' in err)) { - Object.assign(err, { - code: grpc.status.UNKNOWN, - details: err.message, - metadata: new grpc.Metadata(), - }); - } - callback!(err as ServiceError, rows, stats, metadata); + callback!(normalizeError(err), rows, stats, metadata); }) .on('response', response => { if (response.metadata) { @@ -1400,8 +1435,12 @@ export class Snapshot extends EventEmitter { } } }) - .on('data', row => rows.push(row)) - .on('stats', _stats => (stats = _stats)) + .on('data', row => { + rows.push(row); + }) + .on('stats', _stats => { + stats = _stats; + }) .on('end', () => { span.end(); callback!(null, rows, stats, metadata); @@ -1410,6 +1449,373 @@ export class Snapshot extends EventEmitter { ); } + /** + * Executes a query using an optimized streaming model: + * 1. If all results are returned in a single PartialResultSet, the internal streaming + * pipeline is simplified and rows are decoded directly, bypassing Stream overhead. + * 2. If there are more than one PartialResultSets, it seamlessly falls back to + * the standard streaming model. + * + * @private + */ + _run( + queryInput: string | ExecuteSqlRequest, + callback: RunCallback, + options?: {startRunSpan?: boolean}, + ): void { + let query: ExecuteSqlRequest = + typeof queryInput === 'string' ? {sql: queryInput} : queryInput; + + query = Object.assign({}, query) as ExecuteSqlRequest; + query.queryOptions = Object.assign( + Object.assign({}, this.queryOptions), + query.queryOptions, + ); + + const { + gaxOptions, + json, + jsonOptions, + maxResumeRetries, + requestOptions, + columnsMetadata, + types: _omittedTypes, + directedReadOptions: rawDirectedReadOptions, + ...rawQuery + } = query; + void _omittedTypes; + let formattedRequest: google.spanner.v1.IExecuteSqlRequest | undefined; + const seqno = this._seqno++; + + const directedReadOptions = this._getDirectedReadOptions( + rawDirectedReadOptions, + ); + + const sanitizeRequest = () => { + const {params, paramTypes} = Snapshot.encodeParams(query); + const transaction: spannerClient.spanner.v1.ITransactionSelector = {}; + if (this.id) { + transaction.id = this.id as Uint8Array; + } else if (this._options.readWrite) { + transaction.begin = this._options; + } else { + transaction.singleUse = this._options; + } + + if ( + !this.id && + this._options.readWrite && + (this.session.parent as Database).isMuxEnabledForRW_ + ) { + this._setPreviousTransactionId(transaction); + } + + formattedRequest = Object.assign({}, rawQuery, { + session: this.session.formattedName_!, + seqno, + requestOptions: this.configureTagOptions( + typeof transaction.singleUse !== 'undefined', + this.requestOptions?.transactionTag ?? undefined, + requestOptions, + ), + directedReadOptions, + transaction, + params, + paramTypes, + }); + }; + + const headers = Object.assign({}, this.commonHeaders_); + if ( + this._getSpanner().routeToLeaderEnabled && + (this._options.readWrite !== undefined || + this._options.partitionedDml !== undefined) + ) { + addLeaderAwareRoutingHeader(headers); + } + + const traceConfig = { + transactionTag: this.requestOptions?.transactionTag, + requestTag: requestOptions?.requestTag, + ...query, + ...this._traceConfig, + }; + + const executeWithSpans = ( + fn: (runSpan: Span | null, streamSpan: Span) => void, + ) => { + const startRunSpan = options?.startRunSpan !== false; + if (startRunSpan) { + return startTrace('Snapshot.run', traceConfig, runSpan => { + return startTrace('Snapshot.runStream', traceConfig, streamSpan => { + fn(runSpan, streamSpan); + }); + }); + } else { + return startTrace('Snapshot.runStream', traceConfig, streamSpan => { + fn(null, streamSpan); + }); + } + }; + + executeWithSpans((runSpan, streamSpan) => { + let attempt = 0; + const database = this.session.parent as Database; + const nthRequest = nextNthRequest(database); + + let completed = false; + const complete = ( + err: Error | null, + rows: Rows = [], + stats?: ResultStats, + metadata?: ResultMetadata, + ) => { + if (completed) { + return; + } + completed = true; + if (err) { + setSpanError(streamSpan, err); + } + streamSpan.end(); + if (runSpan) { + if (err) { + setSpanError(runSpan, err); + } + runSpan.end(); + } + callback!( + err ? normalizeError(err) : null, + rows, + stats as google.spanner.v1.ResultSetStats, + metadata as google.spanner.v1.ResultSetMetadata, + ); + }; + + const makeRequest = (resumeToken?: ResumeToken): Readable => { + attempt++; + if (!resumeToken) { + if (attempt === 1) { + streamSpan.addEvent('Starting stream'); + } else { + streamSpan.addEvent('Re-attempting start stream', {attempt}); + } + } else { + streamSpan.addEvent('Resuming stream', { + resume_token: resumeToken!.toString(), + attempt, + }); + } + + if ( + !formattedRequest || + (this.id && !formattedRequest.transaction?.id) + ) { + try { + sanitizeRequest(); + } catch (e) { + const errorStream = new PassThrough(); + setImmediate(() => { + errorStream.destroy(e as Error); + }); + return errorStream; + } + } + + const injectedHeaders = injectRequestIDIntoHeaders( + headers, + this.session, + nthRequest, + attempt, + ); + const requestId = injectedHeaders[X_GOOG_SPANNER_REQUEST_ID_HEADER]; + if (runSpan && requestId) { + runSpan.setAttribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, requestId); + } + + return this.requestStream({ + client: 'SpannerClient', + method: 'executeStreamingSql', + reqOpts: Object.assign({}, formattedRequest, { + resumeToken, + }), + gaxOpts: gaxOptions, + headers: injectedHeaders, + }); + }; + + const retryableCodes = [grpc.status.UNAVAILABLE]; + const wrappedMakeRequest = this._wrapWithIdWaiter(makeRequest); + const startTime = Date.now(); + const timeout = gaxOptions?.timeout ?? Infinity; + + const executeRequest = () => { + const requestStream = this.id ? makeRequest() : wrappedMakeRequest(); + + const onError = (err: grpc.ServiceError) => { + const elapsed = Date.now() - startTime; + let maxRetries = Infinity; + if (maxResumeRetries !== undefined) { + maxRetries = maxResumeRetries; + } else if (timeout === Infinity) { + maxRetries = 10; + } + if ( + elapsed < timeout && + attempt <= maxRetries && + err.code && + (retryableCodes.includes(err.code) || isRetryableInternalError(err)) + ) { + requestStream.removeAllListeners(); + requestStream.on('error', () => {}); + requestStream.destroy(); + setImmediate(() => { + executeRequest(); + }); + return; + } + + const wasAborted = isErrorAborted(err); + if (!this.id && this._useInRunner && !wasAborted) { + streamSpan.addEvent('Stream broken. Safe to retry'); + void this.begin(); + } else if (wasAborted) { + streamSpan.addEvent('Stream broken. Not safe to retry', { + 'transaction.id': this.id?.toString(), + }); + } + + const finalError = elapsed >= timeout ? new DeadlineError(err) : err; + complete(finalError); + }; + + const onEnd = () => { + requestStream.removeListener('error', onError); + complete(null, []); + }; + + const onData = (firstChunk: google.spanner.v1.PartialResultSet) => { + requestStream.removeListener('error', onError); + requestStream.removeListener('end', onEnd); + + if (firstChunk.last && !firstChunk.chunkedValue) { + // ⚡ FAST-PATH: Complete result in single chunk! + if (firstChunk.metadata?.transaction && !this.id) { + this._update(firstChunk.metadata.transaction, streamSpan); + } + this._updatePrecommitToken(firstChunk); + + let rows: Rows; + try { + rows = decodeRowsDirect(firstChunk, { + json, + jsonOptions, + columnsMetadata, + }); + } catch (decodeErr) { + requestStream.removeAllListeners(); + requestStream.on('error', () => {}); + requestStream.destroy(); + + complete( + decodeErr as Error, + [], + undefined, + firstChunk.metadata || undefined, + ); + return; + } + + // Query completed successfully. Do not cancel the gRPC call; allow it + // to drain remaining trailers/EOF in the background so it is not marked + // CANCELLED by Spanner or Cloud Monitoring. + requestStream.removeAllListeners('data'); + requestStream.removeAllListeners('error'); + requestStream.on('error', () => {}); + requestStream.resume(); + + complete( + null, + rows, + firstChunk.stats || undefined, + firstChunk.metadata || undefined, + ); + return; + } + + // 🌊 FALLBACK PATH: Multi-chunk query! + if (firstChunk.metadata?.transaction && !this.id) { + this._update(firstChunk.metadata.transaction, streamSpan); + } + this._updatePrecommitToken(firstChunk); + + requestStream.removeAllListeners('data'); + requestStream.removeAllListeners('error'); + + const replayStream = new PassThrough({objectMode: true}); + replayStream.write(firstChunk); + requestStream.pipe(replayStream); + + requestStream.on('error', err => { + replayStream.destroy(err); + }); + replayStream.on('close', () => { + if (!requestStream.destroyed) { + if (requestStream.readableEnded) { + requestStream.resume(); + } else { + requestStream.destroy(); + } + } + }); + + let initialStream: Readable | null = replayStream; + const fallbackMakeRequest = (resumeToken?: ResumeToken): Readable => { + if (!resumeToken && initialStream) { + const stream = initialStream; + initialStream = null; + return stream; + } + return makeRequest(resumeToken); + }; + + const rows: Rows = []; + let stats: google.spanner.v1.ResultSetStats; + let metadata: google.spanner.v1.ResultSetMetadata; + + const fallbackStream = partialResultStream(fallbackMakeRequest, { + json, + jsonOptions, + maxResumeRetries, + columnsMetadata, + gaxOptions, + }); + + fallbackStream + .on('response', response => { + this._updatePrecommitToken(response); + if (response.metadata?.transaction && !this.id) { + this._update(response.metadata.transaction, streamSpan); + } + if (response.metadata) { + metadata = response.metadata; + } + }) + .on('data', row => rows.push(row)) + .on('stats', _stats => (stats = _stats)) + .on('error', err => complete(err as Error, rows, stats, metadata)) + .on('end', () => complete(null, rows, stats, metadata)); + }; + + requestStream.once('error', onError); + requestStream.once('data', onData); + requestStream.once('end', onEnd); + }; + + executeRequest(); + }); + } + /** * ExecuteSql request options. This includes all standard ExecuteSqlRequest * options as well as several convenience properties. @@ -1910,6 +2316,9 @@ export class Snapshot extends EventEmitter { // Queue subsequent requests. return (resumeToken?: ResumeToken): Readable => { + if (this.id) { + return makeRequest(resumeToken); + } const streamProxy = new Readable({ read() {}, }); @@ -2499,8 +2908,8 @@ export class Transaction extends Dml { if ((this.session.parent as Database).isMuxEnabledForRW_) { this._setMutationKey(mutations); } - this.begin().then( - () => { + this.begin() + .then(() => { this.commit(options, (err, resp) => { if (err) { setSpanError(span, err); @@ -2508,13 +2917,13 @@ export class Transaction extends Dml { span.end(); callback(err, resp); }); - }, - err => { + return null; + }) + .catch(err => { setSpanError(span, err); span.end(); callback(err, null); - }, - ); + }); return; } diff --git a/handwritten/spanner/test/database.ts b/handwritten/spanner/test/database.ts index f2f3ea4088b..1539b572de0 100644 --- a/handwritten/spanner/test/database.ts +++ b/handwritten/spanner/test/database.ts @@ -47,7 +47,10 @@ import { BatchWriteOptions, CommitCallback, CommitOptions, + ExecuteSqlRequest, MutationSet, + RunCallback, + RunResponse, } from '../src/transaction'; import {SessionFactory} from '../src/session-factory'; import {RunTransactionOptions} from '../src/transaction-runner'; @@ -184,6 +187,13 @@ class FakeTransaction extends EventEmitter { } begin() {} end() {} + run(_query?: any, _callback?: any): void | Promise {} + _run( + query?: ExecuteSqlRequest | string, + callback?: RunCallback, + ): void | Promise { + return this.run(query, callback); + } runStream(): Transform { return through.obj(); } @@ -1941,6 +1951,204 @@ describe('Database', () => { }); }); + describe('run fast-path', () => { + const QUERY = 'SELECT 1'; + let fakeSessionFactory: FakeSessionFactory; + let fakeSession: FakeSession; + let fakeSnapshot: FakeTransaction; + let getSessionStub: sinon.SinonStub; + let snapshotStub: sinon.SinonStub; + let runStub: sinon.SinonStub; + + beforeEach(() => { + fakeSessionFactory = database.sessionFactory_; + fakeSession = new FakeSession(); + fakeSnapshot = new FakeTransaction( + {} as google.spanner.v1.TransactionOptions.ReadOnly, + ); + + getSessionStub = ( + sandbox.stub(fakeSessionFactory, 'getSession') as sinon.SinonStub + ).callsFake(callback => callback(null, fakeSession)); + + snapshotStub = sandbox + .stub(fakeSession, 'snapshot') + .returns(fakeSnapshot); + + runStub = ( + sandbox.stub(fakeSnapshot, 'run') as sinon.SinonStub + ).callsFake((query, optionsOrCallback, cb) => { + const callback = + typeof optionsOrCallback === 'function' ? optionsOrCallback : cb; + if (callback) { + callback(null, [{id: 1}]); + } + }); + + sandbox.stub(fakeSessionFactory, 'isMultiplexedEnabled').returns(true); + }); + + it('should execute query via snapshot.run fast-path and end snapshot', done => { + const endStub = sandbox.stub(fakeSnapshot, 'end'); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + assert.strictEqual(getSessionStub.callCount, 1); + assert.strictEqual(runStub.callCount, 1); + assert.strictEqual(runStub.lastCall.args[0], QUERY); + assert.strictEqual(endStub.callCount, 1); + done(); + }); + }); + + it('should release the session on snapshot end', done => { + const releaseStub = sandbox.stub( + fakeSessionFactory, + 'release', + ) as sinon.SinonStub; + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + fakeSnapshot.emit('end'); + assert.strictEqual(releaseStub.callCount, 1); + assert.strictEqual(releaseStub.lastCall.args[0], fakeSession); + done(); + }); + }); + + it('should propagate getSession error', done => { + const fakeError = new Error('No session'); + getSessionStub.callsFake(callback => callback(fakeError)); + + database.run(QUERY, (err, rows) => { + assert.strictEqual(err, fakeError); + assert.deepStrictEqual(rows, []); + done(); + }); + }); + + it('should pass options to session.snapshot', done => { + const options = {strong: true}; + database.run(QUERY, options, (err, rows) => { + assert.ifError(err); + assert.strictEqual(snapshotStub.lastCall.args[0], options); + done(); + }); + }); + + it('should propagate runMethod error', done => { + const queryError = Object.assign(new Error('Query execution failed'), { + code: grpc.status.INVALID_ARGUMENT, + }); + runStub.callsFake((query, optionsOrCallback, cb) => { + const callback = + typeof optionsOrCallback === 'function' ? optionsOrCallback : cb; + if (callback) { + callback(queryError); + } + }); + + database.run(QUERY, (err, rows) => { + assert.strictEqual(err, queryError); + assert.deepStrictEqual(rows, []); + done(); + }); + }); + + it('should emit error if session release throws on snapshot end', done => { + const releaseError = new Error('Release failed'); + const releaseStub = sandbox + .stub(fakeSessionFactory, 'release') + .throws(releaseError); + + database.once('error', err => { + assert.strictEqual(err, releaseError); + assert.strictEqual(releaseStub.callCount, 1); + done(); + }); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + fakeSnapshot.emit('end'); + }); + }); + + it('should fall back to streaming path when multiplexed session is disabled', done => { + (fakeSessionFactory.isMultiplexedEnabled as sinon.SinonStub).returns( + false, + ); + const runLegacyStub = sandbox + .stub(database as any, '_runLegacy') + .callsFake((query, options, callback: any) => { + assert.strictEqual(query, QUERY); + callback(null, [{id: 99}]); + }); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 99}]); + assert.strictEqual(runLegacyStub.callCount, 1); + assert.strictEqual(getSessionStub.callCount, 0); + done(); + }); + }); + + it('should fall back to streaming path when runStream is overridden', done => { + database.runStream = () => through.obj() as any; + const runLegacyStub = sandbox + .stub(database as any, '_runLegacy') + .callsFake((query, options, callback: any) => { + assert.strictEqual(query, QUERY); + callback(null, [{id: 100}]); + }); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 100}]); + assert.strictEqual(runLegacyStub.callCount, 1); + assert.strictEqual(getSessionStub.callCount, 0); + done(); + }); + }); + + it('should support Promise-based execution via promisify', async () => { + const runPromise = pfy.promisify(database.run.bind(database)); + const [rows] = await runPromise(QUERY); + assert.deepStrictEqual(rows, [{id: 1}]); + assert.strictEqual(runStub.callCount, 1); + }); + + it('should fall back to snapshot.run when snapshot.runStream is overridden', done => { + fakeSnapshot.runStream = () => through.obj() as any; + const snapshotRunStub = sandbox.stub(fakeSnapshot, '_run'); + + database.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 1}]); + assert.strictEqual(runStub.callCount, 1); + assert.strictEqual(snapshotRunStub.callCount, 0); + assert.strictEqual(runStub.lastCall.args.length, 2); + done(); + }); + }); + + it('should catch synchronous error in runMethod, end snapshot and propagate error', done => { + const syncError = new Error('Synchronous parameter failure'); + const endStub = sandbox.stub(fakeSnapshot, 'end'); + runStub.throws(syncError); + + database.run(QUERY, (err, rows) => { + assert.strictEqual(err, syncError); + assert.deepStrictEqual(rows, []); + assert.strictEqual(endStub.callCount, 1); + done(); + }); + }); + }); + describe('runStream', () => { const QUERY = { sql: 'SELECT * FROM table', diff --git a/handwritten/spanner/test/partial-result-stream.ts b/handwritten/spanner/test/partial-result-stream.ts index bb4334c0c38..abc5fd04bf6 100644 --- a/handwritten/spanner/test/partial-result-stream.ts +++ b/handwritten/spanner/test/partial-result-stream.ts @@ -614,6 +614,89 @@ describe('PartialResultStream', () => { stream.resume(); }); }); + + it('should emit paused event when downstream backpressure is triggered during single-chunk decode', done => { + const stream = new PartialResultStream({}); + let pausedEmitted = false; + stream.on('paused', () => { + pausedEmitted = true; + }); + + sandbox.stub(stream, 'push').callsFake(data => { + if (data === undefined || data === null) { + return true; + } + return false; + }); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('row1')], + last: true, + }); + + assert.strictEqual(pausedEmitted, true); + done(); + }); + + it('should route first chunk with last=true to _addSingleChunk', done => { + const stream = new PartialResultStream({}); + const addSingleChunkSpy = sandbox.spy(stream as any, '_addSingleChunk'); + const rows: any[] = []; + stream + .on('data', row => rows.push(row)) + .on('end', () => { + try { + assert.strictEqual(addSingleChunkSpy.calledOnce, true); + assert.strictEqual(rows.length, 1); + done(); + } catch (err) { + done(err); + } + }) + .on('error', done); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('row1')], + last: true, + }); + stream.end(); + }); + + it('should not route subsequent chunks to _addSingleChunk even if last=true', done => { + const stream = new PartialResultStream({}); + const addSingleChunkSpy = sandbox.spy(stream as any, '_addSingleChunk'); + const rows: any[] = []; + stream + .on('data', row => rows.push(row)) + .on('end', () => { + try { + assert.strictEqual(addSingleChunkSpy.called, false); + assert.strictEqual(rows.length, 2); + done(); + } catch (err) { + done(err); + } + }) + .on('error', done); + + const fields = [{name: NAME, type: {code: 'STRING'}}]; + // First chunk: not last + stream.write({ + metadata: {rowType: {fields}}, + values: [convertToIValue('row1')], + last: false, + }); + // Second chunk: last + stream.write({ + values: [convertToIValue('row2')], + last: true, + }); + stream.end(); + }); }); describe('partialResultStream', () => { @@ -1482,9 +1565,174 @@ describe('PartialResultStream', () => { }); }); }); + + describe('decodeRowsDirect & createFieldDecoders', () => { + it('should return empty array if fields or values are empty', () => { + assert.deepStrictEqual(prs.decodeRowsDirect({values: []} as any), []); + assert.deepStrictEqual( + prs.decodeRowsDirect({ + metadata: {rowType: {fields: []}}, + values: [convertToIValue('test')], + } as any), + [], + ); + }); + + it('should decode basic rows in standard RowImpl mode', () => { + const chunk: any = { + metadata: { + rowType: { + fields: [ + {name: 'id', type: {code: 'INT64'}}, + {name: 'name', type: {code: 'STRING'}}, + ], + }, + }, + values: [convertToIValue('101'), convertToIValue('Alice')], + }; + + const rows = prs.decodeRowsDirect(chunk); + assert.strictEqual(rows.length, 1); + assert.strictEqual(Array.isArray(rows[0]), true); + assert.strictEqual(rows[0][0].name, 'id'); + assert.strictEqual(rows[0][0].value.value, '101'); + assert.strictEqual(rows[0][1].name, 'name'); + assert.strictEqual(rows[0][1].value, 'Alice'); + assert.deepStrictEqual(rows[0].toJSON(), {id: 101, name: 'Alice'}); + }); + + it('should decode rows in JSON mode directly', () => { + const chunk: any = { + metadata: { + rowType: { + fields: [ + {name: 'id', type: {code: 'INT64'}}, + {name: 'name', type: {code: 'STRING'}}, + ], + }, + }, + values: [convertToIValue('101'), convertToIValue('Alice')], + }; + + const rows = prs.decodeRowsDirect(chunk, {json: true}); + assert.strictEqual(rows.length, 1); + assert.deepStrictEqual(rows[0], {id: 101, name: 'Alice'}); + }); + + it('should handle nameless columns in JSON mode', () => { + const chunk: any = { + metadata: { + rowType: { + fields: [ + {name: '', type: {code: 'INT64'}}, + {name: 'name', type: {code: 'STRING'}}, + ], + }, + }, + values: [convertToIValue('101'), convertToIValue('Alice')], + }; + + // Default: omit nameless columns + const rowsOmitted = prs.decodeRowsDirect(chunk, {json: true}); + assert.deepStrictEqual(rowsOmitted[0], {name: 'Alice'}); + + // With includeNameless: true + const rowsIncluded = prs.decodeRowsDirect(chunk, { + json: true, + jsonOptions: {includeNameless: true}, + }); + assert.deepStrictEqual(rowsIncluded[0], {_0: 101, name: 'Alice'}); + }); + + it('should wrap serialization errors in JSON mode with actionable error message', () => { + const chunk: any = { + metadata: { + rowType: { + fields: [{name: 'large_num', type: {code: 'INT64'}}], + }, + }, + values: [convertToIValue('9223372036854775807')], + }; + + assert.throws( + () => { + prs.decodeRowsDirect(chunk, { + json: true, + jsonOptions: {wrapNumbers: false}, + }); + }, + (err: Error) => { + assert( + err.message.includes( + 'Serializing column "large_num" encountered an error', + ), + ); + assert( + err.message.includes( + 'Call row.toJSON({ wrapNumbers: true }) to receive a custom type.', + ), + ); + return true; + }, + ); + }); + + it('should respect custom columnsMetadata in decoders', () => { + const getDecoderSpy = sandbox.spy(codec, 'getDecoder'); + const mockProtoType = {decode: () => {}, toObject: () => {}}; + const chunk: any = { + metadata: { + rowType: { + fields: [{name: 'protoCol', type: {code: 'PROTO'}}], + }, + }, + values: [convertToIValue(Buffer.from('test').toString('base64'))], + }; + + const columnsMetadata = { + protoCol: mockProtoType, + }; + + prs.decodeRowsDirect(chunk, { + columnsMetadata, + }); + assert.strictEqual(getDecoderSpy.called, true); + const [, columnMetadataArg] = getDecoderSpy.lastCall.args; + assert.strictEqual(columnMetadataArg, mockProtoType); + }); + + it('should call custom codec.decode when codec.decode is stubbed', () => { + const stub = sandbox.stub(codec, 'decode').returns('custom_decoded'); + const fields = [{name: 'col', type: {code: 'STRING'}}] as any; + const decoders = prs.createFieldDecoders(fields); + + const result = decoders[0]('raw'); + assert.strictEqual(result, 'custom_decoded'); + assert.strictEqual(stub.callCount, 1); + }); + + it('should fall back to RowImpl toJSON when codec.convertFieldsToJson is stubbed in JSON mode', () => { + const stub = sandbox + .stub(codec, 'convertFieldsToJson') + .returns({mocked: true} as any); + + const chunk: any = { + metadata: { + rowType: { + fields: [{name: 'id', type: {code: 'INT64'}}], + }, + }, + values: [convertToIValue('101')], + }; + + const rows = prs.decodeRowsDirect(chunk, {json: true}); + assert.deepStrictEqual(rows[0], {mocked: true}); + assert.strictEqual(stub.callCount, 1); + }); + }); }); -function convertToIValue(value) { +export function convertToIValue(value) { let kind: string; if (typeof value === 'number') { diff --git a/handwritten/spanner/test/transaction.ts b/handwritten/spanner/test/transaction.ts index 1e9794de5f8..4d0be66e437 100644 --- a/handwritten/spanner/test/transaction.ts +++ b/handwritten/spanner/test/transaction.ts @@ -41,8 +41,13 @@ import { BatchUpdateOptions, ExecuteSqlRequest, ReadRequest, + RunCallback, } from '../src/transaction'; +import {Row} from '../src/partial-result-stream'; import {grpc} from 'google-gax'; +import * as through from 'through2'; + +import {convertToIValue} from './partial-result-stream'; describe('Transaction', () => { const sandbox = sinon.createSandbox(); @@ -109,7 +114,10 @@ describe('Transaction', () => { const txns = proxyquire('../src/transaction', { '@google-cloud/promisify': {promisifyAll: PROMISIFY_ALL}, './codec': {codec}, - './partial-result-stream': {partialResultStream: PARTIAL_RESULT_STREAM}, + './partial-result-stream': { + ...require('../src/partial-result-stream'), + partialResultStream: PARTIAL_RESULT_STREAM, + }, }); Snapshot = txns.Snapshot; @@ -719,6 +727,805 @@ describe('Transaction', () => { }); }); + describe('run fast-path', () => { + const QUERY = {sql: 'SELECT * FROM `MyTable`'}; + + beforeEach(() => { + REQUEST_STREAM.resetHistory(); + PARTIAL_RESULT_STREAM.resetHistory(); + }); + + it('should execute single-chunk query via fast-path without partialResultStream', done => { + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(rows.length, 1); + const row = rows[0] as Row; + assert.strictEqual(row[0].name, 'col1'); + assert.strictEqual(row[0].value, 'val1'); + assert.deepStrictEqual(row.toJSON(), {col1: 'val1'}); + assert.strictEqual(PARTIAL_RESULT_STREAM.called, false); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('val1')], + last: true, + }); + }); + + it('should accept plain SQL string input on fast-path', done => { + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run('SELECT 1', (error, rows) => { + try { + assert.ifError(error); + assert.strictEqual(rows!.length, 1); + const row = rows![0] as Row; + assert.strictEqual(row[0].name, 'col1'); + assert.strictEqual(row[0].value, 'val1'); + assert.strictEqual(PARTIAL_RESULT_STREAM.called, false); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('val1')], + last: true, + }); + }); + + it('should seamlessly fall back to streaming for multi-chunk query', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const { + partialResultStream: realPartialResultStream, + } = require('../src/partial-result-stream'); + PARTIAL_RESULT_STREAM.callsFake((makeRequest: any, options: any) => + realPartialResultStream(makeRequest, options), + ); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(rows.length, 2); + assert.strictEqual(PARTIAL_RESULT_STREAM.called, true); + done(); + } catch (error) { + done(error); + } + }); + + // Chunk 1: last = false + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('val1')], + last: false, + resumeToken: 'token1', + }); + + // Chunk 2: last = true + fakeRequestStream.push({ + values: [convertToIValue('val2')], + last: true, + }); + fakeRequestStream.push(null); + }); + + it('should fall back to streaming when first chunk has chunkedValue=true', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const { + partialResultStream: realPartialResultStream, + } = require('../src/partial-result-stream'); + PARTIAL_RESULT_STREAM.callsFake((makeRequest: any, options: any) => + realPartialResultStream(makeRequest, options), + ); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run(QUERY, (error, rows) => { + try { + assert.ifError(error); + assert.strictEqual(PARTIAL_RESULT_STREAM.called, true); + assert.strictEqual(rows!.length, 1); + assert.strictEqual((rows![0] as Row)[0].value, 'hello world'); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('hello ')], + chunkedValue: true, + last: false, + }); + + setImmediate(() => { + fakeRequestStream.push({ + values: [convertToIValue('world')], + last: true, + }); + fakeRequestStream.push(null); + }); + }); + + it('should update inline transaction ID and precommitToken on fast-path', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + const fakeTxId = Buffer.from('tx-id-123'); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(snapshot.id, fakeTxId); + assert.strictEqual(PARTIAL_RESULT_STREAM.called, false); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.push({ + metadata: { + transaction: {id: fakeTxId}, + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + precommitToken: {precommitToken: 'token-abc'}, + values: [convertToIValue('val1')], + last: true, + }); + }); + + it('should drain trailers in background without canceling request stream on chunk.last', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + const resumeSpy = sandbox.spy(fakeRequestStream, 'resume'); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(rows.length, 1); + assert.strictEqual(resumeSpy.called, true); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('val1')], + last: true, + }); + }); + + it('should return decorated error with UNKNOWN code if decoding throws', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run( + {sql: 'SELECT * FROM `MyTable`', json: true}, + (err, rows) => { + try { + assert(err); + assert.strictEqual((err as any).code, grpc.status.UNKNOWN); + assert.strictEqual(rows.length, 0); + done(); + } catch (error) { + done(error); + } + }, + ); + + // Push chunk that causes integer overflow decoding error + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'INT64'}}], + }, + }, + values: [{stringValue: '9223372036854775807', kind: 'stringValue'}], + last: true, + }); + }); + + it('should return error for invalid parameters', done => { + const fakeQuery = { + sql: 'SELECT * FROM `MyTable`', + params: {a: undefined}, + }; + + snapshot.run(fakeQuery, (error, rows) => { + try { + assert(error); + assert.strictEqual( + error!.message, + 'Value of type undefined not recognized.', + ); + assert.strictEqual(rows!.length, 0); + assert.strictEqual(REQUEST_STREAM.called, false); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + }); + + it('should return error on non-retryable gRPC error', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + const testError = Object.assign(new Error('Invalid SQL syntax'), { + code: grpc.status.INVALID_ARGUMENT, + }); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.strictEqual(err, testError); + assert.strictEqual(rows.length, 0); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.emit('error', testError); + }); + + it('should trigger begin when non-aborted error occurs without transaction id in runner', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.id = undefined; + snapshot._useInRunner = true; + const beginStub = sandbox.stub(snapshot, 'begin').resolves(); + + const testError = Object.assign(new Error('Internal server error'), { + code: grpc.status.INTERNAL, + }); + + snapshot.run(QUERY, (error, rows) => { + try { + assert.strictEqual(error, testError); + assert.strictEqual(rows!.length, 0); + assert.strictEqual(beginStub.calledOnce, true); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream.emit('error', testError); + }); + + it('should retry retryable error up to maxResumeRetries', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + let attempt = 0; + REQUEST_STREAM.callsFake(() => { + attempt++; + const stream = through.obj(); + setImmediate(() => { + if (attempt === 1) { + stream.emit( + 'error', + Object.assign(new Error('Unavailable'), { + code: grpc.status.UNAVAILABLE, + }), + ); + } else { + stream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('success')], + last: true, + }); + } + }); + return stream; + }); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(attempt, 2); + assert.strictEqual(rows.length, 1); + done(); + } catch (error) { + done(error); + } + }); + }); + + it('should handle empty result set when stream ends without chunks', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(rows.length, 0); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.end(); + }); + + it('should preserve the same seqno across retries', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + let attempt = 0; + REQUEST_STREAM.callsFake(() => { + attempt++; + const stream = through.obj(); + setImmediate(() => { + if (attempt === 1) { + stream.emit( + 'error', + Object.assign(new Error('Unavailable'), { + code: grpc.status.UNAVAILABLE, + }), + ); + } else { + stream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('success')], + last: true, + }); + } + }); + return stream; + }); + + const initialSeqno = snapshot._seqno; + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(attempt, 2); + assert.strictEqual(rows.length, 1); + const firstSeqno = REQUEST_STREAM.firstCall.args[0].reqOpts.seqno; + const secondSeqno = REQUEST_STREAM.secondCall.args[0].reqOpts.seqno; + assert.strictEqual(firstSeqno, secondSeqno); + assert.strictEqual(firstSeqno, initialSeqno); + assert.strictEqual(snapshot._seqno, initialSeqno + 1); + done(); + } catch (error) { + done(error); + } + }); + }); + + it('should fall back to multi-chunk partialResultStream when result spans multiple chunks', done => { + const { + partialResultStream: realPartialResultStream, + } = require('../src/partial-result-stream'); + PARTIAL_RESULT_STREAM.callsFake((makeRequest: any, options: any) => + realPartialResultStream(makeRequest, options), + ); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(rows.length, 2); + assert.strictEqual(rows[0][0].value, 'first'); + assert.strictEqual(rows[1][0].value, 'second'); + assert.strictEqual(PARTIAL_RESULT_STREAM.callCount, 1); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('first')], + last: false, + resumeToken: 'tok1', + }); + fakeRequestStream.push({ + values: [convertToIValue('second')], + last: true, + }); + fakeRequestStream.push(null); + }); + + it('should propagate error from multi-chunk fallback stream', done => { + const { + partialResultStream: realPartialResultStream, + } = require('../src/partial-result-stream'); + PARTIAL_RESULT_STREAM.callsFake((makeRequest: any, options: any) => + realPartialResultStream(makeRequest, options), + ); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + const testError = Object.assign(new Error('Mid-stream error'), { + code: grpc.status.CANCELLED, + }); + + snapshot.run(QUERY, (err, rows) => { + try { + assert(err); + assert.strictEqual((err as any).code, grpc.status.CANCELLED); + done(); + } catch (error) { + done(error); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('first')], + last: false, + }); + fakeRequestStream.destroy(testError); + }); + + it('should retry multi-chunk fallback stream with resumeToken on retryable error', done => { + const { + partialResultStream: realPartialResultStream, + } = require('../src/partial-result-stream'); + PARTIAL_RESULT_STREAM.callsFake((makeRequest: any, options: any) => + realPartialResultStream(makeRequest, options), + ); + const fakeRequestStream1 = through.obj(); + const fakeRequestStream2 = through.obj(); + let requestAttempt = 0; + + REQUEST_STREAM.callsFake(() => { + requestAttempt++; + if (requestAttempt === 1) { + return fakeRequestStream1; + } + return fakeRequestStream2; + }); + + snapshot.run(QUERY, (error, rows) => { + try { + assert.ifError(error); + assert.strictEqual(requestAttempt, 2); + assert.strictEqual(rows!.length, 2); + assert.strictEqual((rows![0] as Row)[0].value, 'first'); + assert.strictEqual((rows![1] as Row)[0].value, 'second'); + const secondCallReqOpts = REQUEST_STREAM.secondCall.args[0].reqOpts; + assert.strictEqual(secondCallReqOpts.resumeToken, 'resume-token-1'); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream1.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('first')], + last: false, + resumeToken: 'resume-token-1', + }); + + setImmediate(() => { + fakeRequestStream1.emit( + 'error', + Object.assign(new Error('Unavailable'), { + code: grpc.status.UNAVAILABLE, + }), + ); + + setImmediate(() => { + fakeRequestStream2.push({ + values: [convertToIValue('second')], + last: true, + }); + fakeRequestStream2.push(null); + }); + }); + }); + + it('should retry when isRetryableInternalError occurs', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + let attempt = 0; + REQUEST_STREAM.callsFake(() => { + attempt++; + const stream = through.obj(); + setImmediate(() => { + if (attempt === 1) { + stream.emit( + 'error', + Object.assign( + new Error( + 'Received unexpected EOS on DATA frame from server', + ), + {code: grpc.status.INTERNAL}, + ), + ); + } else { + stream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('recovered')], + last: true, + }); + } + }); + return stream; + }); + + snapshot.run(QUERY, (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(attempt, 2); + assert.strictEqual(rows.length, 1); + done(); + } catch (error) { + done(error); + } + }); + }); + + it('should stop retrying when attempts exceed maxResumeRetries', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + let attempts = 0; + REQUEST_STREAM.callsFake(() => { + attempts++; + const stream = through.obj(); + setImmediate(() => { + stream.emit( + 'error', + Object.assign(new Error('Unavailable'), { + code: grpc.status.UNAVAILABLE, + }), + ); + }); + return stream; + }); + + snapshot.run({sql: 'SELECT 1', maxResumeRetries: 1}, (err, rows) => { + try { + assert(err); + assert.strictEqual((err as any).code, grpc.status.UNAVAILABLE); + assert.strictEqual(attempts, 2); // attempt 1 + 1 retry + done(); + } catch (error) { + done(error); + } + }); + }); + + it('should not retry when elapsed time exceeds timeout', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + let attempts = 0; + REQUEST_STREAM.callsFake(() => { + attempts++; + const stream = through.obj(); + setImmediate(() => { + stream.emit( + 'error', + Object.assign(new Error('Unavailable'), { + code: grpc.status.UNAVAILABLE, + }), + ); + }); + return stream; + }); + + snapshot.run( + {sql: 'SELECT 1', gaxOptions: {timeout: -1}}, + (error, rows) => { + try { + assert(error); + assert.strictEqual( + (error as any).code, + grpc.status.DEADLINE_EXCEEDED, + ); + assert.strictEqual(attempts, 1); // no retries allowed because elapsed > timeout + done(); + } catch (assertionError) { + done(assertionError); + } + }, + ); + }); + + it('should fall back to _runLegacy when runStream is overridden on Snapshot', done => { + snapshot.runStream = () => through.obj() as any; + const runLegacyStub = sandbox + .stub(snapshot as any, '_runLegacy') + .callsFake((query, callback: any) => { + assert.strictEqual(query, QUERY); + callback(null, [{id: 'from_stream'}]); + }); + + snapshot.run(QUERY, (err, rows) => { + assert.ifError(err); + assert.deepStrictEqual(rows, [{id: 'from_stream'}]); + assert.strictEqual(runLegacyStub.callCount, 1); + done(); + }); + }); + + it('should support options.startRunSpan = false', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + snapshot._run( + QUERY, + (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(rows.length, 1); + done(); + } catch (error) { + done(error); + } + }, + {startRunSpan: false}, + ); + + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('ok')], + last: true, + }); + }); + + it('should ensure complete is idempotent in fallbackStream path', done => { + let callbackCalls = 0; + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + const fakeFallbackStream = through.obj(); + PARTIAL_RESULT_STREAM.returns(fakeFallbackStream); + + snapshot.run(QUERY, (err, rows) => { + callbackCalls++; + try { + assert(err); + assert.strictEqual((err as any).message, 'Stream failure'); + assert.strictEqual(callbackCalls, 1); + } catch (e) { + done(e); + } + }); + + // First chunk triggers fallback path + fakeRequestStream.push({ + metadata: { + rowType: { + fields: [{name: 'id', type: {code: 'INT64'}}], + }, + }, + values: [convertToIValue(1)], + last: false, + }); + + setImmediate(() => { + // Emit error, then end + fakeFallbackStream.emit('error', new Error('Stream failure')); + fakeFallbackStream.emit('end'); + + setImmediate(() => { + assert.strictEqual(callbackCalls, 1); + done(); + }); + }); + }); + + it('should allow retries beyond 10 attempts as long as elapsed time is within timeout', done => { + PARTIAL_RESULT_STREAM.resetHistory(); + let attempts = 0; + REQUEST_STREAM.callsFake(() => { + attempts++; + const stream = through.obj(); + setImmediate(() => { + if (attempts < 15) { + stream.emit( + 'error', + Object.assign(new Error('Unavailable'), { + code: grpc.status.UNAVAILABLE, + }), + ); + } else { + stream.push({ + metadata: { + rowType: { + fields: [{name: 'id', type: {code: 'INT64'}}], + }, + }, + values: [convertToIValue(42)], + last: true, + }); + } + }); + return stream; + }); + + snapshot.run( + {sql: 'SELECT 1', gaxOptions: {timeout: 5000}}, + (err, rows) => { + try { + assert.ifError(err); + assert.strictEqual(attempts, 15); + assert.strictEqual(rows!.length, 1); + done(); + } catch (error) { + done(error); + } + }, + ); + }); + }); + describe('runStream', () => { const QUERY = { sql: 'SELECT * FROM `MyTable`', @@ -1365,6 +2172,32 @@ describe('Transaction', () => { assert.strictEqual(callback.callCount, 1); assert.strictEqual(callback.args[0][1], Math.floor(fakeRowCount)); }); + + it('should execute end-to-end via fast-path and return rowCountExact', done => { + const fakeRequestStream = through.obj(); + REQUEST_STREAM.returns(fakeRequestStream); + + dml.runUpdate(SQL, (error, rowCount) => { + try { + assert.ifError(error); + assert.strictEqual(rowCount, 42); + done(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream.push({ + metadata: { + rowType: {fields: []}, + }, + stats: { + rowCount: 'rowCountExact', + rowCountExact: 42, + }, + last: true, + }); + }); }); }); @@ -2890,6 +3723,275 @@ describe('Transaction', () => { }); }); + describe('run', () => { + beforeEach(() => { + REQUEST_STREAM.reset(); + REQUEST_STREAM.resetHistory(); + PARTIAL_RESULT_STREAM.reset(); + PARTIAL_RESULT_STREAM.resetHistory(); + }); + + it('should handle parallel queries on an un-begun read-write transaction', done => { + const fakeRequestStream1 = through.obj(); + const fakeRequestStream2 = through.obj(); + let requestCount = 0; + + REQUEST_STREAM.callsFake(() => { + requestCount++; + if (requestCount === 1) { + return fakeRequestStream1; + } + return fakeRequestStream2; + }); + + const fakeTransactionId = Buffer.from('tx-id-parallel'); + let query1Done = false; + let query2Done = false; + + const checkBothDone = () => { + if (query1Done && query2Done) { + try { + assert.strictEqual(transaction.id, fakeTransactionId); + assert.strictEqual(PARTIAL_RESULT_STREAM.called, false); + assert.strictEqual(requestCount, 2); + const firstCallReqOpts = REQUEST_STREAM.firstCall.args[0].reqOpts; + const secondCallReqOpts = + REQUEST_STREAM.secondCall.args[0].reqOpts; + assert.deepStrictEqual(firstCallReqOpts.transaction, { + begin: { + readWrite: {}, + isolationLevel: IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED, + }, + }); + assert.deepStrictEqual(secondCallReqOpts.transaction, { + id: fakeTransactionId, + }); + done(); + } catch (error) { + done(error); + } + } + }; + + transaction.run({sql: 'SELECT 1'}, (error: any, rows: any) => { + try { + assert.ifError(error); + assert.strictEqual(rows!.length, 1); + assert.strictEqual((rows![0] as Row)[0].value, 'first'); + query1Done = true; + checkBothDone(); + } catch (assertionError) { + done(assertionError); + } + }); + + transaction.run({sql: 'SELECT 2'}, (error: any, rows: any) => { + try { + assert.ifError(error); + assert.strictEqual(rows!.length, 1); + assert.strictEqual((rows![0] as Row)[0].value, 'second'); + query2Done = true; + checkBothDone(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream1.push({ + metadata: { + transaction: {id: fakeTransactionId}, + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('first')], + last: true, + }); + + fakeRequestStream2.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('second')], + last: true, + }); + }); + + it('should fall back to stream for queued parallel request returning multiple chunks', done => { + const { + partialResultStream: realPartialResultStream, + } = require('../src/partial-result-stream'); + PARTIAL_RESULT_STREAM.callsFake((makeRequest: any, options: any) => + realPartialResultStream(makeRequest, options), + ); + const fakeRequestStream1 = through.obj(); + const fakeRequestStream2 = through.obj(); + let requestCount = 0; + + REQUEST_STREAM.callsFake(() => { + requestCount++; + if (requestCount === 1) { + return fakeRequestStream1; + } + return fakeRequestStream2; + }); + + const fakeTransactionId = Buffer.from('tx-id-queued-multi-chunk'); + let query1Done = false; + let query2Done = false; + + const checkBothDone = () => { + if (query1Done && query2Done) { + try { + assert.strictEqual(transaction.id, fakeTransactionId); + assert.strictEqual(requestCount, 2); + done(); + } catch (assertionError) { + done(assertionError); + } + } + }; + + transaction.run({sql: 'SELECT 1'}, (error: any, rows: any) => { + try { + assert.ifError(error); + assert.strictEqual(rows!.length, 1); + assert.strictEqual((rows![0] as Row)[0].value, 'first'); + query1Done = true; + checkBothDone(); + } catch (assertionError) { + done(assertionError); + } + }); + + transaction.run({sql: 'SELECT 2'}, (error: any, rows: any) => { + try { + assert.ifError(error); + assert.strictEqual(rows!.length, 2); + assert.strictEqual((rows![0] as Row)[0].value, 'chunk1-val'); + assert.strictEqual((rows![1] as Row)[0].value, 'chunk2-val'); + query2Done = true; + checkBothDone(); + } catch (assertionError) { + done(assertionError); + } + }); + + // Query 1 completes via fast-path with inline transaction id + fakeRequestStream1.push({ + metadata: { + transaction: {id: fakeTransactionId}, + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('first')], + last: true, + }); + + // Query 2 is released from queue and returns multiple chunks, falling back to stream + setImmediate(() => { + fakeRequestStream2.push({ + metadata: { + rowType: { + fields: [{name: 'col2', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('chunk1-val')], + last: false, + resumeToken: 'resume-token-query2', + }); + setImmediate(() => { + fakeRequestStream2.push({ + values: [convertToIValue('chunk2-val')], + last: true, + }); + fakeRequestStream2.push(null); + }); + }); + }); + + it('should update inline transaction ID and release waiting queries even if row decoding throws', done => { + const fakeRequestStream1 = through.obj(); + const fakeRequestStream2 = through.obj(); + let requestCount = 0; + + REQUEST_STREAM.callsFake(() => { + requestCount++; + if (requestCount === 1) { + return fakeRequestStream1; + } + return fakeRequestStream2; + }); + + const fakeTransactionId = Buffer.from('tx-id-decode-error'); + let query1Done = false; + let query2Done = false; + + const checkBothDone = () => { + if (query1Done && query2Done) { + try { + assert.strictEqual(transaction.id, fakeTransactionId); + assert.strictEqual(requestCount, 2); + done(); + } catch (assertionError) { + done(assertionError); + } + } + }; + + transaction.run( + {sql: 'SELECT 1', json: true}, + (error: any, rows: any) => { + try { + assert(error); + assert.strictEqual((error as any).code, grpc.status.UNKNOWN); + assert.strictEqual(rows!.length, 0); + query1Done = true; + checkBothDone(); + } catch (assertionError) { + done(assertionError); + } + }, + ); + + transaction.run({sql: 'SELECT 2'}, (error: any, rows: any) => { + try { + assert.ifError(error); + assert.strictEqual(rows!.length, 1); + assert.strictEqual((rows![0] as Row)[0].value, 'second'); + query2Done = true; + checkBothDone(); + } catch (assertionError) { + done(assertionError); + } + }); + + fakeRequestStream1.push({ + metadata: { + transaction: {id: fakeTransactionId}, + rowType: { + fields: [{name: 'col1', type: {code: 'INT64'}}], + }, + }, + values: [{stringValue: '9223372036854775807', kind: 'stringValue'}], + last: true, + }); + + fakeRequestStream2.push({ + metadata: { + rowType: { + fields: [{name: 'col1', type: {code: 'STRING'}}], + }, + }, + values: [convertToIValue('second')], + last: true, + }); + }); + }); + describe('runStream', () => { before(() => { PARTIAL_RESULT_STREAM.callsFake(makeRequest => makeRequest());