diff --git a/handwritten/firestore/dev/src/index.ts b/handwritten/firestore/dev/src/index.ts index 9889dd797090..ff9c710d810f 100644 --- a/handwritten/firestore/dev/src/index.ts +++ b/handwritten/firestore/dev/src/index.ts @@ -181,6 +181,11 @@ setLibVersion(libVersion); */ const CLOUD_RESOURCE_HEADER = 'google-cloud-resource-prefix'; +/*! + * HTTP header to indicate whether an RPC is part of a transaction. + */ +const IS_TRANSACTION_HEADER = 'x-goog-is-transaction'; + /** * The maximum number of times to retry idempotent requests. * @private @@ -1624,11 +1629,24 @@ export class Firestore implements firestore.Firestore { private createCallOptions( methodName: string, retryCodes?: number[], + request?: { + transaction?: Uint8Array | null; + newTransaction?: api.ITransactionOptions | null; + }, ): CallOptions { + const isTransaction = + methodName === 'beginTransaction' || + methodName === 'rollback' || + Boolean( + request && + ((request.transaction && request.transaction.length > 0) || + request.newTransaction), + ); const callOptions: CallOptions = { otherArgs: { headers: { [CLOUD_RESOURCE_HEADER]: this.formattedName, + [IS_TRANSACTION_HEADER]: isTransaction ? 'true' : 'false', ...this._settings.customHeaders, ...this._settings[methodName]?.customHeaders, }, @@ -1841,7 +1859,14 @@ export class Firestore implements firestore.Firestore { requestTag: string, retryCodes?: number[], ): Promise { - const callOptions = this.createCallOptions(methodName, retryCodes); + const callOptions = this.createCallOptions( + methodName, + retryCodes, + request as { + transaction?: Uint8Array | null; + newTransaction?: api.ITransactionOptions | null; + }, + ); return this._clientPool.run( requestTag, @@ -1895,7 +1920,7 @@ export class Firestore implements firestore.Firestore { request: {}, requestTag: string, ): Promise { - const callOptions = this.createCallOptions(methodName); + const callOptions = this.createCallOptions(methodName, undefined, request); const bidirectional = methodName === 'listen'; let numResponses = 0; diff --git a/handwritten/firestore/dev/test/index.ts b/handwritten/firestore/dev/test/index.ts index 54e6822ec652..d3237efaf694 100644 --- a/handwritten/firestore/dev/test/index.ts +++ b/handwritten/firestore/dev/test/index.ts @@ -1073,6 +1073,43 @@ describe('instantiation', () => { ); }); }); + + it('sets x-goog-is-transaction routing header', async () => { + let commitHeaders: {[k: string]: string} | undefined; + let batchGetHeaders: {[k: string]: string} | undefined; + const overrides: ApiOverride = { + commit: (_request, options) => { + commitHeaders = options?.otherArgs?.headers; + return response({ + commitTime: {}, + writeResults: [{updateTime: {}}], + }); + }, + batchGetDocuments: (_request, options) => { + batchGetHeaders = options?.otherArgs?.headers; + return stream( + {transaction: Buffer.from('tx1')}, + found(document('documentId')), + ); + }, + }; + + const firestore = await createInstance(overrides); + + await firestore + .collection('collectionId') + .doc('documentId') + .set({foo: 'bar'}); + expect(commitHeaders?.['x-goog-is-transaction']).to.equal('false'); + + await firestore.runTransaction(async transaction => { + const docRef = firestore.doc('collectionId/documentId'); + await transaction.get(docRef); + transaction.set(docRef, {foo: 'baz'}); + }); + expect(batchGetHeaders?.['x-goog-is-transaction']).to.equal('true'); + expect(commitHeaders?.['x-goog-is-transaction']).to.equal('true'); + }); }); describe('serializer', () => { diff --git a/handwritten/firestore/dev/test/transaction.ts b/handwritten/firestore/dev/test/transaction.ts index 175055ada1b3..8398a5b21f3e 100644 --- a/handwritten/firestore/dev/test/transaction.ts +++ b/handwritten/firestore/dev/test/transaction.ts @@ -374,6 +374,9 @@ function runTransaction( // Ensure that we do not specify custom retry behavior for transactional // commits. expect(options!.retry).to.be.undefined; + expect(options!.otherArgs!.headers!['x-goog-is-transaction']).to.equal( + actual.transaction && actual.transaction.length > 0 ? 'true' : 'false', + ); const request = expectedRequests.shift()!; expect(request.type).to.equal('commit'); @@ -384,7 +387,10 @@ function runTransaction( return response(request.response as api.ICommitResponse); } }, - rollback: actual => { + rollback: (actual, options) => { + expect(options!.otherArgs!.headers!['x-goog-is-transaction']).to.equal( + 'true', + ); const request = expectedRequests.shift()!; expect(request.type).to.equal('rollback'); expect(actual).to.deep.eq(request.request); @@ -394,13 +400,25 @@ function runTransaction( return response({}); } }, - batchGetDocuments: actual => { + batchGetDocuments: (actual, options) => { + expect(options!.otherArgs!.headers!['x-goog-is-transaction']).to.equal( + (actual!.transaction && actual!.transaction.length > 0) || + actual!.newTransaction + ? 'true' + : 'false', + ); const request = expectedRequests.shift()!; expect(request.type).to.equal('getDocument'); expect(actual).to.deep.eq(request.request); return request.stream!; }, - runQuery: actual => { + runQuery: (actual, options) => { + expect(options!.otherArgs!.headers!['x-goog-is-transaction']).to.equal( + (actual!.transaction && actual!.transaction.length > 0) || + actual!.newTransaction + ? 'true' + : 'false', + ); const request = expectedRequests.shift()!; expect(request.type).to.equal('query'); actual = extend(true, {}, actual); // Remove undefined properties