From 0ee4527cc4f9cd27757195feca931a0058cb085a Mon Sep 17 00:00:00 2001 From: Nikhil Chennupati Date: Sat, 24 Oct 2026 06:59:15 -0700 Subject: [PATCH] feat(datastore): add x-goog-is-transaction routing header Adds the `x-goog-is-transaction: true` HTTP header to all transactional Datastore RPCs, ensuring proper routing of transactional requests. Changes: - Modified `prepareGaxRequest_` in `request.ts` to calculate `isTransactionHeader` based on the RPC method (`beginTransaction`, `rollback`, or `commit`, `runQuery`, etc.) and the transaction ID presence. - Appended the header `x-goog-is-transaction` to GAX options for all calls (set to `true` or `false`). - Added exhaustive unit tests in `request.ts` that verify the `x-goog-is-transaction` header behavior for both transactional and non-transactional RPC calls. --- handwritten/datastore/src/request.ts | 52 +++++++++++- handwritten/datastore/test/request.ts | 113 ++++++++++++++++++++++++++ 2 files changed, 162 insertions(+), 3 deletions(-) diff --git a/handwritten/datastore/src/request.ts b/handwritten/datastore/src/request.ts index e78ec19f62df..fa087b1c77e7 100644 --- a/handwritten/datastore/src/request.ts +++ b/handwritten/datastore/src/request.ts @@ -1232,7 +1232,11 @@ class DatastoreRequest { prepareGaxRequest_(config: RequestConfig, callback: Function): void { const datastore = this.datastore; - const isTransaction = this.id ? true : false; + const hasTransactionBytes = (val?: string | Uint8Array | null): boolean => + val instanceof Uint8Array || Buffer.isBuffer(val) + ? val.length > 0 + : Boolean(val); + const isTransaction = hasTransactionBytes(this.id); const method = config.method; const reqOpts = extend(true, {}, config.reqOpts); @@ -1241,8 +1245,15 @@ class DatastoreRequest { if (isTransaction) { reqOpts.mode = 'TRANSACTIONAL'; reqOpts.transaction = this.id; - } else { - reqOpts.mode = 'NON_TRANSACTIONAL'; + } else if (!reqOpts.mode) { + reqOpts.mode = + hasTransactionBytes(reqOpts.transaction) || + Boolean( + (reqOpts as any).singleUseTransaction || + (reqOpts as any).single_use_transaction, + ) + ? 'TRANSACTIONAL' + : 'NON_TRANSACTIONAL'; } } @@ -1283,9 +1294,44 @@ class DatastoreRequest { } const gaxClient = datastore.clients_.get(clientName); reqOpts.projectId = projectId!; + + const isTransactionalRequest = (): boolean => { + switch (method) { + case 'beginTransaction': + case 'rollback': + return true; + case 'commit': { + const mode = reqOpts.mode as unknown; + const isTransactionalMode = + mode === 'TRANSACTIONAL' || + mode === google.datastore.v1.CommitRequest.Mode.TRANSACTIONAL; + return ( + isTransactionalMode || + hasTransactionBytes(reqOpts.transaction) || + Boolean( + (reqOpts as any).singleUseTransaction || + (reqOpts as any).single_use_transaction, + ) + ); + } + case 'lookup': + case 'runQuery': + case 'runAggregationQuery': + return Boolean( + reqOpts.readOptions && + (hasTransactionBytes(reqOpts.readOptions.transaction) || + reqOpts.readOptions.newTransaction || + (reqOpts.readOptions as any).new_transaction), + ); + default: + return false; + } + }; + const gaxOpts = extend(true, {}, config.gaxOpts, { headers: { 'google-cloud-resource-prefix': `projects/${projectId}`, + 'x-goog-is-transaction': isTransactionalRequest() ? 'true' : 'false', }, }); const requestFn = gaxClient![method].bind(gaxClient, reqOpts, gaxOpts); diff --git a/handwritten/datastore/test/request.ts b/handwritten/datastore/test/request.ts index 97366be73bbe..c7dd423278b4 100644 --- a/handwritten/datastore/test/request.ts +++ b/handwritten/datastore/test/request.ts @@ -1697,6 +1697,7 @@ describe('Request', () => { [CONFIG.method](_: object, gaxO: any) { assert.deepStrictEqual(gaxO.headers, { 'google-cloud-resource-prefix': 'projects/' + PROJECT_ID, + 'x-goog-is-transaction': 'false', }); done(); }, @@ -1708,6 +1709,118 @@ describe('Request', () => { }); }); + describe('x-goog-is-transaction header', () => { + const methodsForTrue = [ + {method: 'beginTransaction', reqOpts: {}}, + {method: 'rollback', reqOpts: {}}, + {method: 'commit', reqOpts: {mode: 'TRANSACTIONAL'}}, + { + method: 'commit', + reqOpts: {mode: google.datastore.v1.CommitRequest.Mode.TRANSACTIONAL}, + }, + {method: 'commit', reqOpts: {transaction: 'abc'}}, + {method: 'commit', reqOpts: {transaction: Buffer.from('abc')}}, + {method: 'commit', reqOpts: {singleUseTransaction: {}}}, + {method: 'commit', reqOpts: {single_use_transaction: {}}}, + {method: 'lookup', reqOpts: {readOptions: {transaction: 'abc'}}}, + { + method: 'lookup', + reqOpts: {readOptions: {transaction: Buffer.from('abc')}}, + }, + {method: 'lookup', reqOpts: {readOptions: {newTransaction: {}}}}, + {method: 'lookup', reqOpts: {readOptions: {new_transaction: {}}}}, + {method: 'runQuery', reqOpts: {readOptions: {transaction: 'abc'}}}, + {method: 'runQuery', reqOpts: {readOptions: {newTransaction: {}}}}, + {method: 'runQuery', reqOpts: {readOptions: {new_transaction: {}}}}, + { + method: 'runAggregationQuery', + reqOpts: {readOptions: {transaction: 'abc'}}, + }, + { + method: 'runAggregationQuery', + reqOpts: {readOptions: {newTransaction: {}}}, + }, + { + method: 'runAggregationQuery', + reqOpts: {readOptions: {new_transaction: {}}}, + }, + { + method: 'commit', + reqOpts: { + mode: google.datastore.v1.CommitRequest.Mode.NON_TRANSACTIONAL, + transaction: 'abc', + }, + }, + ]; + + const methodsForFalse = [ + {method: 'commit', reqOpts: {mode: 'NON_TRANSACTIONAL'}}, + { + method: 'commit', + reqOpts: { + mode: google.datastore.v1.CommitRequest.Mode.NON_TRANSACTIONAL, + }, + }, + {method: 'commit', reqOpts: {transaction: Buffer.alloc(0)}}, + {method: 'lookup', reqOpts: {}}, + { + method: 'lookup', + reqOpts: {readOptions: {transaction: Buffer.alloc(0)}}, + }, + {method: 'runQuery', reqOpts: {}}, + {method: 'runAggregationQuery', reqOpts: {}}, + {method: 'allocateIds', reqOpts: {}}, + {method: 'reserveIds', reqOpts: {}}, + ]; + + methodsForTrue.forEach(({method, reqOpts}) => { + it(`should be true for ${method} with ${JSON.stringify(reqOpts)}`, done => { + request.datastore.clients_ = new Map(); + request.datastore.clients_.set(CONFIG.client, { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + [method](_: object, gaxO: any) { + assert.strictEqual(gaxO.headers['x-goog-is-transaction'], 'true'); + done(); + }, + }); + const config = Object.assign({}, CONFIG, { + method, + reqOpts, + }); + request.id = undefined; + request.prepareGaxRequest_( + config, + (err: Error, requestFn: Function) => { + assert.ifError(err); + requestFn(); + }, + ); + }); + }); + + methodsForFalse.forEach(({method, reqOpts}) => { + it(`should be false for ${method} with ${JSON.stringify(reqOpts)}`, done => { + request.datastore.clients_ = new Map(); + request.datastore.clients_.set(CONFIG.client, { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + [method](_: object, gaxO: any) { + assert.strictEqual(gaxO.headers['x-goog-is-transaction'], 'false'); + done(); + }, + }); + const config = Object.assign({}, CONFIG, { + method, + reqOpts, + }); + request.id = undefined; // Ensure no transaction ID for non-transactional + request.prepareGaxRequest_(config, (err: Error, requestFn: Function) => { + assert.ifError(err); + requestFn(); + }); + }); + }); + }); + describe('commit', () => { it('should set the mode', done => { request.datastore.clients_ = new Map();