feat(firestore): add x-goog-is-transaction routing header

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<Resp> {
-    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<Duplex> {
-    const callOptions = this.createCallOptions(methodName);
+    const callOptions = this.createCallOptions(methodName, undefined, request);
 
     const bidirectional = methodName === 'listen';
     let numResponses = 0;

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', () => {

handwritten/firestore/dev/test/transaction.ts

@@ -374,6 +374,9 @@ function runTransaction<T>(
       // 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<T>(
         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<T>(
         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