From b1f43e9fd72901795256cd4ffd7308472bb14ae5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 14 Sep 2026 10:16:30 +0200 Subject: [PATCH] chore(spanner): clean up leader routing headers and fix commonHeaders_ mutation - Precomputes 'x-goog-spanner-route-to-leader: true' once in Transaction and PartitionedDml constructors when routeToLeaderEnabled is true. - Removes redundant per-request header mutations and option checks across Snapshot and Transaction methods (begin, read, runStream, batchUpdate, commit, rollback). - Uses safe navigation in Snapshot.prototype._getSpanner to support lightweight sessions without full parent chains. - Fixes latent in-place mutation in Database.prototype.batchCreateSessions where database.commonHeaders_ was mutated, which caused subsequent read-only snapshots to inadvertently route to the leader. - Passes a shallow copy of commonHeaders_ in Transaction.prototype.rollback to protect against downstream in-place header modifications. --- handwritten/spanner/src/database.ts | 10 +- handwritten/spanner/src/transaction.ts | 72 ++----- handwritten/spanner/test/database.ts | 29 ++- handwritten/spanner/test/transaction.ts | 275 +++++++++++++++++++++++- 4 files changed, 321 insertions(+), 65 deletions(-) diff --git a/handwritten/spanner/src/database.ts b/handwritten/spanner/src/database.ts index b6ba8347e4c5..40b203037188 100644 --- a/handwritten/spanner/src/database.ts +++ b/handwritten/spanner/src/database.ts @@ -719,16 +719,14 @@ class Database extends common.GrpcServiceObject { 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( diff --git a/handwritten/spanner/src/transaction.ts b/handwritten/spanner/src/transaction.ts index 490406d9fa95..4fc60166b493 100644 --- a/handwritten/spanner/src/transaction.ts +++ b/handwritten/spanner/src/transaction.ts @@ -697,15 +697,6 @@ export class Snapshot extends EventEmitter { 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 @@ export class Snapshot extends EventEmitter { method: 'beginTransaction', reqOpts, gaxOpts, - headers: injectRequestIDIntoHeaders(headers, this.session), + headers: injectRequestIDIntoHeaders( + this.commonHeaders_, + this.session, + ), }, ( err: null | grpc.ServiceError, @@ -974,15 +968,6 @@ export class Snapshot extends EventEmitter { }, ); - 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 @@ export class Snapshot extends EventEmitter { reqOpts: Object.assign({}, reqOpts, {resumeToken}), gaxOpts: gaxOptions, headers: injectRequestIDIntoHeaders( - headers, + this.commonHeaders_, this.session, nthRequest, attempt, @@ -1525,15 +1510,6 @@ export class Snapshot extends EventEmitter { }); }; - 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 @@ export class Snapshot extends EventEmitter { } const injectedHeaders = injectRequestIDIntoHeaders( - headers, + this.commonHeaders_, this.session, nthRequest, attempt, @@ -1983,15 +1959,6 @@ export class Snapshot extends EventEmitter { }); }; - 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 @@ export class Snapshot extends EventEmitter { reqOpts: Object.assign({}, reqOpts, {resumeToken}), gaxOpts: gaxOptions, headers: injectRequestIDIntoHeaders( - headers, + this.commonHeaders_, this.session, nthRequest, attempt, @@ -2349,7 +2316,7 @@ export class Snapshot extends EventEmitter { * @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 @@ export class Transaction extends Dml { this._options.isolationLevel = IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED; this.requestOptions = requestOptions; this._retryCommit = false; + if (this._getSpanner()?.routeToLeaderEnabled) { + addLeaderAwareRoutingHeader(this.commonHeaders_); + } } /** @@ -2707,9 +2677,6 @@ export class Transaction extends Dml { nextNthRequest(database), 1, ); - if (this._getSpanner().routeToLeaderEnabled) { - addLeaderAwareRoutingHeader(headers); - } const traceConfig: traceConfig = { ...this._traceConfig, @@ -2946,11 +2913,6 @@ export class Transaction extends Dml { 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 @@ export class Transaction extends Dml { reqOpts, gaxOpts, headers: injectRequestIDIntoHeaders( - headers, + this.commonHeaders_, this.session, nextNthRequest(database), 1, @@ -3341,11 +3303,6 @@ export class Transaction extends Dml { 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 @@ export class Transaction extends Dml { method: 'rollback', reqOpts, gaxOpts, - headers: headers, + headers: {...this.commonHeaders_}, }, (err: null | ServiceError) => { if (err) { @@ -3894,6 +3851,9 @@ export class PartitionedDml extends Dml { ) { super(session); this._options = {partitionedDml: options}; + if (this._getSpanner()?.routeToLeaderEnabled) { + addLeaderAwareRoutingHeader(this.commonHeaders_); + } } /** * Use option excludeTxnFromChangeStreams to exclude partitionedDml diff --git a/handwritten/spanner/test/database.ts b/handwritten/spanner/test/database.ts index 1539b572de0e..22b93823491e 100644 --- a/handwritten/spanner/test/database.ts +++ b/handwritten/spanner/test/database.ts @@ -489,7 +489,7 @@ 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_, @@ -497,6 +497,33 @@ describe('Database', () => { ); }); + 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; diff --git a/handwritten/spanner/test/transaction.ts b/handwritten/spanner/test/transaction.ts index 4d0be66e437a..911ce315f01e 100644 --- a/handwritten/spanner/test/transaction.ts +++ b/handwritten/spanner/test/transaction.ts @@ -41,8 +41,8 @@ import { BatchUpdateOptions, ExecuteSqlRequest, ReadRequest, - RunCallback, } from '../src/transaction'; +import {Session} from '../src/session'; import {Row} from '../src/partial-result-stream'; import {grpc} from 'google-gax'; import * as through from 'through2'; @@ -85,6 +85,21 @@ describe('Transaction', () => { }, }; + function createSession(options?: {routeToLeaderEnabled?: boolean}): Session { + return Object.assign({}, SESSION, { + parent: { + ...DATABASE, + parent: { + ...INSTANCE, + parent: { + ...SPANNER, + routeToLeaderEnabled: options?.routeToLeaderEnabled, + }, + }, + }, + }) as unknown as Session; + } + const PARTIAL_RESULT_STREAM = sandbox.stub(); const PROMISIFY_ALL = sandbox.stub(); @@ -220,6 +235,22 @@ describe('Transaction', () => { [AFE_SERVER_TIMING_HEADER]: 'true', }); }); + + it('should not include leader-aware routing header in commonHeaders_ even when routeToLeaderEnabled is true', () => { + const s = new Snapshot(SESSION); + assert.strictEqual( + s.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should not include leader-aware routing header in commonHeaders_ when routeToLeaderEnabled is false', () => { + const s = new Snapshot(createSession({routeToLeaderEnabled: false})); + assert.strictEqual( + s.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); }); describe('begin', () => { @@ -240,6 +271,13 @@ describe('Transaction', () => { assert.deepStrictEqual(headers, snapshot.commonHeaders_); }); + it('should not include leader-aware routing header in beginTransaction request', () => { + snapshot.begin(); + + const {headers} = REQUEST.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should accept gaxOptions', done => { const gaxOptions = { timeout: 1000, @@ -362,6 +400,21 @@ describe('Transaction', () => { }); }); + it('should not include leader-aware routing header in single-use read', () => { + snapshot.createReadStream(TABLE); + + const {headers} = REQUEST_STREAM.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + + it('should not include leader-aware routing header in multi-use read', () => { + snapshot.id = 'transaction-id-123'; + snapshot.createReadStream(TABLE); + + const {headers} = REQUEST_STREAM.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should use the transaction id if present', () => { const id = 'transaction-id-123'; const expectedTransaction = {id}; @@ -1548,6 +1601,21 @@ describe('Transaction', () => { }); }); + it('should not include leader-aware routing header in single-use read-only query', () => { + snapshot.runStream(QUERY); + + const {headers} = REQUEST_STREAM.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + + it('should not include leader-aware routing header in multi-use read-only query', () => { + snapshot.id = 'transaction-id-123'; + snapshot.runStream(QUERY); + + const {headers} = REQUEST_STREAM.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should use the transaction id if present', () => { const id = 'transaction-id-123'; const expectedTransaction = {id}; @@ -2222,6 +2290,63 @@ describe('Transaction', () => { it('should inherit from Dml', () => { assert(transaction instanceof Dml); }); + + it('should precompute leader-aware routing header in commonHeaders_ when routeToLeaderEnabled is true', () => { + const txn = new Transaction(SESSION); + assert.strictEqual( + txn.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + 'true', + ); + }); + + it('should not mutate session.commonHeaders_ when routeToLeaderEnabled is true', () => { + const session = createSession({routeToLeaderEnabled: true}); + const originalHeaders = {...session.commonHeaders_}; + const txn = new Transaction(session); + assert.strictEqual( + txn.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + 'true', + ); + assert.deepStrictEqual(session.commonHeaders_, originalHeaders); + assert.strictEqual( + session.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should not include leader-aware routing header in commonHeaders_ when routeToLeaderEnabled is false', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + assert.strictEqual( + txn.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should not include leader-aware routing header when routeToLeaderEnabled is undefined', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: undefined}), + ); + assert.strictEqual( + txn.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should safely initialize when session lacks parent hierarchy', () => { + const mockSession = { + request: sandbox.stub(), + requestStream: sandbox.stub(), + parent: {formattedName_: 'database-name'}, + commonHeaders_: {}, + }; + const txn = new Transaction(mockSession as unknown as Session); + assert.strictEqual( + txn.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); }); describe('batchUpdate', () => { @@ -2380,6 +2505,17 @@ describe('Transaction', () => { ); }); + it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + const stub = sandbox.stub(txn, 'request'); + txn.batchUpdate(STRING_STATEMENTS, assert.ifError); + + const {headers} = stub.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should encode sql string statements', () => { const stub = sandbox.stub(transaction, 'request'); const expectedStatements = STRING_STATEMENTS.map(sql => ({sql})); @@ -2544,6 +2680,17 @@ describe('Transaction', () => { ); }); + it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + const stub = sandbox.stub(txn, 'request'); + txn.begin(); + + const {headers} = stub.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should accept gaxOptions', done => { const gaxOptions = {}; transaction.request = config => { @@ -2761,6 +2908,17 @@ describe('Transaction', () => { ); }); + it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + const stub = sandbox.stub(txn, 'request'); + txn.commit(); + + const {headers} = stub.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should inject _unbindGaxOpts for commit if _affinityKey is present', () => { const stub = sandbox.stub(transaction, 'request'); (transaction as any)._affinityKey = 'mux-affinity-1'; @@ -3384,10 +3542,23 @@ describe('Transaction', () => { assert.deepStrictEqual( headers, Object.assign( - {[LEADER_AWARE_ROUTING_HEADER]: true}, + {[LEADER_AWARE_ROUTING_HEADER]: 'true'}, transaction.commonHeaders_, ), ); + assert.notStrictEqual(headers, transaction.commonHeaders_); + }); + + it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + const stub = sandbox.stub(txn, 'request'); + txn.id = 'transaction-id-123'; + txn.rollback(); + + const {headers} = stub.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); }); it('should inject _unbindGaxOpts for rollback if _affinityKey is present', () => { @@ -4021,6 +4192,25 @@ describe('Transaction', () => { transaction.runStream(QUERY); }); + it('should not include leader-aware routing header when routeToLeaderEnabled is false', done => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + const QUERY: ExecuteSqlRequest = { + sql: 'SELET * FROM `MyTable`', + }; + + txn.requestStream = config => { + assert.strictEqual( + config.headers[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + done(); + }; + + txn.runStream(QUERY); + }); + it('should set transaction tag when not `singleUse`', done => { const QUERY: ExecuteSqlRequest = { sql: 'SELET * FROM `MyTable`', @@ -4093,6 +4283,17 @@ describe('Transaction', () => { ); }); + it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => { + const txn = new Transaction( + createSession({routeToLeaderEnabled: false}), + ); + const TABLE = 'my-table-123'; + txn.createReadStream(TABLE); + + const {headers} = REQUEST_STREAM.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); + it('should set transaction tag if not `singleUse`', () => { const TABLE = 'my-table-123'; const transactionTag = 'bar'; @@ -4229,6 +4430,65 @@ describe('Transaction', () => { it('should inherit from Dml', () => { assert(pdml instanceof Dml); }); + + it('should precompute leader-aware routing header in commonHeaders_ when routeToLeaderEnabled is true', () => { + const partitionedDml = new PartitionedDml(SESSION); + assert.strictEqual( + partitionedDml.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + 'true', + ); + }); + + it('should not mutate session.commonHeaders_ when routeToLeaderEnabled is true', () => { + const session = createSession({routeToLeaderEnabled: true}); + const originalHeaders = {...session.commonHeaders_}; + const partitionedDml = new PartitionedDml(session); + assert.strictEqual( + partitionedDml.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + 'true', + ); + assert.deepStrictEqual(session.commonHeaders_, originalHeaders); + assert.strictEqual( + session.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should not include leader-aware routing header in commonHeaders_ when routeToLeaderEnabled is false', () => { + const partitionedDml = new PartitionedDml( + createSession({routeToLeaderEnabled: false}), + ); + assert.strictEqual( + partitionedDml.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should not include leader-aware routing header when routeToLeaderEnabled is undefined', () => { + const partitionedDml = new PartitionedDml( + createSession({routeToLeaderEnabled: undefined}), + ); + assert.strictEqual( + partitionedDml.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); + + it('should safely initialize when session lacks parent hierarchy', () => { + const mockSession = { + request: sandbox.stub(), + requestStream: sandbox.stub(), + parent: {formattedName_: 'database-name'}, + commonHeaders_: {}, + }; + const partitionedDml = new PartitionedDml( + mockSession as unknown as Session, + ); + assert.strictEqual( + partitionedDml.commonHeaders_[LEADER_AWARE_ROUTING_HEADER], + undefined, + ); + }); }); describe('begin', () => { @@ -4249,6 +4509,17 @@ describe('Transaction', () => { ), ); }); + + it('should not include leader-aware routing header when routeToLeaderEnabled is false', () => { + const partitionedDml = new PartitionedDml( + createSession({routeToLeaderEnabled: false}), + ); + const stub = sandbox.stub(partitionedDml, 'request'); + partitionedDml.begin(); + + const {headers} = stub.lastCall.args[0]; + assert.strictEqual(headers[LEADER_AWARE_ROUTING_HEADER], undefined); + }); }); describe('runUpdate', () => {