chore(spanner): clean up leader routing headers and fix commonHeaders_ mutation

handwritten/spanner/src/database.ts

@@ -719,16 +719,14 @@
       sessionCount: count,
     };
 
-    const headers = this.commonHeaders_;
-    if (this._getSpanner().routeToLeaderEnabled) {
-      addLeaderAwareRoutingHeader(headers);
-    }
-
     const allHeaders = this._metadataWithRequestId(
       this._nextNthRequest(),
       1,
-      headers,
+      this.commonHeaders_,
     );
+    if (this._getSpanner().routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(allHeaders);
+    }
 
     startTrace('Database.batchCreateSessions', this._traceConfig, span => {
       this.request<google.spanner.v1.IBatchCreateSessionsResponse>(
@@ -2483,11 +2481,11 @@
     if (gaxOpts) {
       const gax = gaxOpts as GetDatabaseRolesOptions;
       if (gax.pageSize !== undefined) {
         (reqOpts as any).pageSize ??= gax.pageSize;
         delete gax.pageSize;
       }
       if (gax.pageToken !== undefined) {
         (reqOpts as any).pageToken ??= gax.pageToken;
         delete gax.pageToken;
       }
     }

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,10 +438,10 @@
           },
         },
       };
       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,
@@ -453,10 +453,10 @@
         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,
@@ -697,15 +697,6 @@
       reqOpts.requestOptions = this.requestOptions;
     }
 
-    const headers = this.commonHeaders_;
-    if (
-      this._getSpanner().routeToLeaderEnabled &&
-      (this._options.readWrite !== undefined ||
-        this._options.partitionedDml !== undefined)
-    ) {
-      addLeaderAwareRoutingHeader(headers);
-    }
-
     return startTrace(
       'Snapshot.begin',
       {
@@ -721,7 +712,10 @@
             method: 'beginTransaction',
             reqOpts,
             gaxOpts,
-            headers: injectRequestIDIntoHeaders(headers, this.session),
+            headers: injectRequestIDIntoHeaders(
+              this.commonHeaders_,
+              this.session,
+            ),
           },
           (
             err: null | grpc.ServiceError,
@@ -974,15 +968,6 @@
       },
     );
 
-    const headers = this.commonHeaders_;
-    if (
-      this._getSpanner().routeToLeaderEnabled &&
-      (this._options.readWrite !== undefined ||
-        this._options.partitionedDml !== undefined)
-    ) {
-      addLeaderAwareRoutingHeader(headers);
-    }
-
     const traceConfig: traceConfig = {
       ...this._traceConfig,
       tableName: table,
@@ -1021,7 +1006,7 @@
           reqOpts: Object.assign({}, reqOpts, {resumeToken}),
           gaxOpts: gaxOptions,
           headers: injectRequestIDIntoHeaders(
-            headers,
+            this.commonHeaders_,
             this.session,
             nthRequest,
             attempt,
@@ -1137,7 +1122,7 @@
     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)
@@ -1525,15 +1510,6 @@
       });
     };
 
-    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,
@@ -1623,7 +1599,7 @@
         }
 
         const injectedHeaders = injectRequestIDIntoHeaders(
-          headers,
+          this.commonHeaders_,
           this.session,
           nthRequest,
           attempt,
@@ -1983,15 +1959,6 @@
       });
     };
 
