chore(spanner): de-duplicate multiplexed session affinity gax options

handwritten/spanner/src/transaction.ts

@@ -79,7 +79,7 @@
  * Injects a key-value pair into the gaxOpts.otherArgs.options object
  * without mutating the original.
  */
 function injectGaxOpt(existingOpts: any, key: string, value: any): any {
   return Object.assign({}, existingOpts, {
     otherArgs: Object.assign({}, existingOpts?.otherArgs, {
       options: Object.assign({}, existingOpts?.otherArgs?.options, {
@@ -438,35 +438,12 @@
           },
         },
       };
-      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 @@
     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 ||
@@ -1137,11 +1146,11 @@
     if (this._affinityKey) {
       const database = this.session?.parent as Database;
       const spanner = database?.parent?.parent as Spanner;
       const client = spanner?.clients_?.get('SpannerClient') as any;
 
       if (client?.spannerStub) {
         Promise.resolve(client.spannerStub)
           .then((stub: any) => {
             return stub?.getChannel?.()?.unbind?.(this._affinityKey);
           })
           .catch(() => {});
@@ -2919,14 +2928,14 @@
                   setSpanError(span, err);
                 }
                 span.end();
                 callback(err, resp);
               });
               return null;
             })
             .catch(err => {
               setSpanError(span, err);
               span.end();
               callback(err, null);
             });
           return;
         }
@@ -2960,7 +2969,7 @@
         const database = this.session.parent as Database;
         if (this._affinityKey) {
           if (!gaxOpts || Object.keys(gaxOpts).length === 0) {
             gaxOpts = this._unbindGaxOpts as any;
           } else {
             gaxOpts = injectGaxOpt(gaxOpts, 'unbind', true);
           }
@@ -3352,7 +3361,7 @@
 
       if (this._affinityKey) {
         if (!gaxOpts || Object.keys(gaxOpts).length === 0) {
           gaxOpts = this._unbindGaxOpts as any;
         } else {
           gaxOpts = injectGaxOpt(gaxOpts, 'unbind', true);
         }

handwritten/spanner/test/transaction.ts

@@ -41,7 +41,7 @@
   BatchUpdateOptions,
   ExecuteSqlRequest,
   ReadRequest,
   RunCallback,
 } from '../src/transaction';
 import {Row} from '../src/partial-result-stream';
 import {grpc} from 'google-gax';
@@ -171,6 +171,25 @@
         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 @@
         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_,