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