diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index f84b73362eb..a41aeab27cf 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -438,35 +438,12 @@ export class Snapshot extends EventEmitter { }, }, }; - this.request = (config: any, callback?: Function) => { - let gaxOpts; - if (!config.gaxOpts || Object.keys(config.gaxOpts).length === 0) { - gaxOpts = this._bindGaxOpts as any; - } else { - gaxOpts = injectGaxOpt( - config.gaxOpts, - 'affinityKey', - this._affinityKey, - ); - } - config = Object.assign({}, config, {gaxOpts}); - return session.request(config, callback); - }; - - this.requestStream = (config: any) => { - let gaxOpts; - if (!config.gaxOpts || Object.keys(config.gaxOpts).length === 0) { - gaxOpts = this._bindGaxOpts as any; - } else { - gaxOpts = injectGaxOpt( - config.gaxOpts, - 'affinityKey', - this._affinityKey, - ); - } - config = Object.assign({}, config, {gaxOpts}); - return session.requestStream(config); - }; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + this.request = (config: any, callback?: Function) => + session.request(this._applyAffinityGaxOpts(config), callback); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + this.requestStream = (config: any) => + session.requestStream(this._applyAffinityGaxOpts(config)); } else { this.request = session.request.bind(session); this.requestStream = session.requestStream.bind(session); @@ -487,6 +464,38 @@ export class Snapshot extends EventEmitter { this._mutationKey = null; } + /** + * Binds the multiplexed session affinity key to the gax options of an + * outgoing request, so that all requests of this transaction are routed to + * the same gRPC channel. + * + * `config` is always a request descriptor that was freshly constructed by the + * caller for this one RPC (and {@link Spanner#prepareGapicRequest_} already + * modifies `config.headers` in place), so the affinity key is assigned + * directly instead of allocating a copy of the descriptor per request. + * + * @private + * + * @param {object} config The request configuration. + * @returns {object} The same request configuration. + */ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + private _applyAffinityGaxOpts(config: any): any { + if (!config) { + return config; + } + if (!config.gaxOpts || Object.keys(config.gaxOpts).length === 0) { + config.gaxOpts = this._bindGaxOpts; + } else { + config.gaxOpts = injectGaxOpt( + config.gaxOpts, + 'affinityKey', + this._affinityKey, + ); + } + return config; + } + protected _updatePrecommitToken(resp: PrecommitTokenProvider): void { if ( this._latestPreCommitToken === null || diff --git a/handwritten/spanner/test/transaction.ts b/handwritten/spanner/test/transaction.ts index aad6c9517b4..647f6c763d8 100644 --- a/handwritten/spanner/test/transaction.ts +++ b/handwritten/spanner/test/transaction.ts @@ -171,6 +171,25 @@ describe('Transaction', () => { assert.strictEqual(REQUEST_STREAM.callCount, 1); }); + it('should keep `request` and `requestStream` usable when detached', () => { + REQUEST.resetHistory(); + REQUEST_STREAM.resetHistory(); + const multiplexedSession = Object.assign({}, SESSION, { + metadata: {multiplexed: true}, + }); + const txn = new Snapshot(multiplexedSession); + + // `TransactionRunner#_interceptErrors` and user code hold on to these + // methods without their receiver, so they must be pre-bound. + const {request, requestStream} = txn; + + request({client: 'SpannerClient'}, () => {}); + requestStream({client: 'SpannerClient'}); + + assert.strictEqual(REQUEST.callCount, 1); + assert.strictEqual(REQUEST_STREAM.callCount, 1); + }); + it('should generate _affinityKey for multiplexed sessions', () => { const multiplexedSession = Object.assign({}, SESSION, { metadata: {multiplexed: true}, @@ -214,6 +233,33 @@ describe('Transaction', () => { assert.deepStrictEqual(arg.gaxOpts, txn._bindGaxOpts); }); + it('should merge the affinity key into caller supplied gaxOpts', () => { + REQUEST.resetHistory(); + const multiplexedSession = Object.assign({}, SESSION, { + metadata: {multiplexed: true}, + }); + const txn = new Snapshot(multiplexedSession); + const gaxOpts = { + timeout: 1000, + otherArgs: {options: {unbind: true}}, + }; + + txn.request({client: 'SpannerClient', gaxOpts}, () => {}); + + const arg = REQUEST.lastCall.args[0]; + assert.deepStrictEqual(arg.gaxOpts, { + timeout: 1000, + otherArgs: { + options: {unbind: true, affinityKey: txn._affinityKey}, + }, + }); + // The caller supplied gax options must not be modified. + assert.deepStrictEqual(gaxOpts, { + timeout: 1000, + otherArgs: {options: {unbind: true}}, + }); + }); + it('should set the commonHeaders_', () => { assert.deepStrictEqual(snapshot.commonHeaders_, { [CLOUD_RESOURCE_HEADER]: snapshot.session.parent.formattedName_,