fix: retry error with timeout

Author: surbhigarg92

src/partial-result-stream.ts

@@ -22,7 +22,7 @@ import mergeStream = require('merge-stream');
 import {common as p} from 'protobufjs';
 import {Readable, Transform} from 'stream';
 import * as streamEvents from 'stream-events';
-import {grpc} from 'google-gax';
+import {grpc, CallOptions} from 'google-gax';
 import {isRetryableInternalError} from './transaction-runner';
 
 import {codec, JSONOptions, Json, Field, Value} from './codec';
@@ -96,6 +96,7 @@ export interface RowOptions {
    *     };
    */
   columnsMetadata?: object;
+  gaxOptions?: CallOptions;
 }
 
 /**
@@ -491,6 +492,8 @@ export function partialResultStream(
   const maxQueued = 10;
   let lastResumeToken: ResumeToken;
   let lastRequestStream: Readable;
+  const startTime = Date.now();
+  const timeout = options?.gaxOptions?.timeout ?? Infinity;
 
   // mergeStream allows multiple streams to be connected into one. This is good;
   // if we need to retry a request and pipe more data to the user's stream.
@@ -541,14 +544,17 @@ export function partialResultStream(
   };
 
   const retry = (err: grpc.ServiceError): void => {
+    const elapsed = Date.now() - startTime;
     if (
       !(
         err.code &&
         (retryableCodes!.includes(err.code) || isRetryableInternalError(err))
       ) ||
       // If we have received too many chunks without a resume token, it is not
       // safe to retry.
-      withoutCheckpointCount > maxQueued
+      withoutCheckpointCount > maxQueued ||
+      // Do not retry if the timeout has been reached
+      elapsed >= timeout!
     ) {
       // This is not a retryable error so this will flush any rows the
       // checkpoint stream has queued. After that, we will destroy the

src/transaction.ts

@@ -704,6 +704,7 @@ export class Snapshot extends EventEmitter {
       jsonOptions,
       maxResumeRetries,
       columnsMetadata,
+      gaxOptions,
     })
       ?.on('response', response => {
         if (response.metadata && response.metadata!.transaction && !this.id) {
@@ -1210,6 +1211,7 @@ export class Snapshot extends EventEmitter {
       jsonOptions,
       maxResumeRetries,
       columnsMetadata,
+      gaxOptions,
     })
       .on('response', response => {
         if (response.metadata && response.metadata!.transaction && !this.id) {

test/partial-result-stream.ts

@@ -311,6 +311,38 @@ describe('PartialResultStream', () => {
         );
     });
 
+    it('should not retry if the initial call returned a retryable error but timeout has reached', done => {
+      const fakeCheckpointStream = through.obj();
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (fakeCheckpointStream as any).reset = () => {};
+
+      sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream);
+
+      const firstFakeRequestStream = through.obj();
+
+      const requestFnStub = sandbox.stub();
+
+      requestFnStub.onCall(0).callsFake(() => {
+        setTimeout(() => {
+          // This causes a new request stream to be created.
+          firstFakeRequestStream.emit('error', {
+            code: grpc.status.UNAVAILABLE,
+            message: 'Error.',
+          } as grpc.ServiceError);
+        }, 50);
+
+        return firstFakeRequestStream;
+      });
+
+      partialResultStream(requestFnStub, {gaxOptions: {timeout: 0}})
+        .on('data', row => {})
+        .on('error', err => {
+          assert.strictEqual(err.code, grpc.status.UNAVAILABLE);
+          assert.strictEqual(requestFnStub.callCount, 1);
+          done();
+        });
+    });
+
     it('should resume if there was a retryable error', done => {
       // This test will emit four rows total:
       // - Two rows

test/transaction.ts

@@ -391,11 +391,15 @@ describe('Transaction', () => {
       });
 
       it('should pass along row options', () => {
+        const gaxOptions = {
+          timeout: 60,
+        };
         const fakeOptions = {
           json: true,
           jsonOptions: {a: 'b'},
           maxResumeRetries: 10,
           columnsMetadata: {column1: {test: 'ss'}, column2: Function},
+          gaxOptions: gaxOptions,
         };
 
         snapshot.createReadStream(TABLE, fakeOptions);
@@ -766,11 +770,15 @@ describe('Transaction', () => {
       });
 
       it('should pass along row options', () => {
+        const gaxOptions = {
+          timeout: 60,
+        };
         const expectedOptions = {
           json: true,
           jsonOptions: {a: 'b'},
           maxResumeRetries: 10,
           columnsMetadata: {column1: {test: 'ss'}, column2: Function},
+          gaxOptions: gaxOptions,
         };
 
         const fakeQuery = Object.assign({}, QUERY, expectedOptions);