-    const headers = this.commonHeaders_;
-    if (
-      this._getSpanner().routeToLeaderEnabled &&
-      (this._options.readWrite !== undefined ||
-        this._options.partitionedDml !== undefined)
-    ) {
-      addLeaderAwareRoutingHeader(headers);
-    }
-
     const traceConfig: traceConfig = {
       transactionTag: this.requestOptions?.transactionTag,
       requestTag: requestOptions?.requestTag,
@@ -2036,7 +2003,7 @@
           reqOpts: Object.assign({}, reqOpts, {resumeToken}),
           gaxOpts: gaxOptions,
           headers: injectRequestIDIntoHeaders(
-            headers,
+            this.commonHeaders_,
             this.session,
             nthRequest,
             attempt,
@@ -2349,7 +2316,7 @@
    * @returns {Spanner}
    */
   protected _getSpanner(): Spanner {
-    return this.session.parent.parent.parent as Spanner;
+    return (this.session?.parent as Database)?.parent?.parent as Spanner;
   }
 }
 
@@ -2553,6 +2520,9 @@
     this._options.isolationLevel = IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED;
     this.requestOptions = requestOptions;
     this._retryCommit = false;
+    if (this._getSpanner()?.routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(this.commonHeaders_);
+    }
   }
 
   /**
@@ -2707,9 +2677,6 @@
       nextNthRequest(database),
       1,
     );
-    if (this._getSpanner().routeToLeaderEnabled) {
-      addLeaderAwareRoutingHeader(headers);
-    }
 
     const traceConfig: traceConfig = {
       ...this._traceConfig,
@@ -2946,11 +2913,6 @@
           this.requestOptions,
         );
 
-        const headers = this.commonHeaders_;
-        if (this._getSpanner().routeToLeaderEnabled) {
-          addLeaderAwareRoutingHeader(headers);
-        }
-
         span.addEvent('Starting Commit');
 
         const database = this.session.parent as Database;
@@ -2969,7 +2931,7 @@
             reqOpts,
             gaxOpts,
             headers: injectRequestIDIntoHeaders(
-              headers,
+              this.commonHeaders_,
               this.session,
               nextNthRequest(database),
               1,
@@ -3341,11 +3303,6 @@
         transactionId,
       };
 
-      const headers = this.commonHeaders_;
-      if (this._getSpanner().routeToLeaderEnabled) {
-        addLeaderAwareRoutingHeader(headers);
-      }
-
       if (this._affinityKey) {
         if (!gaxOpts || Object.keys(gaxOpts).length === 0) {
           gaxOpts = this._unbindGaxOpts as any;
@@ -3360,7 +3317,7 @@
           method: 'rollback',
           reqOpts,
           gaxOpts,
-          headers: headers,
+          headers: {...this.commonHeaders_},
         },
         (err: null | ServiceError) => {
           if (err) {
@@ -3894,6 +3851,9 @@
   ) {
     super(session);
     this._options = {partitionedDml: options};
+    if (this._getSpanner()?.routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(this.commonHeaders_);
+    }
   }
   /**
    * Use option excludeTxnFromChangeStreams to exclude partitionedDml

handwritten/spanner/test/database.ts

@@ -489,14 +489,41 @@ describe('Database', () => {
         headers,
         Object.assign(
           {
-            [LEADER_AWARE_ROUTING_HEADER]: true,
+            [LEADER_AWARE_ROUTING_HEADER]: 'true',
             [X_GOOG_SPANNER_REQUEST_ID_HEADER]: craftRequestId(1, 1, 1, 1),
           },
           database.commonHeaders_,
         ),
       );
     });
 
+    it('should not mutate commonHeaders_', () => {
+      sandbox.stub(database, 'request');
+      const expectedCommonHeaders = Object.assign({}, database.commonHeaders_);
+
+      database.batchCreateSessions({count: 10}, assert.ifError);
+
+      assert.deepStrictEqual(database.commonHeaders_, expectedCommonHeaders);
+      assert.strictEqual(
+        database.commonHeaders_[LEADER_AWARE_ROUTING_HEADER],
+        undefined,
+      );
+    });
+
+    it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => {
+      const stub = sandbox.stub(database, 'request');
+      database.instance = {
+        parent: {
+          routeToLeaderEnabled: false,
+        },
+      } as {} as Instance;
+
+      database.batchCreateSessions({count: 10}, assert.ifError);
+
+      const {headers} = stub.lastCall.args[0];
+      assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined);
+    });
+
     it('should accept just a count number', () => {
       const stub = sandbox.stub(database, 'request');
       const count = 10;