refactor: createReadPartitions signatures

Author: alkatrivedi

src/batch-transaction.ts

@@ -18,11 +18,12 @@ import {PreciseDate} from '@google-cloud/precise-date';
 import {promisifyAll} from '@google-cloud/promisify';
 import * as extend from 'extend';
 import * as is from 'is';
-import {Snapshot} from './transaction';
+import {ReadRequest, Snapshot} from './transaction';
 import {google} from '../protos/protos';
 import {Session, Database} from '.';
 import {
   CLOUD_RESOURCE_HEADER,
+  ResourceCallback,
   addLeaderAwareRoutingHeader,
 } from '../src/common';
 import {startTrace, setSpanError, traceConfig} from './instrument';
@@ -34,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
+>;
+
 /**
  * Use a BatchTransaction object to create partitions and read/query against
  * your Cloud Spanner database.
@@ -266,7 +277,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(),
@@ -276,7 +297,7 @@ class BatchTransaction extends Snapshot {
       'BatchTransaction.createReadPartitions',
       traceConfig,
       span => {
-        const reqOpts = Object.assign({}, options, {
+        const reqOpts: ReadRequest = Object.assign({}, options, {
           keySet: Snapshot.encodeKeySet(options),
         });
 
@@ -303,7 +324,7 @@ class BatchTransaction extends Snapshot {
             }
 
             span.end();
-            callback(err, partitions, resp);
+            cb!(err, partitions, resp);
           },
         );
       },

test/batch-transaction.ts

@@ -341,7 +341,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,
@@ -368,6 +368,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

@@ -267,6 +267,7 @@ export class MockSpanner {
 
     this.read = this.read.bind(this);
     this.streamingRead = this.streamingRead.bind(this);
+    this.partitionRead = this.partitionRead.bind(this);
   }
 
   /**
@@ -982,11 +983,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,25 @@ 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');
+        transaction.close();
+        await database.close();
+      });
+    });
+
     describe('pdml', () => {
       it('should retry on aborted error', async () => {
         const database = newTestDatabase();