feat: add promise based signatures on createReadPartitions (#2300)

* refactor: createReadPartitions signatures

* refactor

* fix-tests

Author: alkatrivedi

src/batch-transaction.ts

@@ -18,7 +18,7 @@ import {PreciseDate} from '@google-cloud/precise-date';
 import {promisifyAll} from '@google-cloud/promisify';
 import * as extend from 'extend';
 import * as is from 'is';
-import {ExecuteSqlRequest, Snapshot} from './transaction';
+import {ReadRequest, ExecuteSqlRequest, Snapshot} from './transaction';
 import {google} from '../protos/protos';
 import {Session, Database} from '.';
 import {
@@ -35,6 +35,16 @@ export interface TransactionIdentifier {
   timestamp?: google.protobuf.ITimestamp;
 }
 
+export type CreateReadPartitionsResponse = [
+  google.spanner.v1.IPartitionReadRequest,
+  google.spanner.v1.IPartitionResponse,
+];
+
+export type CreateReadPartitionsCallback = ResourceCallback<
+  google.spanner.v1.IPartitionReadRequest,
+  google.spanner.v1.IPartitionResponse
+>;
+
 export type CreateQueryPartitionsResponse = [
   google.spanner.v1.IPartitionQueryRequest,
   google.spanner.v1.IPartitionResponse,
@@ -284,7 +294,17 @@ class BatchTransaction extends Snapshot {
    * @param {CreateReadPartitionsCallback} [callback] Callback function.
    * @returns {Promise<CreateReadPartitionsResponse>}
    */
-  createReadPartitions(options, callback) {
+  createReadPartitions(
+    options: ReadRequest,
+  ): Promise<CreateReadPartitionsResponse>;
+  createReadPartitions(
+    options: ReadRequest,
+    callback: CreateReadPartitionsCallback,
+  ): void;
+  createReadPartitions(
+    options: ReadRequest,
+    cb?: CreateReadPartitionsCallback,
+  ): void | Promise<CreateReadPartitionsResponse> {
     const traceConfig: traceConfig = {
       opts: this._observabilityOptions,
       dbName: this.getDBName(),
@@ -321,7 +341,7 @@ class BatchTransaction extends Snapshot {
             }
 
             span.end();
-            callback(err, partitions, resp);
+            cb!(err, partitions, resp);
           },
         );
       },

test/batch-transaction.ts

@@ -371,7 +371,7 @@ describe('BatchTransaction', () => {
       directedReadOptions: fakeDirectedReadOptionsForRequest,
     };
 
-    it('should make the correct request', () => {
+    it('should make the correct request using callback', () => {
       const fakeKeySet = {};
       const expectedQuery = {
         table: QUERY.table,
@@ -398,6 +398,34 @@ describe('BatchTransaction', () => {
         Object.assign({[LEADER_AWARE_ROUTING_HEADER]: 'true'}),
       );
     });
+
+    it('should make the correct request using await', async () => {
+      const fakeKeySet = {};
+      const expectedQuery = {
+        table: QUERY.table,
+        keySet: fakeKeySet,
+        dataBoostEnabled: true,
+        directedReadOptions: fakeDirectedReadOptionsForRequest,
+      };
+
+      const stub = sandbox.stub(batchTransaction, 'createPartitions_');
+
+      (sandbox.stub(FakeTransaction, 'encodeKeySet') as sinon.SinonStub)
+        .withArgs(QUERY)
+        .returns(fakeKeySet);
+
+      await batchTransaction.createReadPartitions(QUERY);
+
+      const {client, method, reqOpts, gaxOpts, headers} = stub.lastCall.args[0];
+      assert.strictEqual(client, 'SpannerClient');
+      assert.strictEqual(method, 'partitionRead');
+      assert.deepStrictEqual(reqOpts, expectedQuery);
+      assert.strictEqual(gaxOpts, GAX_OPTS);
+      assert.deepStrictEqual(
+        headers,
+        Object.assign({[LEADER_AWARE_ROUTING_HEADER]: 'true'}),
+      );
+    });
   });
 
   describe('execute', () => {

test/mockserver/mockspanner.ts

@@ -268,6 +268,7 @@ export class MockSpanner {
 
     this.read = this.read.bind(this);
     this.streamingRead = this.streamingRead.bind(this);
+    this.partitionRead = this.partitionRead.bind(this);
   }
 
   /**
@@ -993,11 +994,21 @@ export class MockSpanner {
   }
 
   partitionRead(
-    call: grpc.ServerUnaryCall<protobuf.PartitionReadRequest, {}>,
+    call: grpc.ServerUnaryCall<
+      protobuf.PartitionReadRequest,
+      protobuf.PartitionResponse
+    >,
     callback: protobuf.Spanner.PartitionReadCallback,
   ) {
     this.pushRequest(call.request!, call.metadata);
-    callback(createUnimplementedError('PartitionQuery is not yet implemented'));
+    this.simulateExecutionTime(this.partitionRead.name)
+      .then(() => {
+        const response = protobuf.PartitionResponse.create({
+          partitions: [{partitionToken: Buffer.from('mock-token')}],
+        });
+        callback(null, response);
+      })
+      .catch(err => callback(err));
   }
 
   private _updateTransaction(

test/spanner.ts

@@ -3651,6 +3651,23 @@ describe('Spanner with mock server', () => {
       });
     });
 
+    describe('createReadPartitions', () => {
+      it('should create set of read partitions', async () => {
+        const database = newTestDatabase({min: 0, incStep: 1});
+        const query = {
+          table: 'abc',
+          keys: ['a', 'b'],
+          ranges: [{}, {}],
+          gaxOptions: {},
+          dataBoostEnabled: true,
+        };
+        const [transaction] = await database.createBatchTransaction();
+        const [readPartitions] = await transaction.createReadPartitions(query);
+        assert.strictEqual(Object.keys(readPartitions).length, 1);
+        assert.strictEqual(readPartitions[0].table, 'abc');
+      });
+    });
+
     describe('createQueryPartitions', () => {
       it('should create set of query partitions', async () => {
         const database = newTestDatabase({min: 0, incStep: 1});