feat: add promise based signatures for execute method

Author: alkatrivedi

src/batch-transaction.ts

@@ -18,7 +18,14 @@ import {PreciseDate} from '@google-cloud/precise-date';
 import {promisifyAll} from '@google-cloud/promisify';
 import * as extend from 'extend';
 import * as is from 'is';
-import {ReadRequest, ExecuteSqlRequest, Snapshot} from './transaction';
+import {
+  ExecuteSqlRequest,
+  ReadCallback,
+  ReadRequest,
+  Rows,
+  RunCallback,
+  Snapshot,
+} from './transaction';
 import {google} from '../protos/protos';
 import {Session, Database} from '.';
 import {
@@ -55,6 +62,12 @@ export type CreateQueryPartitionsCallback = ResourceCallback<
   google.spanner.v1.IPartitionResponse
 >;
 
+type executeResponse = [
+  Rows,
+  google.spanner.v1.ResultSetStats?,
+  google.spanner.v1.ResultSetMetadata?,
+];
+
 /**
  * Use a BatchTransaction object to create partitions and read/query against
  * your Cloud Spanner database.
@@ -364,12 +377,27 @@ class BatchTransaction extends Snapshot {
    * @example <caption>include:samples/batch.js</caption>
    * region_tag:spanner_batch_execute_partitions
    */
-  execute(partition, callback) {
-    if (is.string(partition.table)) {
-      this.read(partition.table, partition, callback);
+  execute(partition: ReadRequest | ExecuteSqlRequest): Promise<executeResponse>;
+  execute(
+    partition: ReadRequest | ExecuteSqlRequest,
+    callback: ReadCallback | RunCallback,
+  ): void;
+  execute(
+    partition: ReadRequest | ExecuteSqlRequest,
+    cb?: ReadCallback | RunCallback,
+  ): void | Promise<executeResponse> {
+    const isRead = typeof (partition as ReadRequest).table === 'string';
+
+    if (isRead) {
+      this.read(
+        (partition as ReadRequest).table!,
+        partition as ReadRequest,
+        cb as ReadCallback,
+      );
       return;
     }
-    this.run(partition, callback);
+
+    this.run(partition as ExecuteSqlRequest, cb as RunCallback);
   }
   /**
    * Executes partition in streaming mode.

test/batch-transaction.ts

@@ -441,7 +441,7 @@ describe('BatchTransaction', () => {
       },
     };
 
-    it('should make read requests for read partitions', () => {
+    it('should make read requests for read partitions using callback', () => {
       const partition = {table: 'abc'};
       const stub = sandbox.stub(batchTransaction, 'read');
 
@@ -490,6 +490,17 @@ describe('BatchTransaction', () => {
       const query = stub.lastCall.args[0];
       assert.strictEqual(query, partition);
     });
+
+    it('should make read requests for read partitions using await', async () => {
+      const partition = {table: 'abc'};
+      const stub = sandbox.stub(batchTransaction, 'read');
+
+      await batchTransaction.execute(partition);
+
+      const [table, options] = stub.lastCall.args;
+      assert.strictEqual(table, partition.table);
+      assert.strictEqual(options, partition);
+    });
   });
 
   describe('executeStream', () => {

test/mockserver/mockspanner.ts

@@ -863,11 +863,102 @@ export class MockSpanner {
 
   streamingRead(call: grpc.ServerWritableStream<protobuf.ReadRequest, {}>) {
     this.pushRequest(call.request!, call.metadata);
-    call.emit(
-      'error',
-      createUnimplementedError('StreamingRead is not yet implemented'),
-    );
-    call.end();
+
+    this.simulateExecutionTime(this.streamingRead.name)
+      .then(() => {
+        if (call.request!.transaction) {
+          const fullTransactionId = `${call.request!.session}/transactions/${
+            call.request!.transaction.id
+          }`;
+          if (this.abortedTransactions.has(fullTransactionId)) {
+            call.sendMetadata(new Metadata());
+            call.emit(
+              'error',
+              MockSpanner.createTransactionAbortedError(`${fullTransactionId}`),
+            );
+            call.end();
+            return;
+          }
+        }
+        const key = `${call.request!.table}|${JSON.stringify(call.request!.keySet)}`;
+        const res = this.statementResults.get(key);
+        if (res) {
+          if (call.request!.transaction?.begin) {
+            const txn = this._updateTransaction(
+              call.request!.session,
+              call.request!.transaction.begin,
+            );
+            if (txn instanceof Error) {
+              call.sendMetadata(new Metadata());
+              call.emit('error', txn);
+              call.end();
+              return;
+            }
+            if (res.type === StatementResultType.RESULT_SET) {
+              (res.resultSet as protobuf.ResultSet).metadata!.transaction = txn;
+            }
+          }
+          let partialResultSets;
+          let resumeIndex;
+          let streamErr;
+          switch (res.type) {
+            case StatementResultType.RESULT_SET:
+              if (Array.isArray(res.resultSet)) {
+                partialResultSets = res.resultSet;
+              } else {
+                partialResultSets = MockSpanner.toPartialResultSets(
+                  res.resultSet,
+                  'NORMAL',
+                );
+              }
+              // Resume on the next index after the last one seen by the client.
+              resumeIndex =
+                call.request!.resumeToken.length === 0
+                  ? 0
+                  : Number.parseInt(call.request!.resumeToken.toString(), 10) +
+                    1;
+              for (
+                let index = resumeIndex;
+                index < partialResultSets.length;
+                index++
+              ) {
+                const streamErr = this.shiftStreamError(
+                  this.executeStreamingSql.name,
+                  index,
+                );
+                if (streamErr) {
+                  call.sendMetadata(new Metadata());
+                  call.emit('error', streamErr);
+                  break;
+                }
+                call.write(partialResultSets[index]);
+              }
+              break;
+            case StatementResultType.ERROR:
+              call.sendMetadata(new Metadata());
+              call.emit('error', res.error);
+              break;
+            default:
+              call.emit(
+                'error',
+                new Error(`Unknown StatementResult type: ${res.type}`),
+              );
+          }
+        } else {
+          call.emit(
+            'error',
+            new Error(
+              `There is no result registered for ${call.request!.table}`,
+            ),
+          );
+        }
+        call.end();
+      })
+      .catch(err => {
+        call.sendMetadata(new Metadata());
+        call.emit('error', err);
+        call.end();
+      });
   }
 
   beginTransaction(

test/spanner.ts

@@ -40,7 +40,7 @@ import {TEST_INSTANCE_NAME} from './mockserver/mockinstanceadmin';
 import * as mockDatabaseAdmin from './mockserver/mockdatabaseadmin';
 import * as sinon from 'sinon';
 import {google} from '../protos/protos';
-import {ExecuteSqlRequest, RunResponse} from '../src/transaction';
+import {ExecuteSqlRequest, ReadRequest, RunResponse} from '../src/transaction';
 import {Row} from '../src/partial-result-stream';
 import {GetDatabaseOperationsOptions} from '../src/instance';
 import {
@@ -420,9 +420,7 @@ describe('Spanner with mock server', () => {
           },
         });
       } catch (e) {
-        // Ignore the fact that streaming read is unimplemented on the mock
-        // server. We just want to verify that the correct request is sent.
-        assert.strictEqual((e as ServiceError).code, Status.UNIMPLEMENTED);
+        assert.strictEqual((e as ServiceError).code, Status.UNKNOWN);
         assert.deepStrictEqual(
           (e as RequestIDError).requestID,
           `1.${randIdForProcess}.1.1.3.1`,
@@ -583,9 +581,7 @@ describe('Spanner with mock server', () => {
               },
             });
           } catch (e) {
-            // Ignore the fact that streaming read is unimplemented on the mock
-            // server. We just want to verify that the correct request is sent.
-            assert.strictEqual((e as ServiceError).code, Status.UNIMPLEMENTED);
+            assert.strictEqual((e as ServiceError).code, Status.UNKNOWN);
             assert.deepStrictEqual(
               (e as RequestIDError).requestID,
               `1.${randIdForProcess}.1.1.2.1`,
@@ -3683,6 +3679,16 @@ describe('Spanner with mock server', () => {
       });
     });
 
+    describe('execute', () => {
+      it('should create set of read partitions', async () => {
+        const database = newTestDatabase({min: 0, incStep: 1});
+        const [transaction] = await database.createBatchTransaction();
+        const [partitions] = await transaction.createQueryPartitions(selectSql);
+        const [resp] = await transaction.execute(partitions[0]);
+        assert.strictEqual(resp.length, 3);
+      });
+    });
+
     describe('pdml', () => {
       it('should retry on aborted error', async () => {
         const database = newTestDatabase();