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

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);

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();