diff --git a/test/readrows.ts b/test/readrows.ts index c1a92c88c..75ac8e472 100644 --- a/test/readrows.ts +++ b/test/readrows.ts @@ -13,7 +13,7 @@ // limitations under the License. import {before, describe, it} from 'mocha'; -import {Bigtable, protos, Row, Table} from '../src'; +import {Bigtable, Row, Table} from '../src'; import * as assert from 'assert'; import {Transform, PassThrough, pipeline} from 'stream'; @@ -21,12 +21,22 @@ import {GoogleError} from 'google-gax'; import {MockServer} from '../src/util/mock-servers/mock-server'; import {BigtableClientMockService} from '../src/util/mock-servers/service-implementations/bigtable-client-mock-service'; import {MockService} from '../src/util/mock-servers/mock-service'; -import {debugLog, readRowsImpl} from './utils/readRowsImpl'; -import {ServerWritableStream} from '@grpc/grpc-js'; +import {readRowsImpl} from './utils/readRowsImpl'; -import {ReadRowsServiceParameters} from '../test/utils/readRowsServiceParameters'; +import { + ReadRowsServiceParameters, + ReadRowsWritableStream, +} from '../test/utils/readRowsServiceParameters'; import * as mocha from 'mocha'; +const DEBUG = process.env.BIGTABLE_TEST_DEBUG === 'true'; + +function debugLog(text: string) { + if (DEBUG) { + console.log(text); + } +} + // Define parameters for a standard Bigtable Mock service const VALUE_SIZE = 1024 * 1024; // we want each row to be split into 2 chunks of different sizes @@ -41,16 +51,12 @@ const STANDARD_SERVICE_WITHOUT_ERRORS: ReadRowsServiceParameters = { valueSize: VALUE_SIZE, chunkSize: CHUNK_SIZE, chunksPerResponse: CHUNKS_PER_RESPONSE, + debugLog, }; type PromiseVoid = Promise; interface ServerImplementationInterface { - ( - server: ServerWritableStream< - protos.google.bigtable.v2.IReadRowsRequest, - protos.google.bigtable.v2.IReadRowsResponse - > - ): PromiseVoid; + (server: ReadRowsWritableStream): PromiseVoid; } describe('Bigtable/ReadRows', () => { @@ -316,6 +322,7 @@ describe('Bigtable/ReadRows', () => { chunkSize: CHUNK_SIZE, chunksPerResponse: CHUNKS_PER_RESPONSE, errorAfterChunkNo, + debugLog, }) as ServerImplementationInterface, }); let receivedRowCount = 0; @@ -367,6 +374,7 @@ describe('Bigtable/ReadRows', () => { valueSize: 1, chunkSize: 1, chunksPerResponse: 1, + debugLog, }) as ServerImplementationInterface, }); const sleep = (ms: number) => { diff --git a/test/utils/readRowsImpl.ts b/test/utils/readRowsImpl.ts index 7176eb667..87ed756b5 100644 --- a/test/utils/readRowsImpl.ts +++ b/test/utils/readRowsImpl.ts @@ -12,28 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -import {ServerWritableStream} from '@grpc/grpc-js'; import {protos} from '../../src'; import {GoogleError, Status} from 'google-gax'; import { ChunkGeneratorParameters, + DebugLog, ReadRowsServiceParameters, + ReadRowsWritableStream, } from './readRowsServiceParameters'; - -const DEBUG = process.env.BIGTABLE_TEST_DEBUG === 'true'; - -export function debugLog(text: string) { - if (DEBUG) { - console.log(text); - } -} +import {google} from '../../protos/protos'; +import IRowRange = google.bigtable.v2.IRowRange; // Generate documentation for this function /** Pretty prints the request object. * @param request The request object to pretty print. + * @param debugLog The logging function for printing test output. */ -export function prettyPrintRequest( - request: protos.google.bigtable.v2.IReadRowsRequest +function prettyPrintRequest( + request: protos.google.bigtable.v2.IReadRowsRequest, + debugLog: DebugLog ) { // pretty-printing important parts of the request. // doing it field by field because we want to apply .toString() to all key fields @@ -79,8 +76,13 @@ export function prettyPrintRequest( * The fake table contains monotonically increasing zero padded rows * in the range [keyFrom, keyTo). * @param chunkGeneratorParameters The parameters for generating chunks. + * @param debugLog The logging function for printing test output. + * @returns {protos.google.bigtable.v2.ReadRowsResponse.ICellChunk[]} The generated chunks. */ -function generateChunks(chunkGeneratorParameters: ChunkGeneratorParameters) { +function generateChunks( + chunkGeneratorParameters: ChunkGeneratorParameters, + debugLog: DebugLog +): protos.google.bigtable.v2.ReadRowsResponse.ICellChunk[] { const keyFrom = chunkGeneratorParameters.keyFrom; const keyTo = chunkGeneratorParameters.keyTo; debugLog(`generating chunks from ${keyFrom} to ${keyTo}`); @@ -125,6 +127,11 @@ function generateChunks(chunkGeneratorParameters: ChunkGeneratorParameters) { return chunks; } +/** Checks if the given key is in the provided RowSet. + * @param stringKey The key to check. + * @param rowSet The RowSet to check against. + * @returns {boolean} True if the key is in the RowSet, false otherwise. + */ function isKeyInRowSet( stringKey: string, rowSet?: protos.google.bigtable.v2.IRowSet | null @@ -163,6 +170,84 @@ function isKeyInRowSet( return true; } +/** Gets the key value for the given property specified in the request. + * @param request The request object to get the key value from. + * @param property The property from the request to get the value from. + * @returns {string | undefined} The key value from the request. + */ +function getKeyValue( + request: protos.google.bigtable.v2.IReadRowsRequest, + property: keyof IRowRange +) { + if ( + request?.rows?.rowRanges && + request?.rows?.rowRanges[0] && + request?.rows?.rowRanges[0][property]?.toString() + ) { + return request?.rows?.rowRanges[0][property]?.toString(); + } + return undefined; +} + +/** Gets the key from the request object. + * @param stream The stream object to get the key from. + * @param keySelectionParameters The parameters for selecting the key. + * @returns {number} The selected key for generating chunks + */ +function getSelectedKey( + request: protos.google.bigtable.v2.IReadRowsRequest, + keySelectionParameters: { + keyOpenProperty: keyof IRowRange; + keyClosedProperty: keyof IRowRange; + defaultKey?: number; + } +) { + const keyRequestOpen = getKeyValue( + request, + keySelectionParameters.keyOpenProperty + ); + const keyRequestClosed = getKeyValue( + request, + keySelectionParameters.keyClosedProperty + ); + const defaultKey = keySelectionParameters.defaultKey; + return defaultKey !== undefined + ? defaultKey + : keyRequestClosed + ? parseInt(keyRequestClosed as string) + : parseInt(keyRequestOpen as string) + 1; +} + +/** Generates chunks for rows in a fake table that match the provided RowSet. + * The fake table contains monotonically increasing zero padded rows + * in the range [keyFrom, keyTo). + * @param request The request object to generate chunks from. + * @param serviceParameters The parameters for generating chunks. + * @returns {protos.google.bigtable.v2.ReadRowsResponse.ICellChunk[]} The generated chunks. + */ +function generateChunksFromRequest( + request: protos.google.bigtable.v2.IReadRowsRequest, + serviceParameters: ReadRowsServiceParameters +) { + return generateChunks( + { + keyFrom: getSelectedKey(request, { + keyOpenProperty: 'startKeyOpen', + keyClosedProperty: 'startKeyClosed', + defaultKey: serviceParameters.keyFrom, + }), + keyTo: getSelectedKey(request, { + keyOpenProperty: 'endKeyOpen', + keyClosedProperty: 'endKeyClosed', + defaultKey: serviceParameters.keyTo, + }), + chunkSize: serviceParameters.chunkSize, + valueSize: serviceParameters.valueSize, + }, + serviceParameters.debugLog + ); +} + // Returns an implementation of the server streaming ReadRows call that would return // monotonically increasing zero padded rows in the range [keyFrom, keyTo). // The returned implementation can be passed to gRPC server. @@ -171,20 +256,11 @@ function isKeyInRowSet( // TODO: Perhaps group the if statements into classes so that they can be unit tested. export function readRowsImpl( serviceParameters: ReadRowsServiceParameters -): ( - stream: ServerWritableStream< - protos.google.bigtable.v2.IReadRowsRequest, - protos.google.bigtable.v2.IReadRowsResponse - > -) => Promise { +): (stream: ReadRowsWritableStream) => Promise { let errorAfterChunkNo = serviceParameters.errorAfterChunkNo; - return async ( - stream: ServerWritableStream< - protos.google.bigtable.v2.IReadRowsRequest, - protos.google.bigtable.v2.IReadRowsResponse - > - ): Promise => { - prettyPrintRequest(stream.request); + return async (stream: ReadRowsWritableStream): Promise => { + const debugLog = serviceParameters.debugLog; + prettyPrintRequest(stream.request, debugLog); let stopWaiting: () => void = () => {}; let cancelled = false; @@ -231,60 +307,7 @@ export function readRowsImpl( }); let chunksSent = 0; - let keyFromRequestClosed: string | undefined; - if ( - stream?.request?.rows?.rowRanges && - stream?.request?.rows?.rowRanges[0] && - stream?.request?.rows?.rowRanges[0]?.startKeyClosed?.toString() - ) { - keyFromRequestClosed = - stream?.request?.rows?.rowRanges[0]?.startKeyClosed?.toString(); - } - let keyFromRequestOpen: string | undefined; - if ( - stream?.request?.rows?.rowRanges && - stream?.request?.rows?.rowRanges[0] && - stream?.request?.rows?.rowRanges[0]?.startKeyOpen?.toString() - ) { - keyFromRequestOpen = - stream?.request?.rows?.rowRanges[0]?.startKeyOpen?.toString(); - } - let keyToRequestClosed: string | undefined; - if ( - stream?.request?.rows?.rowRanges && - stream?.request?.rows?.rowRanges[0] && - stream?.request?.rows?.rowRanges[0]?.endKeyClosed?.toString() - ) { - keyToRequestClosed = - stream?.request?.rows?.rowRanges[0]?.endKeyClosed?.toString(); - } - let keyToRequestOpen: string | undefined; - if ( - stream?.request?.rows?.rowRanges && - stream?.request?.rows?.rowRanges[0] && - stream?.request?.rows?.rowRanges[0]?.endKeyOpen?.toString() - ) { - keyToRequestOpen = - stream?.request?.rows?.rowRanges[0]?.endKeyOpen?.toString(); - } - const keyFromUsed = - serviceParameters.keyFrom !== undefined - ? serviceParameters.keyFrom - : keyFromRequestClosed - ? parseInt(keyFromRequestClosed as string) - : parseInt(keyFromRequestOpen as string) + 1; - const keyToUsed = - serviceParameters.keyTo !== undefined - ? serviceParameters.keyTo - : keyToRequestClosed - ? parseInt(keyToRequestClosed as string) - : parseInt(keyToRequestOpen as string) + 1; - const chunks = generateChunks({ - keyFrom: keyFromUsed, - keyTo: keyToUsed, - chunkSize: serviceParameters.chunkSize, - valueSize: serviceParameters.valueSize, - }); + const chunks = generateChunksFromRequest(stream.request, serviceParameters); let lastScannedRowKey: string | undefined; let currentResponseChunks: protos.google.bigtable.v2.ReadRowsResponse.ICellChunk[] = []; diff --git a/test/utils/readRowsServiceParameters.ts b/test/utils/readRowsServiceParameters.ts index 2e866e146..bf9b2c076 100644 --- a/test/utils/readRowsServiceParameters.ts +++ b/test/utils/readRowsServiceParameters.ts @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +import {ServerWritableStream} from '@grpc/grpc-js'; +import {protos} from '../../src'; + /** * This file contains the parameters for the readRowsService. */ @@ -21,14 +24,22 @@ interface SharedReadRowsParameters { valueSize: number; // An upper bound on the amount of data included in the chunks } +export type DebugLog = (message: string) => void; + export interface ReadRowsServiceParameters extends SharedReadRowsParameters { 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; } export interface ChunkGeneratorParameters extends SharedReadRowsParameters { keyFrom: number; // The first row in the generated chunks will start with this key keyTo: number; // The last row in the generated chunks will start with this key } + +export type ReadRowsWritableStream = ServerWritableStream< + protos.google.bigtable.v2.IReadRowsRequest, + protos.google.bigtable.v2.IReadRowsResponse +>;