fix: Address assertion error in TestReadRows_Retry_LastScannedRow conformance test (#1521)

* Wrap in event type and now I get failing tests

* Correct the fake chunk transformer mock

* Adjust another mock so that Fake chunk transformer

is correct

* Fix all the tests in the chunk transformer

* Adjust interfaces ready for new push events

* Push lastScannedRow events from the chunk transfo

* Remove TODO

* Remove LastScannedRow from known_failures

* Remove console log

* Fix linting error

* Remove the data event wrapping

* Revert mock

* Revert mock

* Revert mock

* Remove unused data properties

* For data events, don’t send an enum

* Run the linter

* Revert changes in chunk transformer test

* Remove the DataEvent import

* Remove unused dependency

* Update the chunk transformer comment

* Clarify comment

* correct comment

* Correct comment

* Correct comment

* Correct the comments

Author: danieljbruce

src/chunktransformer.ts

@@ -18,6 +18,10 @@ import {TableUtils} from './utils/table';
 
 export type Value = string | number | boolean | Uint8Array;
 
+export enum DataEvent {
+  LAST_ROW_KEY_UPDATE,
+}
+
 export interface Chunk {
   rowContents: Value;
   commitRow: boolean;
@@ -51,6 +55,12 @@ export interface TransformErrorProps {
   message: string;
   chunk: Chunk | null;
 }
+export interface ChunkPushLastScannedRowData {
+  eventType: DataEvent.LAST_ROW_KEY_UPDATE;
+  lastScannedRowKey?: string;
+}
+
+export type ChunkPushData = Row | ChunkPushLastScannedRowData;
 
 class TransformError extends Error {
   constructor(props: TransformErrorProps) {
@@ -159,6 +169,19 @@ export class ChunkTransformer extends Transform {
           userOptions: this.options,
         }
       );
+      /**
+       * Push an event that will update the lastRowKey in the user stream after
+       * all rows ahead of this event have reached the user stream. This will
+       * ensure that a retry excludes the lastScannedRow as this is required
+       * for the TestReadRows_Retry_LastScannedRow conformance test to pass. It
+       * is important to use a 'data' event to update the last row key in order
+       * to allow all the data queued ahead of this event to reach the user
+       * stream first.
+       */
+      this.push({
+        eventType: DataEvent.LAST_ROW_KEY_UPDATE,
+        lastScannedRowKey: this.lastRowKey,
+      });
     }
     next();
   }

src/tabular-api-surface.ts

@@ -26,7 +26,12 @@ import {
 } from './index';
 import {Filter, BoundData, RawFilter} from './filter';
 import {Row} from './row';
-import {ChunkTransformer} from './chunktransformer';
+import {
+  ChunkPushData,
+  ChunkPushLastScannedRowData,
+  ChunkTransformer,
+  DataEvent,
+} from './chunktransformer';
 import {BackoffSettings} from 'google-gax/build/src/gax';
 import {google} from '../protos/protos';
 import {CallOptions, ServiceError} from 'google-gax';
@@ -244,11 +249,27 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
       objectMode: true,
       readableHighWaterMark: 0, // We need to disable readside buffering to allow for acceptable behavior when the end user cancels the stream early.
       writableHighWaterMark: 0, // We need to disable writeside buffering because in nodejs 14 the call to _transform happens after write buffering. This creates problems for tracking the last seen row key.
-      transform(row, _encoding, callback) {
+      transform(event, _encoding, callback) {
         if (userCanceled) {
           callback();
           return;
         }
+        if (event.eventType === DataEvent.LAST_ROW_KEY_UPDATE) {
+          /**
+           * This code will run when receiving an event containing
+           * lastScannedRowKey data that the chunk transformer sent. When the
+           * chunk transformer gets lastScannedRowKey data, this code
+           * updates the lastRowKey to ensure row ids with the lastScannedRowKey
+           * aren't re-requested in retries. The lastRowKey needs to be updated
+           * here and not in the chunk transformer to ensure the update is
+           * queued behind all events that deliver data to the user stream
+           * first.
+           */
+          lastRowKey = event.lastScannedRowKey;
+          callback();
+          return;
+        }
+        const row = event;
         if (TableUtils.lessThanOrEqualTo(row.id, lastRowKey)) {
           /*
           Sometimes duplicate rows reach this point. To avoid delivering
@@ -425,17 +446,34 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
       activeRequestStream = requestStream!;
 
       const toRowStream = new Transform({
-        transform: (rowData, _, next) => {
+        transform: (rowData: ChunkPushData, _, next) => {
           if (
             userCanceled ||
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
             (userStream as any)._writableState.ended
           ) {
             return next();
           }
-          const row = this.row(rowData.key);
-          row.data = rowData.data;
-          next(null, row);
+          if (
+            (rowData as ChunkPushLastScannedRowData).eventType ===
+            DataEvent.LAST_ROW_KEY_UPDATE
+          ) {
+            /**
+             * If the data is the chunk transformer communicating that the
+             * lastScannedRow was received then this message is passed along
+             * to the user stream to update the lastRowKey.
+             */
+            next(null, rowData);
+          } else {
+            /**
+             * If the data is just regular rows being pushed from the
+             * chunk transformer then the rows are encoded so that they
+             * can be consumed by the user stream.
+             */
+            const row = this.row((rowData as Row).key as string);
+            row.data = (rowData as Row).data;
+            next(null, row);
+          }
         },
         objectMode: true,
       });

testproxy/known_failures.txt

@@ -12,7 +12,6 @@ TestReadRow_Retry_WithRetryInfo\|
 TestReadRows_ReverseScans_FeatureFlag_Enabled\|
 TestReadRows_NoRetry_OutOfOrderError_Reverse\|
 TestReadRows_Retry_PausedScan\|
-TestReadRows_Retry_LastScannedRow\|
 TestReadRows_Retry_LastScannedRow_Reverse\|
 TestReadRows_Retry_StreamReset\|
 TestReadRows_Generic_CloseClient\|