feat: add inline begin transaction for batch DML requests

ko3a4ok · ko3a4ok · commit 2caaa0fd2333 · 2022-09-20T00:56:48.000Z
diff --git a/src/transaction.ts b/src/transaction.ts
@@ -1245,7 +1245,7 @@ export class Snapshot extends EventEmitter {
    *
    * @param {spannerClient.spanner.v1.ITransaction} resp Response object.
    */
-  private _update(resp: spannerClient.spanner.v1.ITransaction): void {
+  protected _update(resp: spannerClient.spanner.v1.ITransaction): void {
     const {id, readTimestamp} = resp;
 
     this.id = id!;
@@ -1584,14 +1584,20 @@ export class Transaction extends Dml {
         return {sql, params, paramTypes};
       });
 
+    const transaction: spannerClient.spanner.v1.ITransactionSelector = {};
+    if (this.id) {
+      transaction.id = this.id as Uint8Array;
+    } else {
+      transaction.begin = this._options;
+    }
     const reqOpts: spannerClient.spanner.v1.ExecuteBatchDmlRequest = {
       session: this.session.formattedName_!,
       requestOptions: this.configureTagOptions(
         false,
         this.requestOptions?.transactionTag ?? undefined,
         (options as BatchUpdateOptions).requestOptions
       ),
-      transaction: {id: this.id!},
+      transaction,
       seqno: this._seqno++,
       statements,
     } as spannerClient.spanner.v1.ExecuteBatchDmlRequest;
@@ -1618,6 +1624,11 @@ export class Transaction extends Dml {
         }
 
         const {resultSets, status} = resp;
+        for (const resultSet of resultSets) {
+          if (!this.id && resultSet.metadata && resultSet.metadata.transaction) {
+            this._update(resultSet.metadata.transaction);
+          }
+        }
         const rowCounts: number[] = resultSets.map(({stats}) => {
           return (
             (stats &&
diff --git a/test/mockserver/mockspanner.ts b/test/mockserver/mockspanner.ts
@@ -720,7 +720,7 @@ export class MockSpanner {
     this.pushRequest(call.request!, call.metadata);
     this.simulateExecutionTime(this.executeBatchDml.name)
       .then(() => {
-        if (call.request!.transaction && call.request!.transaction.id) {
+        if (call.request!.transaction) {
           const fullTransactionId = `${call.request!.session}/transactions/${
             call.request!.transaction.id
           }`;
@@ -764,7 +764,21 @@ export class MockSpanner {
                 callback(new Error('Wrong result type for batch DML'));
                 break;
               case StatementResultType.UPDATE_COUNT:
-                results.push(MockSpanner.toResultSet(res.updateCount));
+                let resultSet = MockSpanner.toResultSet(res.updateCount);
+                if (call.request!.transaction!.begin && i === 0) {
+                  const transaction = this._updateTransaction(
+                      call.request!.session,
+                      call.request?.transaction!.begin
+                  );
+                  if (transaction instanceof Error) {
+                    callback(transaction);
+                    break;
+                  }
+                  resultSet.metadata = protobuf.ResultSetMetadata.create({
+                    transaction
+                  });
+                }
+                results.push(resultSet);
                 break;
               case StatementResultType.ERROR:
                 if ((res.error as grpc.ServiceError).code) {
diff --git a/test/spanner.ts b/test/spanner.ts
@@ -289,6 +289,7 @@ describe('Spanner with mock server', () => {
               requestTag: 'request-tag',
             },
           });
+          await tx!.batchUpdate([insertSql, insertSql]);
           return await tx.commit();
         }
       );
@@ -307,6 +308,13 @@ describe('Spanner with mock server', () => {
         request.requestOptions!.transactionTag,
         'transaction-tag'
       );
+      assert.ok(request.transaction?.begin, 'transaction is not empty');
+      const nextBatchRequest = spannerMock.getRequests().reverse().find(val => {
+        return (val as v1.ExecuteBatchDmlRequest).statements;
+      }) as v1.ExecuteBatchDmlRequest;
+      assert.ok(nextBatchRequest, 'no ExecuteBatchDmlRequest found');
+      assert.ok(nextBatchRequest.transaction?.id, 'no transaction ID');
+
       const commitRequest = spannerMock.getRequests().find(val => {
         return (val as v1.CommitRequest).mutations;
       }) as v1.CommitRequest;
