Merge branch '359913994-first-PR-add-infrastructure' of https://githu…

…b.com/googleapis/nodejs-bigtable into 359913994-first-PR-add-infrastructure
googleapis · Jan 20, 2025 · 2457dbe · 2457dbe
2 parents 9ec98df + cd97f35
commit 2457dbe
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 2 deletions.
diff --git a/test/readrows.ts b/test/readrows.ts
@@ -461,6 +461,46 @@ describe('Bigtable/ReadRows', () => {
     })();
   });
 
+  it.skip('pitfall: should not request full table scan during a retry on a transient error', async () => {
+    const requests = [];
+
+    const TRANSIENT_ERROR_SERVICE: ReadRowsServiceParameters = {
+      chunkSize: CHUNK_SIZE,
+      valueSize: VALUE_SIZE,
+      chunksPerResponse: CHUNKS_PER_RESPONSE,
+      keyFrom: STANDARD_KEY_FROM,
+      keyTo: STANDARD_KEY_TO,
+      deadlineExceededError: true,
+      hook: request => {
+        requests.push(request);
+      },
+      debugLog,
+    };
+
+    async function readRowsWithDeadline() {
+      service.setService({
+        ReadRows: ReadRowsImpl.createService(
+          TRANSIENT_ERROR_SERVICE
+        ) as ServerImplementationInterface,
+      });
+
+      const rows = await table.getRows();
+      return rows;
+    }
+
+    try {
+      await readRowsWithDeadline();
+      assert.fail('Should have thrown error');
+    } catch (err) {
+      if (err instanceof GoogleError) {
+        assert.equal(err.code, 'DEADLINE_EXCEEDED');
+      }
+
+      // Assert that no retry attempted.
+      assert.strictEqual(requests.length, 1);
+    }
+  });
+
   after(async () => {
     server.shutdown(() => {});
   });

diff --git a/test/utils/readRowsImpl.ts b/test/utils/readRowsImpl.ts
@@ -349,6 +349,11 @@ export class ReadRowsImpl {
    */
   private async handleRequest(stream: ReadRowsWritableStream) {
     const debugLog = this.serviceParameters.debugLog;
+    const hook = this.serviceParameters.hook;
+    if (hook) {
+      hook(stream.request);
+    }
+
     prettyPrintRequest(stream.request, debugLog);
     const readRowsRequestHandler = new ReadRowsRequestHandler(stream, debugLog);
     stream.on('cancelled', () => {
@@ -374,12 +379,15 @@ export class ReadRowsImpl {
   ) {
     const stream = readRowsRequestHandler.stream;
     const debugLog = readRowsRequestHandler.debugLog;
+    const deadlineExceededError = this.serviceParameters.deadlineExceededError;
+
     let chunksSent = 0;
     let lastScannedRowKey: string | undefined;
     let currentResponseChunks: protos.google.bigtable.v2.ReadRowsResponse.ICellChunk[] =
       [];
     let chunkIdx = 0;
     let skipThisRow = false;
+
     for (const chunk of chunks) {
       if (readRowsRequestHandler.cancelled) {
         break;
@@ -409,19 +417,22 @@ export class ReadRowsImpl {
         currentResponseChunks.push(chunk);
         ++chunkIdx;
       }
+
       if (
         currentResponseChunks.length ===
           this.serviceParameters.chunksPerResponse ||
         chunkIdx === this.errorAfterChunkNo ||
         // if we skipped a row and set lastScannedRowKey, dump everything and send a separate message with lastScannedRowKey
-        lastScannedRowKey
+        lastScannedRowKey ||
+        deadlineExceededError
       ) {
         const response: protos.google.bigtable.v2.IReadRowsResponse = {
           chunks: currentResponseChunks,
         };
         chunksSent += currentResponseChunks.length;
         await readRowsRequestHandler.sendResponse(response);
         currentResponseChunks = [];
+
         if (chunkIdx === this.errorAfterChunkNo) {
           debugLog(`sending error after chunk #${chunkIdx}`);
           this.errorAfterChunkNo = undefined; // do not send error for the second time
@@ -431,7 +442,17 @@ export class ReadRowsImpl {
           readRowsRequestHandler.cancelled = true;
           break;
         }
+
+        if (deadlineExceededError) {
+          debugLog('sending deadline exceeded error');
+          const error = new GoogleError('Deadline exceeded');
+          error.code = Status.DEADLINE_EXCEEDED;
+          stream.emit('error', error);
+          readRowsRequestHandler.cancelled = true;
+          break;
+        }
       }
+
       if (lastScannedRowKey) {
         const response: protos.google.bigtable.v2.IReadRowsResponse = {
           lastScannedRowKey,

diff --git a/test/utils/readRowsServiceParameters.ts b/test/utils/readRowsServiceParameters.ts
@@ -27,11 +27,13 @@ interface SharedReadRowsParameters {
 export type DebugLog = (message: string) => void;
 
 export interface ReadRowsServiceParameters extends SharedReadRowsParameters {
+  deadlineExceededError?: boolean; // Send deadline exceeded transient error
+  errorAfterChunkNo?: number; // The chunk that the error should come after
   keyFrom?: number; // The key the data coming from the service will start from
   keyTo?: number; // The key the data coming from the service will end at
-  errorAfterChunkNo?: number; // The chunk that the error should come after
   chunksPerResponse: number; // The total number of chunks the server should send
   debugLog: DebugLog;
+  hook?: (request: protos.google.bigtable.v2.IReadRowsRequest) => void;
 }
 
 export interface ChunkGeneratorParameters extends SharedReadRowsParameters {