fix: retry executeStreamingSql when error code is retryable

src/partial-result-stream.ts

@@ -25,6 +25,7 @@ import * as streamEvents from 'stream-events';
 
 import {codec, JSONOptions, Json, Field, Value} from './codec';
 import {SpannerClient as s} from './v1';
+import {ServiceError, status} from 'grpc';
 
 export type ResumeToken = string | Uint8Array;
 
@@ -342,7 +343,10 @@ export function partialResultStream(
   requestFn: RequestFunction,
   options?: RowOptions
 ): PartialResultStream {
+  const retryableCodes = [status.UNAVAILABLE];
   let lastResumeToken: ResumeToken;
+  let lastRetriedErr: ServiceError | undefined;
+  let lastRequestStream: Readable;
 
   // 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.
@@ -355,20 +359,35 @@ export function partialResultStream(
     },
   });
 
+  // This listener ensures that the last request that executed successfully
+  // after one or more retries will end the requestsStream.
+  const endListener = () => {
+    if (lastRetriedErr) {
+      setImmediate(() => requestsStream.end());
+    }
+  };
   const makeRequest = (): void => {
-    requestsStream.add(requestFn(lastResumeToken));
+    if (lastRequestStream) {
+      lastRequestStream.removeListener('end', endListener);
+    }
+    lastRequestStream = requestFn(lastResumeToken);
+    lastRequestStream.on('end', endListener);
+    requestsStream.add(lastRequestStream);
   };
 
-  const retry = (err: Error): void => {
-    if (!lastResumeToken) {
-      // We won't retry the request, so this will flush any rows the
+  const retry = (err: ServiceError): void => {
+    if (!(err.code && retryableCodes!.includes(err.code))) {
+      // This is not a retryable error, so this will flush any rows the
       // checkpoint stream has queued. After that, we will destroy the
       // user's stream with the same error.
       setImmediate(() => batchAndSplitOnTokenStream.destroy(err));
       return;
     }
 
     // We're going to retry from where we left off.
+    // Keep track of the fact that we retried an error in order to end the
+    // merged result stream.
+    lastRetriedErr = err;
     // Empty queued rows on the checkpoint stream (will not emit them to user).
     batchAndSplitOnTokenStream.reset();
     makeRequest();

test/partial-result-stream.ts

@@ -25,6 +25,8 @@ import * as through from 'through2';
 
 import {codec} from '../src/codec';
 import * as prs from '../src/partial-result-stream';
+import {ServiceError, status} from 'grpc';
+import {Row} from '../src/partial-result-stream';
 
 describe('PartialResultStream', () => {
   const sandbox = sinon.createSandbox();
@@ -251,7 +253,62 @@ describe('PartialResultStream', () => {
       );
     });
 
-    it('should resume if there was an error', done => {
+    it('should retry if the initial call returned a retryable error', done => {
+      // This test will emit two rows total:
+      // - UNAVAILABLE error (should retry)
+      // - Two rows
+      // - Confirm all rows were received.
+      const fakeCheckpointStream = through.obj();
+      // tslint:disable-next-line no-any
+      const resetStub = ((fakeCheckpointStream as any).reset = () => {});
+      sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream);
+
+      const firstFakeRequestStream = through.obj();
+      const secondFakeRequestStream = through.obj();
+
+      const requestFnStub = sandbox.stub();
+
+      requestFnStub.onCall(0).callsFake(() => {
+        setTimeout(() => {
+          // This causes a new request stream to be created.
+          firstFakeRequestStream.emit('error', {
+            code: status.UNAVAILABLE,
+            message: 'Error.',
+          } as ServiceError);
+        }, 50);
+
+        return firstFakeRequestStream;
+      });
+
+      requestFnStub.onCall(1).callsFake(resumeToken => {
+        assert.ok(
+          !resumeToken,
+          'Retry should be called with empty resume token'
+        );
+
+        setTimeout(() => {
+          secondFakeRequestStream.push(RESULT_WITH_TOKEN);
+          fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN);
+          secondFakeRequestStream.push(RESULT_WITH_TOKEN);
+          fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN);
+
+          secondFakeRequestStream.end();
+        }, 500);
+
+        return secondFakeRequestStream;
+      });
+
+      partialResultStream(requestFnStub)
+        .on('error', done)
+        .pipe(
+          concat(rows => {
+            assert.strictEqual(rows.length, 2);
+            done();
+          })
+        );
+    });
+
+    it('should resume if there was a retryable error', done => {
       // This test will emit four rows total:
       // - Two rows
       // - Error event (should retry)
@@ -276,8 +333,10 @@ describe('PartialResultStream', () => {
 
           setTimeout(() => {
             // This causes a new request stream to be created.
-            firstFakeRequestStream.emit('error', new Error('Error.'));
-            firstFakeRequestStream.end();
+            firstFakeRequestStream.emit('error', {
+              code: status.UNAVAILABLE,
+              message: 'Error.',
+            } as ServiceError);
           }, 50);
         }, 50);
 
@@ -309,6 +368,49 @@ describe('PartialResultStream', () => {
         );
     });
 
+    it('should emit non-retryable error', done => {
+      // This test will emit two rows and then an error.
+      const fakeCheckpointStream = through.obj();
+      // tslint:disable-next-line no-any
+      const resetStub = ((fakeCheckpointStream as any).reset = () => {});
+      sandbox.stub(checkpointStream, 'obj').returns(fakeCheckpointStream);
+
+      const fakeRequestStream = through.obj();
+
+      const requestFnStub = sandbox.stub();
+
+      requestFnStub.onCall(0).callsFake(() => {
+        setTimeout(() => {
+          fakeRequestStream.push(RESULT_WITH_TOKEN);
+          fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN);
+          fakeRequestStream.push(RESULT_WITH_TOKEN);
+          fakeCheckpointStream.emit('checkpoint', RESULT_WITH_TOKEN);
+
+          setTimeout(() => {
+            fakeRequestStream.emit('error', {
+              code: status.DATA_LOSS,
+              message: 'Non-retryable error.',
+            } as ServiceError);
+          }, 50);
+        }, 50);
+
+        return fakeRequestStream;
+      });
+
+      const receivedRows: Row[] = [];
+      partialResultStream(requestFnStub)
+        .on('data', row => {
+          receivedRows.push(row);
+        })
+        .on('error', err => {
+          // We should receive two rows before we get an error.
+          assert.strictEqual(receivedRows.length, 2);
+          assert.strictEqual(err.code, status.DATA_LOSS);
+          assert.strictEqual(requestFnStub.callCount, 1);
+          done();
+        });
+    });
+
     it('should emit rows and error when there is no token', done => {
       const expectedRow = sinon.match(EXPECTED_ROW);
       const error = new Error('Error.');