feat: Databoostenabled for Query and Read partitions

src/transaction.ts

@@ -80,6 +80,7 @@ export interface ExecuteSqlRequest extends Statement, RequestOptions {
   seqno?: number;
   queryOptions?: IQueryOptions;
   requestOptions?: Omit<IRequestOptions, 'transactionTag'>;
+  dataBoostEnabled?: boolean | null;
 }
 
 export interface KeyRange {
@@ -100,6 +101,7 @@ export interface ReadRequest extends RequestOptions {
   resumeToken?: Uint8Array | null;
   partitionToken?: Uint8Array | null;
   requestOptions?: Omit<IRequestOptions, 'transactionTag'>;
+  dataBoostEnabled?: boolean | null;
 }
 
 export interface BatchUpdateError extends grpc.ServiceError {

system-test/spanner.ts

@@ -8234,6 +8234,94 @@ describe('Spanner', () => {
         deadlineErrorInsteadOfAbort(done, PG_DATABASE, postgreSqlTable);
       });
     });
+
+    describe('batch transactions', () => {
+      before(done => {
+        if (!IS_EMULATOR_ENABLED) {
+          DATABASE.runTransaction((err, transaction) => {
+            assert.ifError(err);
+
+            transaction!.runUpdate(
+              {
+                sql:
+                  'INSERT INTO ' +
+                  TABLE_NAME +
+                  ' (Key, StringValue) VALUES(@key, @str)',
+                params: {
+                  key: 'k998',
+                  str: 'abc',
+                },
+              },
+              err => {
+                assert.ifError(err);
+                transaction!.commit(done);
+              }
+            );
+          });
+        } else {
+          done();
+        }
+      });
+
+      it('should create and execute a query partition', function (done) {
+        if (IS_EMULATOR_ENABLED) {
+          this.skip();
+        }
+        const selectQuery = {
+          sql: 'SELECT * FROM TxnTable where Key = "k998"',
+        };
+
+        let row_count = 0;
+        DATABASE.createBatchTransaction((err, transaction) => {
+          assert.ifError(err);
+          transaction!.createQueryPartitions(selectQuery, (err, partitions) => {
+            assert.ifError(err);
+            assert.deepStrictEqual(partitions.length, 1);
+            partitions.forEach(partition => {
+              transaction!.execute(partition, (err, results) => {
+                assert.ifError(err);
+                row_count += results.map(row => row.toJSON()).length;
+                assert.deepStrictEqual(row_count, 1);
+                transaction!.close();
+                done();
+              });
+            });
+          });
+        });
+      });
+
+      it('should create and execute a read partition', function (done) {
+        if (IS_EMULATOR_ENABLED) {
+          this.skip();
+        }
+        const key = 'k998';
+        const QUERY = {
+          table: googleSqlTable.name,
+          // Set databoostenabled to true for enabling serveless analytics.
+          dataBoostEnabled: true,
+          keys: [key],
+          columns: ['Key'],
+        };
+
+        let read_row_count = 0;
+        DATABASE.createBatchTransaction((err, transaction) => {
+          assert.ifError(err);
+          transaction!.createReadPartitions(QUERY, (err, partitions) => {
+            assert.ifError(err);
+            assert.deepStrictEqual(partitions.length, 1);
+            partitions.forEach(partition => {
+              transaction!.execute(partition, (err, results) => {
+                assert.ifError(err);
+                read_row_count += results.map(row => row.toJSON()).length;
+                assert.deepStrictEqual(read_row_count, 1);
+                transaction!.close();
+                done();
+              });
+            });
+          });
+        });
+      });
+    });
   });
 });
 

test/batch-transaction.ts

@@ -132,12 +132,14 @@ describe('BatchTransaction', () => {
       gaxOptions: GAX_OPTS,
       params: {},
       types: {},
+      dataBoostEnabled: true,
     };
 
     it('should make the correct request', () => {
       const fakeParams = {
         params: {a: 'b'},
         paramTypes: {a: 'string'},
+        dataBoostEnabled: true,
       };
 
       const expectedQuery = Object.assign({sql: QUERY.sql}, fakeParams);
@@ -283,13 +285,15 @@ describe('BatchTransaction', () => {
       keys: ['a', 'b'],
       ranges: [{}, {}],
       gaxOptions: GAX_OPTS,
+      dataBoostEnabled: true,
     };
 
     it('should make the correct request', () => {
       const fakeKeySet = {};
       const expectedQuery = {
         table: QUERY.table,
         keySet: fakeKeySet,
+        dataBoostEnabled: true,
       };
 
       const stub = sandbox.stub(batchTransaction, 'createPartitions_');
@@ -329,6 +333,30 @@ describe('BatchTransaction', () => {
       const query = stub.lastCall.args[0];
       assert.strictEqual(query, partition);
     });
+
+    it('should make read requests for read partitions with data boost enabled', () => {
+      const partition = {table: 'abc', dataBoostEnabled: true};
+      const stub = sandbox.stub(batchTransaction, 'read');
+
+      batchTransaction.execute(partition, assert.ifError);
+
+      const [table, options] = stub.lastCall.args;
+      assert.strictEqual(table, partition.table);
+      assert.strictEqual(options, partition);
+    });
+
+    it('should make query requests for non-read partitions with data boost enabled', () => {
+      const partition = {
+        sql: 'SELECT * FROM Singers',
+        dataBoostEnabled: true,
+      };
+      const stub = sandbox.stub(batchTransaction, 'run');
+
+      batchTransaction.execute(partition, assert.ifError);
+
+      const query = stub.lastCall.args[0];
+      assert.strictEqual(query, partition);
+    });
   });
 
   describe('executeStream', () => {