From deb0e7c87222465f89bfe117d60939f3a039e552 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Tue, 25 Feb 2025 16:09:14 +0100 Subject: [PATCH 01/14] feat: Execute Query --- package.json | 2 + protos/test_proxy.proto | 88 +- .../api-reference-doc-snippets/instance.js | 77 + src/execute-query/bytebuffertransformer.ts | 150 ++ src/execute-query/executequerystatemachine.ts | 591 +++++ src/execute-query/metadataconsumer.ts | 122 + src/execute-query/namedlist.ts | 98 + src/execute-query/parameterparsing.ts | 308 +++ src/execute-query/preparedstatement.ts | 287 +++ .../protobufreadertransformer.ts | 135 ++ .../queryresultrowtransformer.ts | 255 ++ src/execute-query/types.ts | 83 + src/execute-query/values.ts | 224 ++ src/index.ts | 3 + src/instance.ts | 248 ++ src/tabular-api-surface.ts | 12 + src/utils/retry-options.ts | 77 + system-test/bigtable.ts | 96 + test/base64keymap.ts | 184 ++ test/bytebuffertransformer.ts | 379 +++ test/executequery.ts | 2049 +++++++++++++++++ test/instance.ts | 1430 +++++++++++- test/utils/proto-bytes.ts | 144 ++ testproxy/known_failures.txt | 4 +- testproxy/services/execute-query.js | 64 + testproxy/services/index.js | 2 + .../request/createExecuteQueryResponse.ts | 223 ++ 27 files changed, 7327 insertions(+), 8 deletions(-) create mode 100644 src/execute-query/bytebuffertransformer.ts create mode 100644 src/execute-query/executequerystatemachine.ts create mode 100644 src/execute-query/metadataconsumer.ts create mode 100644 src/execute-query/namedlist.ts create mode 100644 src/execute-query/parameterparsing.ts create mode 100644 src/execute-query/preparedstatement.ts create mode 100644 src/execute-query/protobufreadertransformer.ts create mode 100644 src/execute-query/queryresultrowtransformer.ts create mode 100644 src/execute-query/types.ts create mode 100644 src/execute-query/values.ts create mode 100644 src/utils/retry-options.ts create mode 100644 test/base64keymap.ts create mode 100644 test/bytebuffertransformer.ts create mode 100644 test/executequery.ts create mode 100644 test/utils/proto-bytes.ts create mode 100644 testproxy/services/execute-query.js create mode 100644 testproxy/services/utils/request/createExecuteQueryResponse.ts diff --git a/package.json b/package.json index b78c563fe..87b19c4aa 100644 --- a/package.json +++ b/package.json @@ -58,10 +58,12 @@ "@opentelemetry/sdk-metrics": "^1.30.0", "@types/long": "^4.0.0", "arrify": "2.0.0", + "abort-controller": "^3.0.0", "concat-stream": "^2.0.0", "dot-prop": "6.0.0", "escape-string-regexp": "4.0.0", "extend": "^3.0.2", + "fast-crc32c": "^2.0.0", "google-gax": "^5.0.1-rc.0", "grpc-gcp": "^1.0.1", "is": "^3.3.0", diff --git a/protos/test_proxy.proto b/protos/test_proxy.proto index 551dd4d8b..f919dc476 100644 --- a/protos/test_proxy.proto +++ b/protos/test_proxy.proto @@ -1,4 +1,4 @@ -// Copyright 2023 Google LLC +// Copyright 2024 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -22,12 +22,43 @@ import "google/bigtable/v2/data.proto"; import "google/protobuf/duration.proto"; import "google/rpc/status.proto"; -option go_package = "./testproxypb"; +option go_package = "cloud.google.com/go/bigtable/testproxy/testproxypb;testproxypb"; option java_multiple_files = true; option java_package = "com.google.cloud.bigtable.testproxy"; +// A config flag that dictates how the optional features should be enabled +// during the client creation. The optional features customize how the client +// interacts with the server, and are defined in +// https://github.com/googleapis/googleapis/blob/master/google/bigtable/v2/feature_flags.proto +enum OptionalFeatureConfig { + OPTIONAL_FEATURE_CONFIG_DEFAULT = 0; + + OPTIONAL_FEATURE_CONFIG_ENABLE_ALL = 1; +} + // Request to test proxy service to create a client object. message CreateClientRequest { + message SecurityOptions { + // Access token to use for client credentials. If empty, the client will not + // use any call credentials. Certain implementations may require `use_ssl` + // to be set when using this. + string access_token = 1; + + // Whether to use SSL channel credentials when connecting to the data + // endpoint. + bool use_ssl = 2; + + // If using SSL channel credentials, override the SSL endpoint to match the + // host that is specified in the backend's certificate. Also sets the + // client's authority header value. + string ssl_endpoint_override = 3; + + // PEM encoding of the server root certificates. If not set, the default + // root certs will be used instead. The default can be overridden via the + // GRPC_DEFAULT_SSL_ROOTS_FILE_PATH env var. + string ssl_root_certs_pem = 4; + } + // A unique ID associated with the client object to be created. string client_id = 1; @@ -52,6 +83,21 @@ message CreateClientRequest { // the created client. Otherwise, the default timeout from the client library // will be used. Note that the override applies to all the methods. google.protobuf.Duration per_operation_timeout = 6; + + // Optional config that dictates how the optional features should be enabled + // during the client creation. Please check the enum type's docstring above. + OptionalFeatureConfig optional_feature_config = 7; + + // Options to allow connecting to backends with channel and/or call + // credentials. This is needed internally by Cloud Bigtable's own testing + // frameworks.It is not necessary to support these fields for client + // conformance testing. + // + // WARNING: this allows the proxy to connect to a real production + // CBT backend with the right options, however, the proxy itself is insecure + // so it is not recommended to use it with real credentials or outside testing + // contexts. + SecurityOptions security_options = 8; } // Response from test proxy service for CreateClientRequest. @@ -203,6 +249,39 @@ message ReadModifyWriteRowRequest { google.bigtable.v2.ReadModifyWriteRowRequest request = 2; } +// Request to test proxy service to execute a query. +message ExecuteQueryRequest { + // The ID of the target client object. + string client_id = 1; + + // The raw request to the Bigtable server. + google.bigtable.v2.ExecuteQueryRequest request = 2; +} + +// Response from test proxy service for ExecuteQueryRequest. +message ExecuteQueryResult { + // The RPC status from the client binding. + google.rpc.Status status = 1; + + // Name and type information for the query result. + ResultSetMetadata metadata = 4; + + // Encoded version of the ResultSet. Should not contain type information. + repeated SqlRow rows = 3; +} + +// Schema information for the query result. +message ResultSetMetadata { + // Column metadata for each column inthe query result. + repeated google.bigtable.v2.ColumnMetadata columns = 1; +} + +// Representation of a single row in the query result. +message SqlRow { + // Columnar values returned by the query. + repeated google.bigtable.v2.Value values = 1; +} + // Note that all RPCs are unary, even when the equivalent client binding call // may be streaming. This is an intentional simplification. // @@ -265,4 +344,7 @@ service CloudBigtableV2TestProxy { // Performs a read-modify-write operation with the client. rpc ReadModifyWriteRow(ReadModifyWriteRowRequest) returns (RowResult) {} -} + + // Executes a BTQL query with the client. + rpc ExecuteQuery(ExecuteQueryRequest) returns (ExecuteQueryResult) {} +} \ No newline at end of file diff --git a/samples/api-reference-doc-snippets/instance.js b/samples/api-reference-doc-snippets/instance.js index 7894c8be4..67f116d97 100644 --- a/samples/api-reference-doc-snippets/instance.js +++ b/samples/api-reference-doc-snippets/instance.js @@ -362,6 +362,83 @@ const snippets = { }); // [END bigtable_api_del_instance] }, + + executeQuery: (instanceId, tableId) => { + // [START bigtable_api_execute_query] + const {Bigtable} = require('@google-cloud/bigtable'); + const bigtable = new Bigtable(); + const instance = bigtable.instance(instanceId); + + const query = `SELECT + _key + from \`${tableId}\` WHERE _key=@row_key`; + const parameters = { + row_key: 'alincoln', + }; + + // if query parameter types are ambiguous, you can pass types explicitly + const parameter_types = { + row_key: Bigtable.ExecuteQueryTypes.String(), + }; + + const options = { + query, + parameters, + parameter_types, // optional + }; + + instance + .executeQuery(options) + .then(result => { + const rows = result[0]; + }) + .catch(err => { + // Handle the error. + }); + + // [END bigtable_api_execute_query] + }, + + createExecuteQueryStream: (instanceId, tableId) => { + // [START bigtable_api_create_query_stream] + const {Bigtable} = require('@google-cloud/bigtable'); + const bigtable = new Bigtable(); + const instance = bigtable.instance(instanceId); + + const query = `SELECT + _key + from \`${tableId}\` WHERE _key=@row_key`; + const parameters = { + row_key: 'alincoln', + }; + + const options = { + query, + parameters, + }; + instance + .createExecuteQueryStream(options) + .on('error', err => { + // Handle the error. + }) + .on('data', row => { + // `row` is a QueryResultRow object. + }) + .on('end', () => { + // All rows retrieved. + }); + + // If you anticipate many results, you can end a stream early to prevent + // unnecessary processing. + //- + // instance + // .createExecuteQueryStream(options) + // .on('data', function (row) { + // this.end(); + // }); + + // [END bigtable_api_create_query_stream] + }, }; module.exports = snippets; diff --git a/src/execute-query/bytebuffertransformer.ts b/src/execute-query/bytebuffertransformer.ts new file mode 100644 index 000000000..223328a20 --- /dev/null +++ b/src/execute-query/bytebuffertransformer.ts @@ -0,0 +1,150 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {Transform, TransformCallback} from 'stream'; +// eslint-disable-next-line @typescript-eslint/no-var-requires +import {google} from '../../protos/protos'; +import * as SqlValues from './values'; + +/** + * stream.Transform which buffers bytes from `ExecuteQuery` responses until + * resumeToken is received. At that point all buffered messages are passed + * forward. + */ +export class ByteBufferTransformer extends Transform { + private messageQueue: Buffer[] = []; + private messageBuffer: Uint8Array[] = []; + private protoBytesEncoding?: BufferEncoding; + + constructor(protoBytesEncoding?: BufferEncoding) { + super({objectMode: true, highWaterMark: 0}); + this.protoBytesEncoding = protoBytesEncoding; + } + + private resetQueueAndBuffer = ( + estimatedBatchSize: number | null | undefined, + ): void => { + this.messageQueue = []; + this.messageBuffer = new Array(estimatedBatchSize || 0); + }; + + private flushMessageBuffer = ( + batchChecksum: number, + estimatedBatchSize: number | null | undefined, + ): void => { + if (this.messageBuffer.length === 0) { + throw new Error('Recieved empty batch with non-zero checksum.'); + } + const newBatch = Buffer.concat(this.messageBuffer); + if (!SqlValues.checksumValid(newBatch, batchChecksum)) { + throw new Error('Failed to validate next batch of results'); + } + this.messageQueue.push(newBatch); + this.messageBuffer = new Array(estimatedBatchSize || 0); + }; + + private pushMessages = (resumeToken: string | Uint8Array): void => { + const token = SqlValues.ensureUint8Array( + resumeToken, + this.protoBytesEncoding, + ); + if (this.messageBuffer.length !== 0) { + throw new Error('Recieved incomplete batch of rows.'); + } + this.push([this.messageQueue, token]); + this.messageBuffer = []; + this.messageQueue = []; + }; + + /** + * Process a `PartialResultSet` message from the server. + * For more info refer to the PartialResultSet protobuf definition. + * @param partialResultSet The `PartialResultSet` message to process. + */ + private processProtoRowsBatch = ( + partialResultSet: google.bigtable.v2.IPartialResultSet, + ): void => { + let handled = false; + if (partialResultSet.reset) { + this.resetQueueAndBuffer(partialResultSet.estimatedBatchSize); + handled = true; + } + + if (partialResultSet.protoRowsBatch?.batchData?.length) { + this.messageBuffer.push( + SqlValues.ensureUint8Array( + partialResultSet.protoRowsBatch.batchData, + this.protoBytesEncoding, + ), + ); + handled = true; + } + + if (partialResultSet.batchChecksum) { + this.flushMessageBuffer( + partialResultSet.batchChecksum, + partialResultSet.estimatedBatchSize, + ); + handled = true; + } + + if ( + partialResultSet.resumeToken && + partialResultSet.resumeToken.length > 0 + ) { + this.pushMessages(partialResultSet.resumeToken); + handled = true; + } + + if (!handled) { + throw new Error('Response did not contain any results!'); + } + }; + + _transform( + chunk: google.bigtable.v2.ExecuteQueryResponse, + _encoding: BufferEncoding, + callback: TransformCallback, + ) { + let maybeError: Error | null = null; + const reponse = chunk as google.bigtable.v2.ExecuteQueryResponse; + try { + switch (reponse.response) { + case 'results': { + this.processProtoRowsBatch(reponse.results!); + break; + } + default: + throw Error(`Response contains unknown type ${reponse.response}`); + } + } catch (error) { + maybeError = new Error( + `Internal Error. Failed to process response: ${error}`, + ); + } + callback(maybeError); + } + + _flush(callback: TransformCallback): void { + if (this.messageBuffer.length > 0) { + callback( + new Error( + 'Internal Error. Last message did not contain a resumeToken.', + ), + ); + return; + } + callback(null); + } +} diff --git a/src/execute-query/executequerystatemachine.ts b/src/execute-query/executequerystatemachine.ts new file mode 100644 index 000000000..5f458a47d --- /dev/null +++ b/src/execute-query/executequerystatemachine.ts @@ -0,0 +1,591 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +import { + PreparedStatement, + PreparedStatementDataCallback, +} from './preparedstatement'; +import {Bigtable} from '..'; +import {ServiceError, RetryOptions} from 'google-gax'; +import {google} from '../../protos/protos'; +import * as SqlTypes from './types'; +import {AbortableDuplex} from '..'; +import {ByteBufferTransformer} from './bytebuffertransformer'; +import { + DEFAULT_RETRY_COUNT, + isExpiredQueryError, + RETRYABLE_STATUS_CODES, +} from '../utils/retry-options'; +import {ExecuteQueryStreamWithMetadata} from './values'; +import {ProtobufReaderTransformer} from './protobufreadertransformer'; +import {MetadataConsumer} from './metadataconsumer'; +import { + DEFAULT_BACKOFF_SETTINGS, + isRstStreamError, +} from '../tabular-api-surface'; +const pumpify = require('pumpify'); + +/** + * Creates a stream with some additional functionalities used by + * the ExecuteQueryStateMachine. + */ +interface CallerStream extends ExecuteQueryStreamWithMetadata { + /** + * Sets the metadata used to parse the executeQuery responses. + */ + updateMetadata: (metadata: SqlTypes.ResultSetMetadata) => void; + /** + * @returns the latest resumeToken before which all data was processed. + */ + getLatestResumeToken: () => Uint8Array | string | null; + /** + * @param callback guaranteed to be called *after* the last message + * was processed. + */ + onDrain: (callback: () => void) => void; + /** + * No other data event will be emitted after this method is called. + */ + close: () => void; + /** + * keeps a reference to the state machine. + */ + _stateMachine: ExecuteQueryStateMachine; +} + +interface StreamRetryOptions { + maxRetries: number; + totalTimeout: number; + retryCodes: Set; + initialRetryDelayMillis: number; + retryDelayMultiplier: number; + maxRetryDelayMillis: number; +} + +export type State = + /** + * This is the starting state. When the executeQuery starts we try to + * fetch the query plan from the PreparedStatement object. It is done via a callback + * If the query plan is expired, it might take time to refresh and the callback + * won't be called immediately. + */ + | 'AwaitingQueryPlan' + /** + * We may have recieved some data from the server, but no resumeToken has beed reached. + * This is an important distinction, because before the resumeToken, if the server + * returns an "expired query plan" error, we can still try to refresh it. + */ + | 'BeforeFirstResumeToken' + /** + * After the first resumeToken has been reached, we can't refresh the query plan, + * because the schema could have changed in the mean time. This would cause the + * new rows to be parsed differently than the previous ones. That's why, in this + * state, we treat the "expired query plan" error as non-retryable. + */ + | 'AfterFirstResumeToken' + /** + * When we need to properly dispose of the* old responseStream and buteBuffer, we + * enter a "Draining..." state. Depending on what we want to do next, we have + * DrainAndRefreshQueryPlan - moves to AwaitingQueryPlan after draining completed + * DrainingBeforeResumeToken - moves to BeforeFirstResumeToken after draining completed + * DrainingAfterResumeToken - moves to AfterFirstResumeToken after draining completed + * + * We want to make a new request only when all requests already written to the Reader + * by our previous active request stream were processed. + * + * For simplicity, we will drop the previous Bigtable stream and ByteBuffer transform + * and recreate them. We could also keep the ByteBuffer alive, but that would require + * us to clean up its internal state and still wait for the entire buffer to be + * read—just one step upstream. + * + * Please note that we cannot use gax's built-in streaming retries, as we have no way + * of informing it that we'd like to wait for an event before retrying. An alternative + * approach would be to purge all buffers of all streams before making a request, but + * there is no standard API for that. Additionally, this would not help us with gax, + * as we cannot traverse all streams upstream of our ByteBuffer to purge their buffers, + * nor can we rely on implementation details. + * + * We cannot simply wait until all events are processed by ByteBuffer's _transform() method, + * as there might still be events left in ByteBuffer's readable buffer that we do not want + * to discard. + * + * Our solution is to wait until all events in the Reader's writable buffer are processed + * and use the last resumeToken seen by the Reader to make a new request. + * + * We will detach (unpipe) the ByteBuffer from the Reader and wait until all requests + * written to theReader by the ByteBuffer are processed using the _transform() method. + * This ensures that all events written before detachment are handled by _transform(), + * and the last resumption token seen by the Reader is the correct one to use. + * + * Thus, we will wait for the buffer to clear before making a new request and use the + * last resumeToken seen by the Reader to determine the correct token for the retry request. + * + * This guarantees that no responses will be lost—everything processed by the + * Reader's `_transform()` method has been pushed to the caller and won't be discarded. + * Additionally, no duplicates will occur, as no more responses will be seen by `_transform()` + * until a new request is made. + */ + | 'DrainAndRefreshQueryPlan' + /** + * Moves to BeforeFirstResumeToken after draining. For more info see 'DrainAndRefreshQueryPlan' state. + */ + | 'DrainingBeforeResumeToken' + /** + * Moves to AfterFirstResumeToken after draining. For more info see 'DrainAndRefreshQueryPlan' state. + */ + | 'DrainingAfterResumeToken' + /** + * This state indicates that the stream has finished without error. + */ + | 'Finished' + /** + * This state indicates that a non-retryable error occured and the stream + * cannot be recovered. + */ + | 'Failed'; + +const DEFAULT_TOTAL_TIMEOUT_MS = 60000; + +/** + * This object handles creating and piping the streams + * which are used to process the responses from the server. + * It's main purpose is to make sure that the callerStream, which + * the user gets as a result of Instance.executeQuery, behaves properly: + * - closes in case of a failure + * - doesn't close in case of a retryable error. + * + * We create the following streams: + * responseStream -> byteBuffer -> readerStream -> resultStream + * + * The last two (readerStream and resultStream) are connected + * and returned to the caller - hence called the callerStream. + * + * When a request is made responseStream and byteBuffer are created, + * connected and piped to the readerStream. + * + * On retry, the old responseStream-byteBuffer pair is discarded and a + * new pair is crated. + * + * For more info please refer to the `State` type + */ +export class ExecuteQueryStateMachine { + private bigtable: Bigtable; + private callerStream: CallerStream; + private originalEnd: (chunk?: any, encoding?: any, cb?: () => void) => void; + private retryOptions: StreamRetryOptions; + private valuesStream: AbortableDuplex | null; + private requestParams: any; + private lastPreparedStatementBytes?: Uint8Array | string; + private preparedStatement: PreparedStatement; + private state: State; + private deadlineTs: number; + private protoBytesEncoding?: BufferEncoding; + private numErrors: number; + private retryTimer: NodeJS.Timeout | null; + private timeoutTimer: NodeJS.Timeout | null; + + constructor( + bigtable: Bigtable, + callerStream: CallerStream, + preparedStatement: PreparedStatement, + requestParams: any, + retryOptions?: Partial | null, + protoBytesEncoding?: BufferEncoding, + ) { + this.bigtable = bigtable; + this.callerStream = callerStream; + this.originalEnd = callerStream.end.bind(callerStream); + this.callerStream.end = this.handleCallersEnd.bind(this); + this.requestParams = requestParams; + this.retryOptions = this.parseRetryOptions(retryOptions); + this.deadlineTs = Date.now() + this.retryOptions.totalTimeout; + this.valuesStream = null; + this.preparedStatement = preparedStatement; + this.protoBytesEncoding = protoBytesEncoding; + this.numErrors = 0; + this.retryTimer = null; + this.timeoutTimer = setTimeout( + this.handleTotalTimeout, + this.calculateTotalTimeout(), + ); + + this.state = 'AwaitingQueryPlan'; + this.preparedStatement.getData( + this.handleQueryPlan, + this.calculateTotalTimeout(), + ); + } + + private parseRetryOptions = ( + input?: Partial | null, + ): StreamRetryOptions => { + const rCodes = input?.retryCodes + ? new Set(input?.retryCodes) + : RETRYABLE_STATUS_CODES; + const backoffSettings = input?.backoffSettings; + const clientTotalTimeout = + this?.bigtable?.options?.BigtableClient?.clientConfig?.interfaces && + this?.bigtable?.options?.BigtableClient?.clientConfig?.interfaces[ + 'google.bigtable.v2.Bigtable' + ]?.methods['ExecuteQuery']?.timeout_millis; + return { + maxRetries: backoffSettings?.maxRetries || DEFAULT_RETRY_COUNT, + totalTimeout: + backoffSettings?.totalTimeoutMillis || + clientTotalTimeout || + DEFAULT_TOTAL_TIMEOUT_MS, + retryCodes: rCodes, + initialRetryDelayMillis: + backoffSettings?.initialRetryDelayMillis || + DEFAULT_BACKOFF_SETTINGS.initialRetryDelayMillis, + retryDelayMultiplier: + backoffSettings?.retryDelayMultiplier || + DEFAULT_BACKOFF_SETTINGS.retryDelayMultiplier, + maxRetryDelayMillis: + backoffSettings?.maxRetryDelayMillis || + DEFAULT_BACKOFF_SETTINGS.maxRetryDelayMillis, + }; + }; + + private calculateTotalTimeout = () => { + return Math.max(this.deadlineTs - Date.now(), 0); + }; + + private fail = (err: Error) => { + if (this.state !== 'Failed' && this.state !== 'Finished') { + this.state = 'Failed'; + this.clearTimers(); + this.callerStream.emit('error', err); + } + }; + + private createValuesStream = (): AbortableDuplex => { + const reqOpts: google.bigtable.v2.IExecuteQueryRequest = { + ...this.requestParams, + preparedStatement: this.lastPreparedStatementBytes, + resumeToken: this.callerStream.getLatestResumeToken(), + }; + + const retryOpts = { + currentRetryAttempt: 0, + // Handling retries in this client. + // Options below prevent gax from retrying. + noResponseRetries: 0, + shouldRetryFn: () => { + return false; + }, + }; + + const responseStream = this.bigtable.request({ + client: 'BigtableClient', + method: 'executeQuery', + reqOpts, + gaxOpts: retryOpts, + }); + + const byteBuffer = new ByteBufferTransformer(this.protoBytesEncoding); + const rowValuesStream = pumpify.obj([responseStream, byteBuffer]); + + let aborted = false; + const abort = () => { + if (!aborted) { + aborted = true; + responseStream.abort(); + } + }; + rowValuesStream.abort = abort; + + return rowValuesStream; + }; + + private makeNewRequest = ( + preparedStatementBytes?: Uint8Array | string, + metadata?: SqlTypes.ResultSetMetadata, + ) => { + if (this.valuesStream !== null) { + // assume old streams were scrached. + this.fail( + new Error( + 'Internal error: making a request before streams from the last one was cleaned up.', + ), + ); + } + + if (preparedStatementBytes) { + this.lastPreparedStatementBytes = preparedStatementBytes; + } + if (metadata) { + this.callerStream.updateMetadata(metadata); + } + this.valuesStream = this.createValuesStream(); + + this.valuesStream + .on('error', this.handleStreamError) + .on('data', this.handleStreamData) + .on('close', this.handleStreamEnd) + .on('end', this.handleStreamEnd); + + this.valuesStream.pipe(this.callerStream, {end: false}); + }; + + private discardOldValueStream = () => { + if (this.valuesStream) { + this.valuesStream.abort(); + this.valuesStream.unpipe(this.callerStream); + this.valuesStream.removeAllListeners('error'); + this.valuesStream.removeAllListeners('data'); + this.valuesStream.removeAllListeners('end'); + this.valuesStream.removeAllListeners('close'); + this.valuesStream.destroy(); + this.valuesStream = null; + } + }; + + private getNextRetryDelay = (): number => { + // 0 - 100 ms jitter + const jitter = Math.floor(Math.random() * 100); + const calculatedNextRetryDelay = + this.retryOptions.initialRetryDelayMillis * + Math.pow(this.retryOptions.retryDelayMultiplier, this.numErrors) + + jitter; + + return Math.min( + calculatedNextRetryDelay, + this.retryOptions.maxRetryDelayMillis, + ); + }; + + private clearTimers = (): void => { + if (this.retryTimer) { + clearTimeout(this.retryTimer); + this.retryTimer = null; + } + if (this.timeoutTimer) { + clearTimeout(this.timeoutTimer); + this.timeoutTimer = null; + } + }; + + // Transitions: + + private startNextAttempt = (): void => { + if (this.state === 'DrainAndRefreshQueryPlan') { + this.state = 'AwaitingQueryPlan'; + this.preparedStatement.getData( + this.handleQueryPlan, + this.calculateTotalTimeout(), + ); + } else if (this.state === 'DrainingBeforeResumeToken') { + this.state = 'BeforeFirstResumeToken'; + this.makeNewRequest(this.lastPreparedStatementBytes); + } else if (this.state === 'DrainingAfterResumeToken') { + this.state = 'AfterFirstResumeToken'; + this.makeNewRequest(this.lastPreparedStatementBytes); + } else { + this.fail( + new Error( + `startNextAttempt can't be invoked on a current state ${this.state}`, + ), + ); + } + }; + + private handleDrainingDone = (): void => { + if ( + this.state === 'DrainAndRefreshQueryPlan' || + this.state === 'DrainingBeforeResumeToken' || + this.state === 'DrainingAfterResumeToken' + ) { + this.retryTimer = setTimeout( + this.startNextAttempt, + this.getNextRetryDelay(), + ); + } else { + this.fail( + new Error( + `handleDrainingDone can't be invoked on a current state ${this.state}`, + ), + ); + } + }; + + private handleTotalTimeout = (): void => { + this.discardOldValueStream(); + if (this.retryTimer) { + clearTimeout(this.retryTimer); + this.retryTimer = null; + } + this.fail(new Error('Deadline exceeded.')); + }; + + private handleStreamError = (err: ServiceError): void => { + this.discardOldValueStream(); + if ( + this.retryOptions.retryCodes.has(err.code) || // retryable error + isRstStreamError(err) + ) { + // We want to make a new request only when all requests already written to the Reader by our + // previous active request stream were processed. + this.numErrors += 1; + if (this.numErrors <= this.retryOptions.maxRetries) { + if (this.state === 'AfterFirstResumeToken') { + this.state = 'DrainingAfterResumeToken'; + this.callerStream.onDrain(this.handleDrainingDone); + } else if (this.state === 'BeforeFirstResumeToken') { + this.state = 'DrainingBeforeResumeToken'; + this.callerStream.onDrain(this.handleDrainingDone); + } else { + this.fail( + new Error( + `Can't handle a stream error in the current state ${this.state}`, + ), + ); + } + } else { + this.fail( + new Error(`Maximum retry limit exeeded. Last error: ${err.message}`), + ); + } + } else if (isExpiredQueryError(err)) { + if (this.state === 'AfterFirstResumeToken') { + this.fail(new Error('Query plan expired during a retry attempt.')); + } else if (this.state === 'BeforeFirstResumeToken') { + this.state = 'DrainAndRefreshQueryPlan'; + // If the server returned the "expired query error" we mark it as expired. + this.preparedStatement.markAsExpired(); + this.callerStream.onDrain(this.handleDrainingDone); + } else { + this.fail( + new Error( + `Can't handle expired query error in the current state ${this.state}`, + ), + ); + } + } else { + this.fail(new Error(`Unexpected error: ${err.message}`)); + } + }; + + private handleQueryPlan: PreparedStatementDataCallback = ( + err?: Error, + preparedStatementBytes?: Uint8Array | string, + metadata?: SqlTypes.ResultSetMetadata, + ) => { + if (this.state === 'AwaitingQueryPlan') { + if (err) { + this.numErrors += 1; + if (this.numErrors <= this.retryOptions.maxRetries) { + this.preparedStatement.getData( + this.handleQueryPlan, + this.calculateTotalTimeout(), + ); + } else { + this.fail( + new Error( + `Failed to get query plan. Maximum retry limit exceeded. Last error: ${err.message}`, + ), + ); + } + } else { + this.state = 'BeforeFirstResumeToken'; + this.makeNewRequest(preparedStatementBytes, metadata); + } + } else { + this.fail( + new Error( + `handleQueryPlan can't be invoked on a current state ${this.state}`, + ), + ); + } + }; + + /** + * This method is called when the valuesStream emits data. + * The valuesStream yelds data only after the resume token + * is recieved, hence the state change. + */ + private handleStreamData = (data: any) => { + if ( + this.state === 'BeforeFirstResumeToken' || + this.state === 'AfterFirstResumeToken' + ) { + this.state = 'AfterFirstResumeToken'; + } else { + this.fail( + new Error( + `Internal Error: recieved data in an invalid state ${this.state}`, + ), + ); + } + }; + + private handleStreamEnd = (): void => { + if ( + this.state === 'AfterFirstResumeToken' || + this.state === 'BeforeFirstResumeToken' + ) { + this.clearTimers(); + this.state = 'Finished'; + this.originalEnd(); + } else if (this.state === 'Finished') { + // noop + } else { + this.fail( + new Error( + `Internal Error: Cannot handle stream end in state: ${this.state}`, + ), + ); + } + }; + + /** + * The caller should be able to call callerStream.end() to stop receiving + * more rows and cancel the stream prematurely. However this has a side effect + * the 'end' event will be emitted. + * We don't want that, because it also gets emitted if the stream ended + * normally. To tell these two situations apart, we'll overwrite the end + * function, but save the "original" end() function which will be called + * on valuesStream.on('end'). + */ + private handleCallersEnd = ( + chunk?: any, + encoding?: any, + cb?: () => void, + ): CallerStream => { + if (this.state !== 'Failed' && this.state !== 'Finished') { + this.clearTimers(); + this.discardOldValueStream(); + this.state = 'Finished'; + this.callerStream.close(); + } + return this.callerStream; + }; +} + +export function createCallerStream( + readerStream: ProtobufReaderTransformer, + resultStream: ExecuteQueryStreamWithMetadata, + metadataConsumer: MetadataConsumer, + setCallerCancelled: (v: boolean) => void, +): CallerStream { + const callerStream = pumpify.obj([readerStream, resultStream]); + callerStream.getMetadata = resultStream.getMetadata.bind(resultStream); + callerStream.updateMetadata = metadataConsumer.consume.bind(metadataConsumer); + callerStream.getLatestResumeToken = () => readerStream.resumeToken; + callerStream.onDrain = readerStream.onDrain.bind(readerStream); + callerStream.close = () => { + setCallerCancelled(true); + callerStream.destroy(); + }; + return callerStream; +} diff --git a/src/execute-query/metadataconsumer.ts b/src/execute-query/metadataconsumer.ts new file mode 100644 index 000000000..8fa3ecc2b --- /dev/null +++ b/src/execute-query/metadataconsumer.ts @@ -0,0 +1,122 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +import * as Types from './types'; +// eslint-disable-next-line @typescript-eslint/no-var-requires +import {google} from '../../protos/protos'; + +/** + * This class keeps and parses the metadata. + */ +export class MetadataConsumer { + private metadata: Types.ResultSetMetadata | null; + + constructor() { + this.metadata = null; + } + + getMetadata = (): Types.ResultSetMetadata | null => { + return this.metadata; + }; + + consume = (new_metadata: Types.ResultSetMetadata) => { + this.metadata = new_metadata; + }; + + static parsePBType(type: google.bigtable.v2.Type): Types.Type { + switch (type.kind) { + case 'bytesType': + return Types.Bytes(); + case 'stringType': + return Types.String(); + case 'int64Type': + return Types.Int64(); + case 'float64Type': + return Types.Float64(); + case 'float32Type': + return Types.Float32(); + case 'boolType': + return Types.Bool(); + case 'timestampType': + return Types.Timestamp(); + case 'dateType': + return Types.Date(); + case 'structType': + return Types.Struct( + ...type.structType!.fields!.map(field => ({ + name: field.fieldName as string | null, + type: MetadataConsumer.parsePBType( + field.type as google.bigtable.v2.Type, + ), + })), + ); + case 'arrayType': + return Types.Array( + MetadataConsumer.parsePBType( + type.arrayType!.elementType! as google.bigtable.v2.Type, + ), + ); + case 'mapType': { + const keyType = MetadataConsumer.parsePBType( + type.mapType!.keyType! as google.bigtable.v2.Type, + ); + if ( + keyType.type !== 'int64' && + keyType.type !== 'string' && + keyType.type !== 'bytes' + ) { + throw new Error( + `Unsupported type of map key received: ${keyType.type}`, + ); + } + return Types.Map( + keyType, + MetadataConsumer.parsePBType( + type.mapType!.valueType! as google.bigtable.v2.Type, + ), + ); + } + default: + throw new Error( + `Type ${type.kind} not supported by current client version`, + ); + } + } + + static parseMetadata( + metadata: google.bigtable.v2.IResultSetMetadata, + ): Types.ResultSetMetadata { + if (!metadata.protoSchema) { + throw new Error('Only protoSchemas are supported.'); + } + const columns = metadata.protoSchema.columns!; + if (columns.length === 0) { + throw new Error('Invalid empty ResultSetMetadata received.'); + } + + return Types.ResultSetMetadata.fromTuples( + columns.map(column => { + if (column.name === null || column.name === '') { + throw new Error(`Invalid column name "${column.name}"`); + } else { + return [ + column.name ?? null, + MetadataConsumer.parsePBType( + column.type! as google.bigtable.v2.Type, + ), + ]; + } + }), + ); + } +} diff --git a/src/execute-query/namedlist.ts b/src/execute-query/namedlist.ts new file mode 100644 index 000000000..b745c537a --- /dev/null +++ b/src/execute-query/namedlist.ts @@ -0,0 +1,98 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/** + * Represents how field names correspond to field indexes in a NamedList. + */ +export type FieldMapping = { + validFieldNames: Map; + duplicateFieldNames: Map; + fieldNames: (string | null)[]; +}; + +function constructFieldMapping(values: (string | null)[]): FieldMapping { + const fieldMapping: Map = new Map(); + for (let i = 0; i < values.length; i++) { + const name = values[i]; + if (name) { + if (!fieldMapping.has(name)) { + fieldMapping.set(name, []); + } + fieldMapping.get(name)!.push(i); + } + } + const validFieldNames = new Map(); + const duplicateFieldNames = new Map(); + for (const [name, indexes] of fieldMapping.entries()) { + if (indexes.length > 1) { + duplicateFieldNames.set(name, indexes); + } else { + validFieldNames.set(name, indexes[0]); + } + } + return { + validFieldNames, + duplicateFieldNames, + fieldNames: values, + }; +} + +/** + * Class representing a list which allows retrieving elements both by index + * and by name. If multiple elements have the same name, they have to be + * retrieved by index. Otherwise an error is thrown. + */ +export class NamedList { + values: Array; + fieldMapping: FieldMapping; + + constructor(values: Array, fieldMapping: FieldMapping) { + this.values = values; + this.fieldMapping = fieldMapping; + } + + protected static _fromTuples, T>( + type: {new (values: Array, fieldMapping: FieldMapping): R}, + tuples: [string | null, T][], + ): R { + return new type( + tuples.map(tuple => tuple[1]), + constructFieldMapping(tuples.map(tuple => tuple[0])), + ); + } + + get(indexOrName: string | number): T { + let index; + if (typeof indexOrName === 'string') { + if (this.fieldMapping.duplicateFieldNames.has(indexOrName)) { + throw new Error( + `Cannot access ${indexOrName} by name because it is available on multiple indexes: ${this.fieldMapping.duplicateFieldNames + .get(indexOrName)! + .join(', ')}`, + ); + } + index = this.fieldMapping.validFieldNames.get(indexOrName); + if (index === undefined) { + throw new Error(`Unknown field name '${indexOrName}'.`); + } + } else { + index = indexOrName; + } + return this.values[index]; + } + + getFieldNameAtIndex(index: number): string | null { + return this.fieldMapping.fieldNames[index]; + } +} diff --git a/src/execute-query/parameterparsing.ts b/src/execute-query/parameterparsing.ts new file mode 100644 index 000000000..d0fc35be2 --- /dev/null +++ b/src/execute-query/parameterparsing.ts @@ -0,0 +1,308 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {PreciseDate} from '@google-cloud/precise-date'; +import {google} from '../../protos/protos'; +import {BigtableDate, ExecuteQueryParameterValue} from './values'; +import * as SqlTypes from './types'; +import * as is from 'is'; +import Long = require('long'); + +/** + * Creates protobuf objects with explicit types from passed parameters. + * The protobuf Value objects have a field describing their type explicitly. + * For each param we create a Value object based on the provided type. + * @param parameters map from parameter name to parameter value + * @param parameter_types map from parameter name to parameter type + * @returns map from parameter name to a Value object + */ +export function parseParameters( + parameters: {[param: string]: ExecuteQueryParameterValue}, + parameterTypes: {[param: string]: SqlTypes.Type}, +): {[param: string]: google.bigtable.v2.IValue} { + // Assert both objects contain the same keys: + const parameterKeys = Object.keys(parameters); + const parameterTypeKeys = Object.keys(parameterTypes); + if (parameterKeys.length !== parameterTypeKeys.length) { + throw new Error( + `Number of parameters (${parameterKeys.length}) does not match number of parameter types (${parameterTypeKeys.length}).`, + ); + } + // if the numbers of keys match, but keys differ we will catch it in the next step + const entries: [string, google.bigtable.v2.IValue][] = []; + for (const [key, value] of Object.entries(parameters)) { + let type: SqlTypes.Type; + if (Object.prototype.hasOwnProperty.call(parameterTypes, key)) { + type = parameterTypes[key]; + } else { + throw new Error(`Unrecognized parameter: ${key}`); + } + + entries.push([key, setTypeField(convertJsValueToValue(value, type), type)]); + } + return Object.fromEntries(entries); +} + +export function parseParameterTypes(parameter_types: { + [param: string]: SqlTypes.Type; +}): {[param: string]: google.bigtable.v2.IType} { + return Object.fromEntries( + Object.entries(parameter_types).map(([key, value]) => [ + key, + executeQueryTypeToPBType(value), + ]), + ); +} + +function inferType(value: ExecuteQueryParameterValue): SqlTypes.Type { + if (is.number(value)) { + return SqlTypes.Float64(); + } else if (typeof value === 'bigint') { + return SqlTypes.Int64(); + } else if (is.string(value)) { + return SqlTypes.String(); + } else if (is.boolean(value)) { + return SqlTypes.Bool(); + } else if (is.array(value)) { + // eslint-disable-next-line + throw new Error( + `Cannot infer type of an array ${value}. Please provide a type hint using parameter_types.`, + ); + } else if (typeof value === 'object') { + if (value instanceof Uint8Array) { + return SqlTypes.Bytes(); + } else if (value instanceof PreciseDate) { + return SqlTypes.Timestamp(); + } else if (value instanceof Date) { + throw new Error( + 'Date is not supported as a parameter type. Please use PreciseDate for Sql TIMESTAMP or BigtableDate for SQL DATE', + ); + } else if (value instanceof BigtableDate) { + return SqlTypes.Date(); + } + } + + const typeString = typeof value; + let prototypeString = null; + if (typeString === 'object') { + if (value === null) { + prototypeString = 'null'; + } else { + prototypeString = `constructor.name = ${value.constructor.name}`; + } + } + const typeInfo = `typeof == ${typeString}${ + prototypeString ? `, ${prototypeString}` : '' + })`; + throw new Error( + `Cannot infer type of ${value} (${typeInfo}). Please provide a type hint using parameter_types.`, + ); +} + +export function setTypeField( + value: google.bigtable.v2.IValue, + type: SqlTypes.Type, +): google.bigtable.v2.IValue { + value.type = executeQueryTypeToPBType(type); + return value; +} + +export function executeQueryTypeToPBType( + type: SqlTypes.Type, +): google.bigtable.v2.IType { + switch (type.type) { + case 'string': + return {stringType: {}}; + case 'int64': + return {int64Type: {}}; + case 'float32': + return {float32Type: {}}; + case 'float64': + return {float64Type: {}}; + case 'bytes': + return {bytesType: {}}; + case 'bool': + return {boolType: {}}; + case 'timestamp': + return {timestampType: {}}; + case 'date': + return {dateType: {}}; + case 'array': + return { + arrayType: {elementType: executeQueryTypeToPBType(type.elementType)}, + }; + case 'struct': + return { + structType: { + fields: type.values.map((value, index) => ({ + fieldName: type.getFieldNameAtIndex(index), + type: executeQueryTypeToPBType(value), + })), + }, + }; + case 'map': + return { + mapType: { + keyType: executeQueryTypeToPBType(type.keyType), + valueType: executeQueryTypeToPBType(type.valueType), + }, + }; + } +} + +export function convertJsValueToValue( + value: ExecuteQueryParameterValue, + type: SqlTypes.Type, +): google.bigtable.v2.IValue { + if (value === null) { + return {}; + } + + switch (type.type) { + case 'string': + return convertToString(value); + case 'int64': + return convertToInt64(value); + case 'float32': + return convertToFloat64(value); + case 'float64': + return convertToFloat64(value); + case 'bytes': + return convertToBytes(value); + case 'bool': + return convertToBool(value); + case 'timestamp': + return convertToTimestamp(value); + case 'date': + return convertToDate(value); + case 'array': + return convertToArray(value, type); + case 'struct': + throw new Error('Struct is not a supported query param type'); + case 'map': + throw new Error('Map is not a supported query param type'); + } +} + +function convertToString( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (is.string(value)) { + return {stringValue: value as string}; + } + throw new Error(`Value ${value} cannot be converted to string.`); +} + +const MAX_LONG = BigInt(Long.MAX_VALUE.toString()); +const MIN_LONG = BigInt(Long.MIN_VALUE.toString()); + +function bigintToLong(value: bigint): Long { + // Long fromString does not check this + if (value > MAX_LONG || value < MIN_LONG) { + throw new Error( + `Value ${value} cannot be converted to int64 - it is out of range.`, + ); + } + return Long.fromString(value.toString()); +} + +function convertToInt64( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (typeof value === 'bigint') { + return { + intValue: bigintToLong(value), + }; + } else if (typeof value === 'number') { + throw new Error( + `Value ${value} cannot be converted to int64 - argument of type INT64 should by passed as BigInt.`, + ); + } + throw new Error(`Value ${value} cannot be converted to int64.`); +} + +function convertToFloat64( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (typeof value === 'number') { + return {floatValue: value}; + } + throw new Error(`Value ${value} cannot be converted to float64.`); +} + +function convertToBytes( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (typeof value === 'object' && value instanceof Uint8Array) { + return {bytesValue: value}; + } + throw new Error(`Value ${value} cannot be converted to bytes.`); +} + +function convertToBool( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (typeof value === 'boolean') { + return {boolValue: value}; + } + throw new Error(`Value ${value} cannot be converted to boolean.`); +} + +function convertToTimestamp( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (typeof value === 'object' && value instanceof PreciseDate) { + return {timestampValue: value.toStruct()}; + } + throw new Error( + `Value ${value} cannot be converted to timestamp, please use PreciseDate instead.`, + ); +} + +function convertToDate( + value: ExecuteQueryParameterValue, +): google.bigtable.v2.IValue { + if (typeof value === 'object' && value instanceof BigtableDate) { + return {dateValue: value}; + } + throw new Error(`Value ${value} cannot be converted to date.`); +} + +function convertToArray( + value: ExecuteQueryParameterValue, + type: SqlTypes.ArrayType, +): google.bigtable.v2.IValue { + if (!is.array(value)) { + throw new Error(`Value ${value} cannot be converted to an array.`); + } + const arrayValue = value as Array; + return { + arrayValue: { + values: arrayValue.map((element, index) => { + try { + return convertJsValueToValue(element, type.elementType); + // eslint-disable-next-line + } catch (conversionError: any) { + if (conversionError instanceof Error) { + throw new Error( + `Error while converting element ${index} of an array: ${conversionError.message}`, + ); + } else { + throw conversionError; + } + } + }), + }, + }; +} diff --git a/src/execute-query/preparedstatement.ts b/src/execute-query/preparedstatement.ts new file mode 100644 index 000000000..8a2dbc040 --- /dev/null +++ b/src/execute-query/preparedstatement.ts @@ -0,0 +1,287 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +import {Bigtable} from '..'; +import * as SqlTypes from './types'; +import {PreciseDate} from '@google-cloud/precise-date'; + +import {google} from '../../protos/protos'; +import {MetadataConsumer} from './metadataconsumer'; +import {EventEmitter} from 'events'; +import {ServiceError, CallOptions} from 'google-gax'; + +export const SHOULD_REFRESH_SOON_PERIOD_MS = 1000; + +export type PreparedStatementDataCallback = ( + err?: Error, + preparedQueryBytes?: Uint8Array | string, + metadata?: SqlTypes.ResultSetMetadata, +) => void; + +interface IRetryRequest { + client: string; + method: string; + reqOpts: google.bigtable.v2.IPrepareQueryRequest; + retryOpts?: CallOptions; +} + +/** + * This object keeps track of the query plan a.k.a. metadata and preparedQuery bytes. + * It provides a way of retrieving last retrieved query plan. + * If a query plan is marked as expired, it will be refreshed. + * You can get the query plan via the getData method. + * If the query plan is not expired, getData will return the value immediately. + * If the object is marked as expired, getting the query plan will wait for + * a refresh to happen. If the refresh fails, all awaiting getData calls + * also return an error. + */ +export class PreparedStatement extends EventEmitter { + private bigtable: Bigtable; + private retryRequest: IRetryRequest; + private metadata: SqlTypes.ResultSetMetadata; + private preparedQueryBytes: Uint8Array | string; + private validUntilTimestamp: number | null; + private forcedExpiration: boolean; + private isRefreshing: boolean; + private timer: NodeJS.Timeout | null; + private lastRefreshError: ServiceError | Error | null; + private parameterTypes: {[param: string]: SqlTypes.Type}; + + constructor( + bigtable: Bigtable, + response: google.bigtable.v2.PrepareQueryResponse, + retryRequest: IRetryRequest, + parameterTypes: {[param: string]: SqlTypes.Type}, + ) { + super(); + this.bigtable = bigtable; + this.metadata = MetadataConsumer.parseMetadata(response.metadata!); + this.preparedQueryBytes = response.preparedQuery; + this.validUntilTimestamp = timestampFromResponse(response); + this.timer = null; + this.isRefreshing = false; + this.lastRefreshError = null; + this.forcedExpiration = false; + this.retryRequest = retryRequest; + this.parameterTypes = parameterTypes; + } + + /** + * Returns true if the validUntilTimestamp is close, + * meaning less than SHOULD_REFRESH_SOON_PERIOD_MS away. + */ + private shouldRefreshSoon = (): boolean => { + if (!this.validUntilTimestamp) { + return false; + } + return ( + Date.now() > this.validUntilTimestamp - SHOULD_REFRESH_SOON_PERIOD_MS + ); + }; + + /** + * Schedules the refresh. It is deffered to the next tick to ensure + * that the current call stack is finished before a request to bigtable is made. + */ + private setupTimer = (): void => { + this.timer = setTimeout(this.handleTimerEnd, 0); + }; + + private discardTimer = (): void => { + if (this.timer) { + clearTimeout(this.timer); + this.timer = null; + } + }; + + /** + * Performs a request to bigtable to get a refreshed query plan. + */ + private startRefreshing = (): void => { + this.isRefreshing = true; + this.bigtable.request(this.retryRequest, this.handlePrepareQueryResponse); + }; + + /** + * Begins the refresh. + */ + private handleTimerEnd = (): void => { + if (!this.isRefreshing) { + this.discardTimer(); + this.startRefreshing(); + } + }; + + /** + * Callback for handling the call to bigtable. + */ + private handlePrepareQueryResponse = ( + err: ServiceError | null, + response?: google.bigtable.v2.PrepareQueryResponse, + ): void => { + if (this.isRefreshing) { + this.isRefreshing = false; + this.discardTimer(); + if (err) { + this.lastRefreshError = err; + this.emit('refreshDone'); + } else { + try { + this.lastRefreshError = null; + this.forcedExpiration = false; + this.validUntilTimestamp = timestampFromResponse(response!); + this.metadata = MetadataConsumer.parseMetadata(response!.metadata!); + this.preparedQueryBytes = response!.preparedQuery; + } catch (err: any) { + this.lastRefreshError = err as Error; + } + this.emit('refreshDone'); + } + } else { + const err = new Error( + 'Invalid state: PrepareQueryResponse recieved when not refreshing.', + ); + console.error(err); + throw err; + } + }; + + /** + * Invoked when the query plan is retrieved from this object. + */ + private scheduleRefreshIfNeeded = (): void => { + if (!this.isRefreshing && this.timer === null) { + if (this.isExpired() || this.shouldRefreshSoon()) { + this.setupTimer(); + } // else noop + } // else noop + }; + + /** + * This function should be called, when the server returns + * the FAILED_PRECONDITION error saying the query plan + * is expired. For more info refer to the ExecuteQueryStateMachine. + */ + markAsExpired = (): void => { + this.forcedExpiration = true; + }; + + /** + * Used for retrieveing the query plan (preparedQuery bytes and metadata) + * @param callback called when query plan is available + * @param timeoutMs when callback should be called with an error. + */ + getData = ( + callback: PreparedStatementDataCallback, + timeoutMs: number, + ): void => { + this.scheduleRefreshIfNeeded(); + if (this.isExpired()) { + const listener = new CallbackWithTimeout(callback, timeoutMs); + this.once('refreshDone', () => { + // If there are many listeners, the query plan could have expired again + // before we got to processing this one, so we have to check it again. + if (this.isExpired() || this.lastRefreshError) { + listener.tryInvoke( + this.lastRefreshError || + new Error('Getting a fresh query plan failed.'), + undefined, + undefined, + ); + } else { + listener.tryInvoke(undefined, this.preparedQueryBytes, this.metadata); + } + }); + } else { + // for the sake of consistency we should call the callback asynchornously + // regardless if the plan needs refreshing or not. + setTimeout( + () => callback(undefined, this.preparedQueryBytes, this.metadata), + 0, + ); + } + }; + + /** + * @returns parameter types used to create the query plan + */ + getParameterTypes = (): {[param: string]: SqlTypes.Type} => + this.parameterTypes; + + /** + * @returns true if the object has been marked as expired. + */ + isExpired = (): boolean => { + return this.forcedExpiration; + }; +} + +/** + * This class makes sure the callback is called only once. + * If the timeout expired, the callback is called with a "Timeout Expired" error. + * Otherwise it is called with provided args. + */ +class CallbackWithTimeout { + private callback: PreparedStatementDataCallback | null; + private timer: NodeJS.Timeout | null; + private isValid: boolean; + + constructor(callback: PreparedStatementDataCallback, timeout: number) { + this.callback = callback; + this.isValid = true; + this.timer = setTimeout(() => { + this.tryInvoke( + new Error( + 'Deadline Exceeded waiting for prepared statement to refresh.', + ), + ); + }, timeout); + } + + /** + * If this object has not yet been invalidated, the callback is called. + * @param args + */ + tryInvoke(...args: Parameters): void { + if (!this.isValid || !this.callback) { + return; + } + const callback = this.callback; + this.invalidate(); + callback(...args); + } + + /** + * After this method is called, the callback can no longer be invoked. + */ + private invalidate(): void { + if (this.timer) { + clearTimeout(this.timer); + this.timer = null; + } + this.callback = null; + this.isValid = false; + } +} + +function timestampFromResponse( + response: google.bigtable.v2.PrepareQueryResponse, +): number | null { + if (!response.validUntil?.seconds) { + return null; + } + return new PreciseDate({ + seconds: response.validUntil?.seconds ?? undefined, + nanos: response.validUntil?.nanos ?? undefined, + }).getTime(); +} diff --git a/src/execute-query/protobufreadertransformer.ts b/src/execute-query/protobufreadertransformer.ts new file mode 100644 index 000000000..7349b095a --- /dev/null +++ b/src/execute-query/protobufreadertransformer.ts @@ -0,0 +1,135 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {Transform, TransformOptions, TransformCallback} from 'stream'; +// eslint-disable-next-line @typescript-eslint/no-var-requires +import {google} from '../../protos/protos'; +import {MetadataConsumer} from './metadataconsumer'; + +class DrainGuard { + callback: () => void; + + constructor(callback: () => void) { + this.callback = callback; + } +} + +export type BatchAndToken = [Buffer[], Uint8Array]; + +/** + * This transformer is responsible for deserializing bytes sent from the + * server to an appropriate object which can be used to construct rows. + * Right now only google.bigtable.v2.ProtoRows is supported. + */ +export class ProtobufReaderTransformer extends Transform { + metadataConsumer: MetadataConsumer; + resumeToken: Uint8Array | null; + + constructor(metadataConsumer: MetadataConsumer, opts?: TransformOptions) { + super({...opts, objectMode: true, highWaterMark: 1024}); + this.metadataConsumer = metadataConsumer; + this.resumeToken = null; + } + + _transform( + batchAndToken: BatchAndToken | DrainGuard, + _encoding: BufferEncoding, + callback: TransformCallback, + ) { + if (batchAndToken instanceof DrainGuard) { + batchAndToken.callback(); + } else { + const maybeMetadata = this.metadataConsumer.getMetadata(); + if (maybeMetadata) { + const [batches, resumeToken] = batchAndToken; + this.resumeToken = resumeToken; + const valuesBuffer: google.bigtable.v2.IValue[] = []; + + const expectedLength = maybeMetadata.columns.length; + + for (let batchIdx = 0; batchIdx < batches.length; batchIdx++) { + const batch = batches[batchIdx]; + const protoRows = google.bigtable.v2.ProtoRows.decode(batch); + + if (protoRows.values.length % expectedLength) { + callback(new Error('Internal error - received incomplete row.')); + return; + } + + for ( + let valueIdx = 0; + valueIdx < protoRows.values.length; + valueIdx++ + ) { + valuesBuffer.push(protoRows.values[valueIdx]); + } + } + + if (valuesBuffer.length > 0) { + for (let i = 0; i < valuesBuffer.length; i += expectedLength) { + this.push(valuesBuffer.slice(i, i + expectedLength)); + } + } + } else { + return callback(new Error('Internal error - missing metadata')); + } + } + + callback(); + } + + /** + * @param callback guaranteed to be called **after** the last message + * was processed. + */ + onDrain = (callback: () => void) => { + // Writable streams keep a buffer of objects to + // process (in case of a Transform processing means calling _transform() method). Readable streams + // keep a buffer of objects to be read by downstream processors. Transforms are both Writable and + // Readable, thus they have one buffer for parameters to, and one buffer for results of + // _transform() method. + // Objects can end up in a writeable's buffer if they are written after previous write call + // returned false but before 'drain' event is emitted or when a write happens while another + // objects is processed. + // Objects can end up in a readable's buffer if there are no downstream processors ready to accept + // new objects (either there are none or at least one of them is paused). + // + // Our data pipeline looks as follows: + // bigtable stream -> ByteBuffer transform -> Reader transform -> ... + // + // But if we include buffers in this diagram this becomes more complicated: + // + // (bigtable stream -> [readable buffer]) + // -> ([writable buffer] -> ByteBuffer transform() -> [readable buffer]) + // -> ([writable buffer] -> Reader transform() -> [readable buffer]) + // -> ... + // + // and each of these buffers can buffer requests that were (in readable) or were not (in readable) + // already passed to _transform() method. + // During the retry we have to recreate bigtable stream and discard all data stored in + // the ByteBuffer, and perform a new request with an appropriate resumeToken. + // + // Passing an appropriate resumeToken is crucial to prevent duplicate or lost responses. + // + // So, how to obtain a resumption token? Let's try a few options: + // We cannot take last resumeToken that was seen by ByteBuffer's _transform() method + // - it is possible that there are some unprocessed events in ByteBuffer's writable buffer + // that will be processed at some point. + // The same applies to Reader's _transform(), writeable buffers are still there. + // + // Thus we see that we have to consider events waiting in the buffers and wait until they are + // processed. + this.write(new DrainGuard(callback)); + }; +} diff --git a/src/execute-query/queryresultrowtransformer.ts b/src/execute-query/queryresultrowtransformer.ts new file mode 100644 index 000000000..3cd277ccb --- /dev/null +++ b/src/execute-query/queryresultrowtransformer.ts @@ -0,0 +1,255 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// eslint-disable-next-line @typescript-eslint/no-var-requires +import {google} from '../../protos/protos'; +import Long = require('long'); + +import { + EncodedKeyMap, + BigtableDate, + ExecuteQueryStreamWithMetadata, + SqlValue, + QueryResultRow, + Struct, + ensureUint8Array, +} from './values'; +import {ResultSetMetadata, Type} from './types'; +import {PreciseDate} from '@google-cloud/precise-date'; +import assert = require('assert'); +import {Transform, TransformCallback, TransformOptions} from 'stream'; +import {MetadataConsumer} from './metadataconsumer'; +import {FieldMapping} from './namedlist'; + +/** + * Class representing a readable stream with ExecuteQuery results + * which also lets the caller get metadata. + */ +export class ExecuteQueryStreamTransformWithMetadata + extends Transform + implements ExecuteQueryStreamWithMetadata +{ + metadataConsumer: MetadataConsumer; + fieldMapping: FieldMapping | null; + hasCallerCancelled: () => boolean; + protoBytesEncoding?: BufferEncoding; + + constructor( + metadataConsumer: MetadataConsumer, + hasCallerCancelled: () => boolean, + protoBytesEncoding?: BufferEncoding, + opts?: TransformOptions, + ) { + super({...opts, objectMode: true, highWaterMark: 0}); + this.fieldMapping = null; + this.metadataConsumer = metadataConsumer; + this.hasCallerCancelled = hasCallerCancelled; + this.protoBytesEncoding = protoBytesEncoding; + } + + valueToJsType(value: google.bigtable.v2.Value, metadata: Type): SqlValue { + if (!value.kind) { + return null; + } + switch (metadata.type) { + case 'bytes': + if (value.kind === 'bytesValue') { + return ensureUint8Array(value.bytesValue!, this.protoBytesEncoding); + } + break; + case 'string': + if (value.kind === 'stringValue') { + return value.stringValue!; + } + break; + case 'int64': + if (value.kind === 'intValue') { + return intValueToBigInt(value.intValue!); + } + break; + case 'bool': + if (value.kind === 'boolValue') { + return value.boolValue!; + } + break; + case 'float32': + case 'float64': + if (value.kind === 'floatValue') { + return value.floatValue!; + } + break; + case 'timestamp': + if (value.kind === 'timestampValue') { + return new PreciseDate({ + seconds: value.timestampValue!.seconds ?? undefined, + nanos: value.timestampValue!.nanos ?? undefined, + }); + } + break; + case 'date': + if (value.kind === 'dateValue') { + return new BigtableDate( + value.dateValue!.year || 0, + value.dateValue!.month || 0, + value.dateValue!.day || 0, + ); + } + break; + case 'array': + return this.valueToJsArray(value, metadata); + case 'struct': + return this.valueToJsStruct(value, metadata); + case 'map': + return this.valueToJsMap(value, metadata); + default: + throw new Error( + `Unexpected type to parse: ${JSON.stringify(metadata)}`, + ); + } + throw new Error(`Metadata and Value not matching. + Metadata:${metadata} + Value:${value}`); + } + + valueToJsArray(value: google.bigtable.v2.Value, metadata: Type): SqlValue { + assert(metadata.type === 'array'); + if ( + value.arrayValue === null || + value.arrayValue === undefined || + value.arrayValue.values === null || + value.arrayValue.values === undefined + ) { + return null; + } + return value.arrayValue.values.map(value => + this.valueToJsType( + value as google.bigtable.v2.Value, + metadata.elementType, + ), + ); + } + + valueToJsStruct(value: google.bigtable.v2.Value, metadata: Type): SqlValue { + assert(metadata.type === 'struct'); + if ( + value.arrayValue === null || + value.arrayValue === undefined || + value.arrayValue.values === null || + value.arrayValue.values === undefined + ) { + return null; + } + if (value.arrayValue.values.length !== metadata.values.length) { + throw new Error( + `Internal error - received Struct with ${value.arrayValue.values.length} values, but metadata has ${metadata.values.length} fields.`, + ); + } + return new Struct( + value.arrayValue.values.map((value, index) => + this.valueToJsType( + value as google.bigtable.v2.Value, + metadata.get(index), + ), + ), + metadata.fieldMapping, + ); + } + + valueToJsMap(value: google.bigtable.v2.Value, metadata: Type): SqlValue { + assert(metadata.type === 'map'); + if ( + value.arrayValue === null || + value.arrayValue === undefined || + value.arrayValue.values === null || + value.arrayValue.values === undefined + ) { + return null; + } + if ( + metadata.keyType.type !== 'int64' && + metadata.keyType.type !== 'string' && + metadata.keyType.type !== 'bytes' + ) { + throw new Error( + `Internal error - unsupported type of key received: ${metadata.keyType.type}`, + ); + } + const values: google.bigtable.v2.Value[] = value.arrayValue + .values as google.bigtable.v2.Value[]; + return new EncodedKeyMap( + values.map(value => { + // Types are ensured by checking metadata.keyType.type earlier. + const pair = value?.arrayValue?.values as google.bigtable.v2.Value[]; + const keyValue = this.valueToJsType(pair[0], metadata.keyType) as + | bigint + | string + | Uint8Array + | null; + return [keyValue, this.valueToJsType(pair[1], metadata.valueType)]; + }), + ); + } + + getFieldMapping(): FieldMapping { + if (this.fieldMapping === null) { + const metadata = this.getMetadata(); + if (metadata === null) { + throw new Error('Metadata was not sent by the server.'); + } + this.fieldMapping = metadata.fieldMapping; + } + return this.fieldMapping; + } + + _transform( + chunk: Array, + _encoding: BufferEncoding, + callback: TransformCallback, + ) { + let error: Error | null = null; + try { + if (!this.hasCallerCancelled()) { + const maybeMetadata = this.metadataConsumer.getMetadata(); + if (maybeMetadata) { + this.push( + new QueryResultRow( + chunk.map((value, index) => + this.valueToJsType(value, maybeMetadata.get(index)), + ), + this.getFieldMapping(), + ), + ); + } else { + throw new Error( + 'Server error - expected to receive metadata by now.', + ); + } + } + } catch (e) { + error = e as Error; + } + callback(error); + } + + getMetadata(): ResultSetMetadata | null { + return this.metadataConsumer.getMetadata(); + } +} + +function intValueToBigInt(intValue: string | number | Long): bigint { + if (intValue instanceof Long) { + return BigInt(intValue.toString()); + } + return BigInt(intValue); +} diff --git a/src/execute-query/types.ts b/src/execute-query/types.ts new file mode 100644 index 000000000..e374bd4bb --- /dev/null +++ b/src/execute-query/types.ts @@ -0,0 +1,83 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {NamedList} from './namedlist'; + +export type ArrayType = {type: 'array'; elementType: Type}; +export type MapType = {type: 'map'; keyType: Type; valueType: Type}; +export type FieldType = {name: string | null; type: Type}; +export class StructType extends NamedList { + type: 'struct' = 'struct' as const; + + static fromTuples(tuples: [string | null, Type][]): StructType { + return NamedList._fromTuples(StructType, tuples); + } +} + +export type Int64Type = ReturnType; +export type Float64Type = ReturnType; +export type Float32Type = ReturnType; +export type BytesType = ReturnType; +export type StringType = ReturnType; +export type BoolType = ReturnType; +export type TimestampType = ReturnType; +export type DateType = ReturnType; + +/** + * Factory functions are provided instead of constants + * for all types for coherence and for extensibility + * (we need parameters at least for arrays, structs and maps) + */ +export const Int64 = () => ({type: 'int64' as const}); +export const Float64 = () => ({type: 'float64' as const}); +export const Float32 = () => ({type: 'float32' as const}); +export const Bytes = () => ({type: 'bytes' as const}); +export const String = () => ({type: 'string' as const}); +export const Bool = () => ({type: 'bool' as const}); +export const Timestamp = () => ({type: 'timestamp' as const}); +export const Date = () => ({type: 'date' as const}); +export const Struct = (...fields: FieldType[]): StructType => + StructType.fromTuples(fields.map(value => [value.name, value.type])); +export const Array = (elementType: Type): ArrayType => ({ + type: 'array' as const, + elementType, +}); +export const Map = (keyType: Type, valueType: Type): MapType => ({ + type: 'map' as const, + keyType, + valueType, +}); + +export type Type = + | Int64Type + | Float32Type + | Float64Type + | BytesType + | StringType + | BoolType + | TimestampType + | DateType + | ArrayType + | StructType + | MapType; + +export class ResultSetMetadata extends NamedList { + get columns(): Array { + return this.values; + } + + static fromTuples(tuples: [string | null, Type][]): ResultSetMetadata { + return NamedList._fromTuples(ResultSetMetadata, tuples); + } +} diff --git a/src/execute-query/values.ts b/src/execute-query/values.ts new file mode 100644 index 000000000..1186c984a --- /dev/null +++ b/src/execute-query/values.ts @@ -0,0 +1,224 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {Duplex, Readable} from 'stream'; +import * as SqlTypes from './types'; +import {PreciseDate} from '@google-cloud/precise-date'; +import {NamedList} from './namedlist'; +const CRC32C = require('fast-crc32c'); + +export type BigtableMap = EncodedKeyMap; + +export type ExecuteQueryParameterValue = + | Uint8Array + | string + | bigint + | number + | boolean + | PreciseDate + | BigtableDate + | Array + | null; + +export type SqlValue = + | ExecuteQueryParameterValue + | Struct + | Array + | BigtableMap; + +/** + * A custom class created to allow setting year, month and date to 0. + */ +export class BigtableDate { + year: number; + month: number; + day: number; + + constructor(year: number, month: number, day: number) { + if (year < 0) { + throw new Error('Invalid year.'); + } + this.year = year; + + if (month < 0 || month > 12) { + throw new Error('Invalid month.'); + } + this.month = month; + + if (day < 0 || day > 31) { + throw new Error('Invalid month.'); + } + this.day = day; + } +} + +export class QueryResultRow extends NamedList {} + +export class Struct extends NamedList { + static fromTuples(tuples: [string | null, SqlValue][]): Struct { + return NamedList._fromTuples(Struct, tuples); + } +} + +export interface ExecuteQueryStreamWithMetadata extends Duplex { + getMetadata: () => SqlTypes.ResultSetMetadata | null; + end: () => this; +} + +export function checksumValid( + buffer: Buffer, + expectedChecksum: number, +): boolean { + return CRC32C.calculate(buffer) === expectedChecksum; +} + +export function ensureUint8Array( + bytes: Uint8Array | string, + encoding?: BufferEncoding, +): Uint8Array { + return bytes instanceof Uint8Array + ? bytes + : Buffer.from(bytes, encoding || 'base64'); +} + +function _parseBufferToMapKey( + key: bigint | string | Uint8Array | null, +): bigint | string | null { + // Uint8Array is always an instance of Buffer, + // but we keep it in the if condition for TS linter's sake + if (key instanceof Buffer || key instanceof Uint8Array) { + const base64Key = key.toString('base64'); + return base64Key; + } else if (key === null) { + return null; + } else { + return key; + } +} + +type MapKey = bigint | string | Uint8Array | null; +// internal representation of the map contains the original key and value tuple +type _MapValue = [MapKey, SqlValue]; + +export class EncodedKeyMap + implements Map +{ + private map_impl: Map; + /** + * Class representing a Map Value returned by ExecuteQuery. Native JS Map + * does not support Buffer comparison - it compares object id and not + * buffer values. This class solves this by encoding Buffer keys as + * base64 strings. + * Please note that an empty string and an empty buffer have the same + * representation, however, we do not ever use mixed key types (all keys are + * always all buffers or all strings) so we don't need to handle this. + */ + constructor( + entries?: ReadonlyArray< + [bigint | string | Uint8Array | null, SqlValue] + > | null, + ) { + if (entries) { + // Process entries to encode Buffer keys as base64 + const processedEntries = entries.map(([key, value]) => { + return [_parseBufferToMapKey(key), [key, value]]; + }) as ReadonlyArray<[bigint | string, _MapValue]>; + this.map_impl = new Map( + processedEntries, + ); + } else { + this.map_impl = new Map(); + } + } + + clear(): void { + this.map_impl.clear(); + } + + delete(key: string | bigint | Uint8Array | null): boolean { + return this.map_impl.delete(_parseBufferToMapKey(key)); + } + + forEach( + callbackfn: ( + value: SqlValue, + key: string | bigint | Uint8Array | null, + map: Map, + ) => void, + thisArg?: any, + ): void { + this.map_impl.forEach((value, key) => { + callbackfn.call(thisArg, value[1], value[0], this.map_impl); + }); + } + + has(key: string | bigint | Uint8Array | null): boolean { + return this.map_impl.has(_parseBufferToMapKey(key)); + } + + get size(): number { + return this.map_impl.size; + } + + entries(): IterableIterator<[string | bigint | Uint8Array | null, SqlValue]> { + return this.map_impl.values(); + } + + keys(): IterableIterator { + const iterator = this.map_impl.values(); + return { + next: () => { + const result = iterator.next(); + if (result.done) return result; + return {value: result.value[0], done: false}; + }, + [Symbol.iterator]() { + return this; + }, + }; + } + + values(): IterableIterator { + const iterator = this.map_impl.values(); + return { + next: () => { + const result = iterator.next(); + if (result.done) return result; + return {value: result.value[1], done: false}; + }, + [Symbol.iterator]() { + return this; + }, + }; + } + + [Symbol.iterator](): IterableIterator< + [string | bigint | Uint8Array | null, SqlValue] + > { + return this.entries(); + } + + get [Symbol.toStringTag](): string { + return 'EncodedKeyMap'; + } + + get(key: string | bigint | Uint8Array | null): SqlValue | undefined { + return this.map_impl.get(_parseBufferToMapKey(key))?.[1]; + } + + set(key: string | bigint | Uint8Array | null, value: SqlValue): this { + this.map_impl.set(_parseBufferToMapKey(key), [key, value]); + return this; + } +} diff --git a/src/index.ts b/src/index.ts index b17fcc7a3..a5ccdf63d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -24,6 +24,7 @@ import { } from 'google-gax'; import * as gax from 'google-gax'; import * as protos from '../protos/protos'; +import * as SqlTypes from './execute-query/types'; import {AppProfile} from './app-profile'; import {Cluster} from './cluster'; @@ -1137,6 +1138,7 @@ promisifyAll(Bigtable, { module.exports = Bigtable; module.exports.v2 = v2; module.exports.Bigtable = Bigtable; +module.exports.SqlTypes = SqlTypes; export {v2}; export {protos}; @@ -1362,3 +1364,4 @@ export { WaitForReplicationCallback, WaitForReplicationResponse, } from './table'; +export {SqlTypes}; diff --git a/src/instance.ts b/src/instance.ts index d1ca5890f..1acceaf6e 100644 --- a/src/instance.ts +++ b/src/instance.ts @@ -19,6 +19,7 @@ import * as is from 'is'; import * as extend from 'extend'; // eslint-disable-next-line @typescript-eslint/no-var-requires const pumpify = require('pumpify'); +const concat = require('concat-stream'); import snakeCase = require('lodash.snakecase'); import { @@ -70,6 +71,25 @@ import {Backup, RestoreTableCallback, RestoreTableResponse} from './backup'; import {ClusterUtils} from './utils/cluster'; import {AuthorizedView} from './authorized-view'; +import * as SqlTypes from './execute-query/types'; +import { + ExecuteQueryParameterValue, + QueryResultRow, +} from './execute-query/values'; +import {ProtobufReaderTransformer} from './execute-query/protobufreadertransformer'; +import {ExecuteQueryStreamTransformWithMetadata} from './execute-query/queryresultrowtransformer'; +import {ExecuteQueryStreamWithMetadata} from './execute-query/values'; +import { + parseParameters, + parseParameterTypes, +} from './execute-query/parameterparsing'; +import {MetadataConsumer} from './execute-query/metadataconsumer'; +import { + createCallerStream, + ExecuteQueryStateMachine, +} from './execute-query/executequerystatemachine'; +import {PreparedStatement} from './execute-query/preparedstatement'; + export interface ClusterInfo extends BasicClusterConfig { id: string; } @@ -154,6 +174,32 @@ export interface CreateTableFromBackupConfig { gaxOptions?: CallOptions; } +export type ExecuteQueryCallback = ( + err: Error | null, + rows?: QueryResultRow[], +) => void; + +export interface ExecuteQueryOptions { + preparedStatement: PreparedStatement; + parameters?: {[param: string]: ExecuteQueryParameterValue}; + retryOptions?: CallOptions; + encoding?: BufferEncoding; +} +export type ExecuteQueryResponse = [QueryResultRow[]]; + +export type PrepareStatementCallback = ( + err: Error | null, + preparedStatement?: PreparedStatement, +) => void; + +export interface PrepareStatementOptions { + query: string; + parameterTypes?: {[param: string]: SqlTypes.Type}; + retryOptions?: CallOptions; + encoding?: BufferEncoding; +} +export type PrepareStatementResponse = [PreparedStatement]; + /** * Create an Instance object to interact with a Cloud Bigtable instance. * @@ -1486,6 +1532,208 @@ Please use the format 'my-instance' or '${bigtable.projectName}/instances/my-ins view(tableName: string, viewName: string): AuthorizedView { return new AuthorizedView(this, tableName, viewName); } + + prepareStatement( + options: PrepareStatementOptions, + ): Promise; + prepareStatement( + options: PrepareStatementOptions, + callback: PrepareStatementCallback, + ): void; + prepareStatement(query: string): Promise; + prepareStatement(query: string, callback: PrepareStatementCallback): void; + /** + * Prepare an SQL query to be executed on an instance. + * + * @param {?string} [query] PreparedStatement object representing a query + * to execute. + * @param {string} [opts.query] Query string for which we want to construct the preparedStatement object. + * @param {object} [opts.parameterTypes] Object mapping names of parameters to their types. + * Type hints should be constructed using factory functions such as {@link Int64} + * @param {CallOptions} [opts.retryOptions] gax's CallOptions wich are passed straight to gax. + * The same retry options are also used when automatically refreshing the PreparedStatement. + * + * @param {function} callback The callback function. + * @param {?error} callback.err An error returned while making this request. + * @param {?PreparedStatement} callback.preparedStatement The preparedStatement object used to perform the executeQuery operation. + * + */ + prepareStatement( + queryOrOpts: string | PrepareStatementOptions, + callback?: PrepareStatementCallback, + ): void | Promise { + const opts: PrepareStatementOptions = + typeof queryOrOpts === 'string' ? {query: queryOrOpts} : queryOrOpts; + + const protoParamTypes = parseParameterTypes(opts.parameterTypes || {}); + const request = { + client: 'BigtableClient', + method: 'prepareQuery', + reqOpts: { + instanceName: this.name, + appProfileId: this.bigtable.appProfileId, + query: opts.query, + protoFormat: google.bigtable.v2.ProtoFormat.create(), + paramTypes: protoParamTypes, + }, + gaxOpts: opts.retryOptions, + }; + this.bigtable.request(request, (...args) => { + if (args[0]) { + callback!(args[0]); + } + try { + callback!( + null, + new PreparedStatement( + this.bigtable, + args[1]!, + request, + opts.parameterTypes || {}, + ), + ); + } catch (err) { + callback!(err as any, undefined); + } + }); + } + + executeQuery(options: ExecuteQueryOptions): Promise; + executeQuery( + options: ExecuteQueryOptions, + callback: ExecuteQueryCallback, + ): void; + executeQuery( + preparedStatement: PreparedStatement, + ): Promise; + executeQuery( + preparedStatement: PreparedStatement, + callback: ExecuteQueryCallback, + ): void; + /** + * Execute a SQL query on an instance. + * + * + * @param {?PreparedStatement} [preparedStatement] PreparedStatement object representing a query + * to execute. + * @param {?object} [options] Configuration object. See + * {@link Instance#createExecuteQueryStream} for a complete list of options. + * + * @param {function} callback The callback function. + * @param {?error} callback.err An error returned while making this request. + * @param {?QueryResultRow[]} callback.rows List of rows. + * @param {?SqlTypes.ResultSetMetadata} callback.metadata Metadata for the response. + * + * @example include:samples/api-reference-doc-snippets/instance.js + * region_tag:bigtable_api_execute_query + */ + executeQuery( + preparedStatementOrOpts: PreparedStatement | ExecuteQueryOptions, + callback?: ExecuteQueryCallback, + ): void | Promise { + let opts: ExecuteQueryOptions; + if (preparedStatementOrOpts instanceof PreparedStatement) { + opts = {preparedStatement: preparedStatementOrOpts}; + } else { + opts = preparedStatementOrOpts; + } + const stream = this.createExecuteQueryStream(opts); + + stream.on('error', callback!).pipe( + concat((rows: QueryResultRow[]) => { + callback!(null, rows); + }), + ); + } + + /** + * Execute a SQL query on an instance. + * + * @param {PreparedStatement} [preparedStatement] SQL query to execute. Parameters can be specified using @name notation. + * @param {object} [opts] Configuration object. + * @param {object} [opts.parameters] Object mapping names of parameters used in the query to JS values. + * @param {object} [opts.retryOptions] Retry options used for executing the query. Note that the only values + * used are: + * - retryOptions.retry.retryCodes + * - retryOptions.retry.backoffSettings.maxRetries + * - retryOptions.retry.backoffSettings.totalTimeoutMillis + * - retryOptions.retry.backoffSettings.maxRetryDelayMillis + * - retryOptions.retry.backoffSettings.retryDelayMultiplier + * - retryOptions.retry.backoffSettings.initialRetryDelayMillis + * @returns {ExecuteQueryStreamWithMetadata} + * + * @example include:samples/api-reference-doc-snippets/instance.js + * region_tag:bigtable_api_create_query_stream + */ + createExecuteQueryStream( + opts: ExecuteQueryOptions, + ): ExecuteQueryStreamWithMetadata { + /** + * We create the following streams: + * responseStream -> byteBuffer -> readerStream -> resultStream + * + * The last two (readerStream and resultStream) are connected using pumpify + * and returned to the caller. + * + * When a request is made responseStream and byteBuffer are created, + * connected using pumpify and piped to the readerStream. + * + * On retry, the old responseStream-byteBuffer pair is discarded and a + * new pair is crated. + * + * For more info please refer to comments in setupRetries function. + * + */ + const metadataConsumer = new MetadataConsumer(); + + let callerCancelled = false; + const setCallerCancelled = (value: boolean) => { + callerCancelled = value; + }; + const hasCallerCancelled = () => callerCancelled; + + const resultStream = new ExecuteQueryStreamTransformWithMetadata( + metadataConsumer, + hasCallerCancelled, + opts.encoding, + ); + const protoParams: {[k: string]: google.bigtable.v2.IValue} | null = + parseParameters( + opts.parameters || {}, + opts.preparedStatement.getParameterTypes(), + ); + + const readerStream = new ProtobufReaderTransformer(metadataConsumer); + + const reqOpts: google.bigtable.v2.IExecuteQueryRequest = { + instanceName: this.name, + appProfileId: this.bigtable.appProfileId, + protoFormat: google.bigtable.v2.ProtoFormat.create(), + params: protoParams, + }; + + // This creates a row stream which is two streams connected in a series. + const callerStream = createCallerStream( + readerStream, + resultStream, + metadataConsumer, + setCallerCancelled, + ); + + const stateMachine = new ExecuteQueryStateMachine( + this.bigtable, + callerStream, + opts.preparedStatement, + reqOpts, + opts.retryOptions?.retry, + opts.encoding, + ); + + // make sure stateMachine is not garbage collected as long as the callerStream. + callerStream._stateMachine = stateMachine; + + return callerStream; + } } /*! Developer Documentation diff --git a/src/tabular-api-surface.ts b/src/tabular-api-surface.ts index 462658e92..72cfe01ef 100644 --- a/src/tabular-api-surface.ts +++ b/src/tabular-api-surface.ts @@ -691,3 +691,15 @@ export class PartialFailureError extends Error { } } } +// Retry on "received rst stream" errors +export function isRstStreamError(error: ServiceError): boolean { + if (error.code === 13 && error.message) { + const error_message = (error.message || '').toLowerCase(); + return ( + error.code === 13 && + (error_message.includes('rst_stream') || + error_message.includes('rst stream')) + ); + } + return false; +} diff --git a/src/utils/retry-options.ts b/src/utils/retry-options.ts new file mode 100644 index 000000000..d50c82fe1 --- /dev/null +++ b/src/utils/retry-options.ts @@ -0,0 +1,77 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import {BackoffSettings} from 'google-gax/build/src/gax'; +import {GoogleError, grpc, ServiceError} from 'google-gax'; + +export const RETRYABLE_STATUS_CODES = new Set([ + grpc.status.DEADLINE_EXCEEDED.valueOf(), + grpc.status.UNAVAILABLE.valueOf(), +]); +export const DEFAULT_RETRY_COUNT = 10; +export const IGNORED_STATUS_CODES = new Set([grpc.status.CANCELLED.valueOf()]); + +interface Violation { + type: string; + description: string; +} + +interface ViolationsList { + violations?: Violation[]; +} + +function containsPreparedQueryExpired(violations: ViolationsList[]): boolean { + if (!Array.isArray(violations) || violations.length === 0) { + return false; + } + + for (const obj of violations) { + if (obj.violations && Array.isArray(obj.violations)) { + for (const violation of obj.violations) { + if (violation.type === 'PREPARED_QUERY_EXPIRED') { + return true; + } + } + } + } + + return false; +} + +/** + * Checks if the error is an "expired query plan" error. + * For more info refer to the ExecuteQueryStateMachine + * @param error + */ +export const isExpiredQueryError = ( + error: GoogleError | ServiceError, +): boolean => { + if ( + error.code === grpc.status.FAILED_PRECONDITION && + Object.hasOwn(error, 'statusDetails') + ) { + const statusDetails = (error as GoogleError) + .statusDetails as ViolationsList[]; + return containsPreparedQueryExpired(statusDetails); + } + return false; +}; + +/** + * Checks if the error is a cancel error - caused by aborting the stream. + * @param error + */ +export function isCancelError(error: ServiceError) { + return error.code === grpc.status.CANCELLED.valueOf(); +} diff --git a/system-test/bigtable.ts b/system-test/bigtable.ts index 52c1369f2..ebcd3c02d 100644 --- a/system-test/bigtable.ts +++ b/system-test/bigtable.ts @@ -25,6 +25,7 @@ import { Instance, InstanceOptions, MutateOptions, + SqlTypes, } from '../src'; import {Mutation} from '../src/mutation'; import {AppProfile} from '../src/app-profile.js'; @@ -37,6 +38,7 @@ import {RawFilter} from '../src/filter'; import {generateId, PREFIX} from './common'; import {BigtableTableAdminClient} from '../src/v2'; import {ServiceError} from 'google-gax'; +import {BigtableDate, QueryResultRow} from '../src/execute-query/values'; describe('Bigtable', () => { const bigtable = new Bigtable(); @@ -788,6 +790,100 @@ describe('Bigtable', () => { }); describe('fetching data', () => { + it('should execute a query', async () => { + const [preparedStatement] = await INSTANCE.prepareStatement({ + query: + 'SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, @doubleParam AS doubleCol,\n' + + '@floatParam AS floatCol, @boolParam AS boolCol, @tsParam AS tsCol, @dateParam AS dateCol,\n' + + '@byteArrayParam AS byteArrayCol, @stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol,\n' + + '@floatArrayParam AS floatArrayCol, @doubleArrayParam AS doubleArrayCol, @boolArrayParam AS boolArrayCol,\n' + + '@tsArrayParam AS tsArrayCol, @dateArrayParam AS dateArrayCol', + parameterTypes: { + bytesParam: SqlTypes.Bytes(), + intArrayParam: SqlTypes.Array(SqlTypes.Int64()), + dateArrayParam: SqlTypes.Array(SqlTypes.Date()), + stringParam: SqlTypes.String(), + byteArrayParam: SqlTypes.Array(SqlTypes.Bytes()), + doubleArrayParam: SqlTypes.Array(SqlTypes.Float64()), + boolArrayParam: SqlTypes.Array(SqlTypes.Bool()), + doubleParam: SqlTypes.Float64(), + floatParam: SqlTypes.Float32(), + dateParam: SqlTypes.Date(), + floatArrayParam: SqlTypes.Array(SqlTypes.Float32()), + tsArrayParam: SqlTypes.Array(SqlTypes.Timestamp()), + int64Param: SqlTypes.Int64(), + boolParam: SqlTypes.Bool(), + tsParam: SqlTypes.Timestamp(), + stringArrayParam: SqlTypes.Array(SqlTypes.String()), + }, + }); + const params = { + bytesParam: Buffer.from('test'), + intArrayParam: [BigInt(1), BigInt(2), BigInt(3)], + dateArrayParam: [ + new BigtableDate(2025, 5, 14), + new BigtableDate(2025, 5, 13), + ], + stringParam: 'test', + byteArrayParam: [Buffer.from('test')], + doubleArrayParam: [1.0, 2.0, 3.0], + boolArrayParam: [true, false], + doubleParam: 1.0, + floatParam: 1.0, + dateParam: new BigtableDate(2025, 5, 14), + floatArrayParam: [1.0, 2.0, 3.0], + tsArrayParam: [ + new PreciseDate(Date.now()), + new PreciseDate(Date.now()), + ], + int64Param: BigInt(123), + boolParam: true, + tsParam: new PreciseDate(Date.now()), + stringArrayParam: ['test', 'test'], + }; + const [rows] = (await INSTANCE.executeQuery({ + preparedStatement, + parameters: params, + })) as any as [Row[]]; + assert(rows[0] instanceof QueryResultRow); + assert.deepStrictEqual(rows[0].get('strCol'), params.stringParam); + assert.deepStrictEqual(rows[0].get('bytesCol'), params.bytesParam); + assert.deepStrictEqual(rows[0].get('intCol'), params.int64Param); + assert.deepStrictEqual(rows[0].get('doubleCol'), params.doubleParam); + assert.deepStrictEqual(rows[0].get('floatCol'), params.floatParam); + assert.deepStrictEqual(rows[0].get('boolCol'), params.boolParam); + assert.deepStrictEqual(rows[0].get('tsCol'), params.tsParam); + assert.deepStrictEqual(rows[0].get('dateCol'), params.dateParam); + assert.deepStrictEqual( + rows[0].get('byteArrayCol'), + params.byteArrayParam, + ); + assert.deepStrictEqual( + rows[0].get('stringArrayCol'), + params.stringArrayParam, + ); + assert.deepStrictEqual( + rows[0].get('intArrayCol'), + params.intArrayParam, + ); + assert.deepStrictEqual( + rows[0].get('floatArrayCol'), + params.floatArrayParam, + ); + assert.deepStrictEqual( + rows[0].get('doubleArrayCol'), + params.doubleArrayParam, + ); + assert.deepStrictEqual( + rows[0].get('boolArrayCol'), + params.boolArrayParam, + ); + assert.deepStrictEqual(rows[0].get('tsArrayCol'), params.tsArrayParam); + assert.deepStrictEqual( + rows[0].get('dateArrayCol'), + params.dateArrayParam, + ); + }); it('should get rows', async () => { const [rows] = await TABLE.getRows(); assert.strictEqual(rows.length, 4); diff --git a/test/base64keymap.ts b/test/base64keymap.ts new file mode 100644 index 000000000..48cb55662 --- /dev/null +++ b/test/base64keymap.ts @@ -0,0 +1,184 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import * as assert from 'assert'; +import {describe, it, before, beforeEach, afterEach} from 'mocha'; +import * as Long from 'long'; +import * as proxyquire from 'proxyquire'; +import * as sn from 'sinon'; + +import {RowStateEnum} from '../src/chunktransformer.js'; +import {Mutation} from '../src/mutation.js'; +import {Row} from '../src/row.js'; +import {EncodedKeyMap, SqlValue} from '../src/execute-query/values.js'; + +describe('Bigtable/EncodedKeyMap', () => { + describe('map tests', () => { + it('test constructor', () => { + const bufferKey = Buffer.from('exampleKey'); + const entries: [string | Buffer, string][] = [ + [bufferKey, 'valueForBufferKey'], + ['stringKey', 'valueForStringKey'], + ]; + + const map = new EncodedKeyMap(entries); + // get works with the same object + assert.deepStrictEqual(map.get(bufferKey), 'valueForBufferKey'); + // get works with a new object + assert.deepStrictEqual( + map.get(Buffer.from('exampleKey')), + 'valueForBufferKey', + ); + // get works with a regular string + assert.deepStrictEqual(map.get('stringKey'), 'valueForStringKey'); + }); + it('test duplicate keys', () => { + const bufferKey1 = Buffer.from('exampleKey'); + const bufferKey2 = Buffer.from('exampleKey'); + const bufferKey3 = Buffer.from('exampleKey'); + const entries: [string | Buffer, string][] = [ + [bufferKey1, 'valueForBufferKey1'], + ['stringKey', 'valueForStringKey1'], + [bufferKey2, 'valueForBufferKey2'], + ['stringKey', 'valueForStringKey2'], + ]; + + const map = new EncodedKeyMap(entries); + // get works with the same object + assert.deepStrictEqual(map.get(bufferKey1), 'valueForBufferKey2'); + assert.deepStrictEqual(map.get(bufferKey2), 'valueForBufferKey2'); + // get works with a new object + assert.deepStrictEqual( + map.get(Buffer.from('exampleKey')), + 'valueForBufferKey2', + ); + // get works with a regular string + assert.deepStrictEqual(map.get('stringKey'), 'valueForStringKey2'); + + // check that old value is replaced + map.set(bufferKey3, 'valueForBufferKey3'); + assert.deepStrictEqual( + map.get(Buffer.from('exampleKey')), + 'valueForBufferKey3', + ); + map.set('stringKey', 'valueForStringKey3'); + assert.deepStrictEqual(map.get('stringKey'), 'valueForStringKey3'); + }); + it('test get/set', () => { + const bufferKey = Buffer.from('exampleKey'); + const map = new EncodedKeyMap(); + map.set(bufferKey, 'valueForBufferKey'); + map.set('stringKey', 'valueForStringKey'); + // get works with the same object + assert.deepStrictEqual(map.get(bufferKey), 'valueForBufferKey'); + // get works with a new object + assert.deepStrictEqual( + map.get(Buffer.from('exampleKey')), + 'valueForBufferKey', + ); + // get works with a regular string + assert.deepStrictEqual(map.get('stringKey'), 'valueForStringKey'); + }); + it('test null vs empty bytes', () => { + const entries: [string | Buffer | null, string][] = [ + [null, 'valueForNull'], + ['', 'valueForEmptyString'], + ]; + + // TS normally would not permit a null key, thus we pass entries as any + const map = new EncodedKeyMap(entries as any); + // get works with the same object + assert.deepStrictEqual(map.get(''), 'valueForEmptyString'); + // get works with a regular string + assert.deepStrictEqual(map.get(null as any), 'valueForNull'); + }); + it('test null vs empty bytes', () => { + const entries: [string | Buffer | null, string][] = [ + [null, 'valueForNull'], + [Buffer.from(''), 'valueForEmptyBuffer'], + ]; + + // TS normally would not permit a null key, thus we pass entries as any + const map = new EncodedKeyMap(entries as any); + // get works with the same object + assert.deepStrictEqual(map.get(Buffer.from('')), 'valueForEmptyBuffer'); + // get works with a regular string + assert.deepStrictEqual(map.get(null as any), 'valueForNull'); + }); + it('map builtin functions', () => { + const entries: [string | Buffer | null, string][] = [ + [Buffer.from('Buffer1'), 'valueForBuffer1'], + ['stringKey1', 'valueForString1'], + ]; + + // TS normally would not permit a null key, thus we pass entries as any + const map = new EncodedKeyMap(entries as any); + + // get works with a buffer + assert.deepStrictEqual( + map.get(Buffer.from('Buffer1')), + 'valueForBuffer1', + ); + // get works with a regular string + assert.deepStrictEqual(map.get('stringKey1'), 'valueForString1'); + + // delete, set, has, size + + map.set(Buffer.from('Buffer2'), 'valueForBuffer2'); + map.set('stringKey2', 'valueForString2'); + + assert.deepStrictEqual(map.size, 4); + + assert.deepStrictEqual( + map.get(Buffer.from('Buffer2')), + 'valueForBuffer2', + ); + assert.deepStrictEqual(map.get('stringKey2'), 'valueForString2'); + + assert.strictEqual(map.has('stringKey2'), true); + assert.strictEqual(map.has(Buffer.from('Buffer2')), true); + + map.delete('stringKey2'); + map.delete(Buffer.from('Buffer2')); + + assert.strictEqual(map.has('stringKey2'), false); + assert.strictEqual(map.has(Buffer.from('Buffer2')), false); + + assert.deepStrictEqual(map.size, 2); + + // iterators + + const keys = [...map.keys()]; + assert.deepStrictEqual(keys[0]?.toString(), 'Buffer1'); + assert.deepStrictEqual(keys[0] instanceof Buffer, true); + assert.deepStrictEqual(keys[1], 'stringKey1'); + + const values = [...map.values()]; + assert.deepStrictEqual(values[0], 'valueForBuffer1'); + assert.deepStrictEqual(values[1], 'valueForString1'); + + const resultForEach: [string | bigint | Uint8Array | null, SqlValue][] = + []; + map.forEach((value, key) => { + resultForEach.push([key, value]); + }); + + assert.deepStrictEqual(resultForEach[0][0]?.toString(), 'Buffer1'); + assert.deepStrictEqual(resultForEach[0][0] instanceof Buffer, true); + assert.deepStrictEqual(resultForEach[0][1], 'valueForBuffer1'); + assert.deepStrictEqual(resultForEach[1][0], 'stringKey1'); + assert.deepStrictEqual(resultForEach[1][1], 'valueForString1'); + }); + }); +}); diff --git a/test/bytebuffertransformer.ts b/test/bytebuffertransformer.ts new file mode 100644 index 000000000..c5f541ad7 --- /dev/null +++ b/test/bytebuffertransformer.ts @@ -0,0 +1,379 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +import * as assert from 'assert'; +import {describe, it} from 'mocha'; +import * as sinon from 'sinon'; +import {google} from '../protos/protos'; +import {createProtoRows} from './utils/proto-bytes'; +import {ByteBufferTransformer} from '../src/execute-query/bytebuffertransformer'; +import * as SqlValues from '../src/execute-query/values'; + +type PublicByteBufferTransformer = { + messageQueue: Buffer[]; + messageBuffer: Uint8Array[]; + push: (data: any) => void; + processProtoRowsBatch: ( + partialResultSet: google.bigtable.v2.IPartialResultSet, + ) => void; +}; + +describe('Bigtable/ExecuteQueryByteBufferTransformer', () => { + let checksumValidStub: any; + let checksumIsValid = true; + let byteBuffer: PublicByteBufferTransformer; + + beforeEach(() => { + checksumIsValid = true; + checksumValidStub = sinon + .stub(SqlValues, 'checksumValid') + .callsFake(() => checksumIsValid); + byteBuffer = + new ByteBufferTransformer() as any as PublicByteBufferTransformer; + }); + + afterEach(() => { + checksumValidStub.restore(); + }); + + describe('processProtoRowsBatch', () => { + it('empty result', done => { + assert.throws(() => { + byteBuffer.processProtoRowsBatch({}); + }, /Error: Response did not contain any results!/); + done(); + }); + + it('just checksum', done => { + const response1 = createProtoRows(undefined, undefined, undefined, { + intValue: 1, + }); + const responseWithChecksum = createProtoRows(undefined, 111, undefined); + + // fill the buffer + byteBuffer.processProtoRowsBatch(response1.results!); + + // check that the buffer is filled + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.strictEqual( + byteBuffer.messageBuffer[0], + response1.results!.protoRowsBatch!.batchData!, + ); + + // send the checksum + byteBuffer.processProtoRowsBatch(responseWithChecksum.results!); + + // check that the buffer is flushed and queue contains the new message + assert.strictEqual(byteBuffer.messageQueue.length, 1); + assert.deepStrictEqual( + byteBuffer.messageQueue[0], + Buffer.concat([ + response1.results!.protoRowsBatch!.batchData! as Buffer, + ]), + ); + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + done(); + }); + + it('checksum flushes the buffer', done => { + const response1 = createProtoRows(undefined, undefined, undefined, { + intValue: 1, + }); + const responseWithChecksum = createProtoRows(undefined, 111, undefined, { + intValue: 2, + }); + + // fill the buffer + byteBuffer.processProtoRowsBatch(response1.results!); + + // check that the buffer is filled + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.strictEqual( + byteBuffer.messageBuffer[0], + response1.results!.protoRowsBatch!.batchData!, + ); + + // send a reset + byteBuffer.processProtoRowsBatch(responseWithChecksum.results!); + + // check that the buffer is flushed and queue contains the new message + // containing both values + assert.strictEqual(byteBuffer.messageQueue.length, 1); + assert.deepStrictEqual( + byteBuffer.messageQueue[0], + Buffer.concat([ + response1.results!.protoRowsBatch!.batchData! as Buffer, + responseWithChecksum.results!.protoRowsBatch!.batchData! as Buffer, + ]), + ); + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + done(); + }); + + it('just reset', done => { + const responseWithReset = createProtoRows(undefined, undefined, true); + + // send a reset + byteBuffer.processProtoRowsBatch(responseWithReset.results!); + + done(); + }); + + it('reset empties the buffer', done => { + // we first prepare the byteBuffer with a few messages + // then we send a reset and observe that the queue and + // buffer have been emptied and only the new message + // is present + const response1 = createProtoRows(undefined, undefined, undefined, { + intValue: 1, + }); + const responseWithReset = createProtoRows(undefined, undefined, true, { + intValue: 4, + }); + + byteBuffer.processProtoRowsBatch(response1.results!); + + // check that the buffer is filled + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.strictEqual( + byteBuffer.messageBuffer[0], + response1.results!.protoRowsBatch!.batchData!, + ); + + // send a reset + byteBuffer.processProtoRowsBatch(responseWithReset.results!); + + // check that the buffer has been emptied and populated with + // the new message after the reset + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.deepStrictEqual( + byteBuffer.messageBuffer[0], + responseWithReset.results!.protoRowsBatch!.batchData!, + ); + done(); + }); + + it('reset empties the queue and buffer', done => { + // we first prepare the byteBuffer with a few messages + // then we send a reset and observe that the queue and + // buffer have been emptied and only the new message + // is present + const responses = [ + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + createProtoRows(undefined, 111, undefined, {intValue: 2}), + createProtoRows(undefined, undefined, undefined, {intValue: 3}), + ]; + const responseWithReset = createProtoRows(undefined, undefined, true, { + intValue: 4, + }); + + // fill the buffer with messages + for (const response of responses) { + byteBuffer.processProtoRowsBatch(response.results!); + } + + // check that the buffer and queue are filled + assert.strictEqual(byteBuffer.messageQueue.length, 1); + assert.deepStrictEqual( + byteBuffer.messageQueue[0], + Buffer.concat([ + responses[0].results!.protoRowsBatch!.batchData! as Buffer, + responses[1].results!.protoRowsBatch!.batchData! as Buffer, + ]), + ); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.strictEqual( + byteBuffer.messageBuffer[0], + responses[2].results!.protoRowsBatch!.batchData!, + ); + + // send a reset + byteBuffer.processProtoRowsBatch(responseWithReset.results!); + + // check that the buffer and queue have been emptied and populated with + // the new message after the reset + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.deepStrictEqual( + byteBuffer.messageBuffer[0], + responseWithReset.results!.protoRowsBatch!.batchData!, + ); + done(); + }); + + it('token triggers push', done => { + let pushedData = null; + byteBuffer.push = (data: any) => { + pushedData = data; + }; + const response1 = createProtoRows(undefined, undefined, undefined, { + intValue: 1, + }); + const responseWithToken = createProtoRows('token', 111, undefined, { + intValue: 2, + }); + + // fill the buffer + byteBuffer.processProtoRowsBatch(response1.results!); + + // check that the buffer is filled + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.strictEqual( + byteBuffer.messageBuffer[0], + response1.results!.protoRowsBatch!.batchData!, + ); + + // send a token + byteBuffer.processProtoRowsBatch(responseWithToken.results!); + + // check that the data was pushed and buffer and queue are empty + // but pushed data contins the value from the 2nd message + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.deepStrictEqual(pushedData, [ + [ + Buffer.concat([ + response1.results!.protoRowsBatch!.batchData! as Buffer, + responseWithToken.results!.protoRowsBatch!.batchData! as Buffer, + ]), + ], + Buffer.from('token'), + ]); + done(); + }); + + it('separate token', done => { + let pushedData = null; + byteBuffer.push = (data: any) => { + pushedData = data; + }; + const response1 = createProtoRows(undefined, 111, undefined, { + intValue: 1, + }); + const responseWithToken = createProtoRows('token', undefined, undefined); + + // fill the buffer + byteBuffer.processProtoRowsBatch(response1.results!); + + // check that the buffer is filled + assert.strictEqual(byteBuffer.messageQueue.length, 1); + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + + // send a token + byteBuffer.processProtoRowsBatch(responseWithToken.results!); + + // check that the data was pushed and buffer and queue are empty + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.deepStrictEqual(pushedData, [ + [response1.results!.protoRowsBatch!.batchData! as Buffer], + Buffer.from('token'), + ]); + done(); + }); + + it('checksum without data throws', done => { + const responseWithChecksum = createProtoRows(undefined, 111, undefined); + + // send a checksum + assert.throws(() => { + byteBuffer.processProtoRowsBatch(responseWithChecksum.results!); + }, /Error: Recieved empty batch with non-zero checksum\./); + + done(); + }); + + it('token without checksum throws', done => { + let pushedData = null; + byteBuffer.push = (data: any) => { + pushedData = data; + }; + const response1 = createProtoRows(undefined, undefined, undefined, { + intValue: 1, + }); + const responseWithToken = createProtoRows('token', undefined, undefined); + + // fill the buffer + byteBuffer.processProtoRowsBatch(response1.results!); + + // check that the buffer is filled + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 1); + assert.strictEqual( + byteBuffer.messageBuffer[0], + response1.results!.protoRowsBatch!.batchData!, + ); + + // send a token + assert.throws(() => { + byteBuffer.processProtoRowsBatch(responseWithToken.results!); + }, /Error: Recieved incomplete batch of rows\./); + + done(); + }); + + it('token without data', done => { + let pushedData = null; + byteBuffer.push = (data: any) => { + pushedData = data; + }; + const responseWithToken = createProtoRows('token', undefined, undefined); + + // check that the buffer and queue are empty + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + + // send a token + byteBuffer.processProtoRowsBatch(responseWithToken.results!); + + // check that the token was pushed even though the buffer and queue are empty + assert.strictEqual(byteBuffer.messageBuffer.length, 0); + assert.strictEqual(byteBuffer.messageQueue.length, 0); + assert.deepStrictEqual(pushedData, [[], Buffer.from('token')]); + done(); + }); + + it('cheksum properly calculated', done => { + checksumValidStub.restore(); + const response = createProtoRows( + 'token1', + 2412835642, + undefined, + {intValue: 1}, + {intValue: 2}, + ); + byteBuffer.processProtoRowsBatch(response.results!); + done(); + }); + + it('invalid cheksum throws', done => { + checksumValidStub.restore(); + const response = createProtoRows( + 'token1', + 111, + undefined, + {intValue: 1}, + {intValue: 2}, + ); + assert.throws(() => { + byteBuffer.processProtoRowsBatch(response.results!); + }, /Error: Failed to validate next batch of results/); + done(); + }); + }); +}); diff --git a/test/executequery.ts b/test/executequery.ts new file mode 100644 index 000000000..2e87ac8f3 --- /dev/null +++ b/test/executequery.ts @@ -0,0 +1,2049 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +import * as promisify from '@google-cloud/promisify'; +import * as assert from 'assert'; +import {before, beforeEach, afterEach, describe, it} from 'mocha'; +import * as sinon from 'sinon'; +import * as proxyquire from 'proxyquire'; +import {grpc} from 'google-gax'; +import * as inst from '../src/instance'; +import {Bigtable} from '../src'; +import {google} from '../protos/protos'; +import * as SqlTypes from '../src/execute-query/types'; +import * as pumpify from 'pumpify'; +import { + ArrayReadableStream, + createMetadata, + createPrepareQueryResponse, + createProtoRows, + pbType, +} from './utils/proto-bytes'; +import {QueryResultRow} from '../src/execute-query/values'; +import { + PreparedStatement, + SHOULD_REFRESH_SOON_PERIOD_MS, +} from '../src/execute-query/preparedstatement'; +import {MetadataConsumer} from '../src/execute-query/metadataconsumer'; +import {PassThrough} from 'stream'; +import * as SqlValues from '../src/execute-query/values'; + +const sandbox = sinon.createSandbox(); + +const fakePromisify = Object.assign({}, promisify, { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + promisifyAll(klass: Function, options: any) { + if (klass.name !== 'Instance') { + return; + } + assert.deepStrictEqual(options.exclude, [ + 'appProfile', + 'cluster', + 'table', + 'getBackupsStream', + 'getTablesStream', + 'getAppProfilesStream', + 'view', + ]); + }, +}); + +class MockPreparedStatement { + callbacks: any[] = []; + markedAsExpired = false; + getData = (cb: any, timeout: any) => { + this.callbacks.push(cb); + }; + getParameterTypes = () => { + return {}; + }; + markAsExpired = () => { + this.markedAsExpired = true; + }; +} + +function createResultSetMetadata( + ...values: [string | null, google.bigtable.v2.Type][] +): SqlTypes.ResultSetMetadata { + return MetadataConsumer.parseMetadata(createMetadata(...values).metadata!); +} + +const performCallbacks = (callbacks: any[], interval: number) => { + let counter = 0; + const performNext = () => { + callbacks[counter++](); + if (counter < callbacks.length) { + setTimeout(performNext, interval); + } + }; + performNext(); +}; + +const createExpiredQueryError = () => { + return { + code: grpc.status.FAILED_PRECONDITION, + details: 'failed precondition', + statusDetails: [ + { + violations: [ + { + type: 'PREPARED_QUERY_EXPIRED', + description: + 'The prepared query has expired. Please re-issue the ExecuteQuery with a valid prepared query.', + }, + ], + }, + ], + }; +}; + +describe('Bigtable/ExecuteQueryStateMachine', () => { + const INSTANCE_ID = 'my-instance'; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const BIGTABLE = { + projectName: 'projects/my-project', + projectId: 'my-project', + request: () => {}, + } as Bigtable; + let Instance: typeof inst.Instance; + let instance: inst.Instance; + let checksumValidStub: any; + + before(() => { + Instance = proxyquire('../src/instance.js', { + '@google-cloud/promisify': fakePromisify, + pumpify, + }).Instance; + }); + + beforeEach(() => { + instance = new Instance(BIGTABLE, INSTANCE_ID); + checksumValidStub = sinon + .stub(SqlValues, 'checksumValid') + .callsFake(() => true); + }); + + afterEach(() => { + sandbox.restore(); + checksumValidStub.restore(); + }); + + describe('happy_path', () => { + it('responses within timeout', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, 111, undefined, {intValue: 2}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', undefined, undefined), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token2', undefined, undefined), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.emit('end'); + bigtableStream.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(responses.length, 2); + assert.equal(responses[0].get(0), 1); + assert.equal(responses[1].get(0), 2); + done(); + }, + ], + 1, + ); + }); + }); + + describe('queryPlanErrors', () => { + it('one query plan error', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0](new Error('fetching QP failed')); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 2); + preparedStatement.callbacks[1]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', 111, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.emit('end'); + bigtableStream.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(responses[0].get(0), 1); + done(); + }, + ], + 1, + ); + }); + + it('query plan expired error', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const expiredError = createExpiredQueryError(); + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', expiredError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainAndRefreshQueryPlan', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + assert.equal(preparedStatement.markedAsExpired, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 2); + preparedStatement.callbacks[1]( + undefined, + 'bytes', + createResultSetMetadata(['f2', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows('token1', 111, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream2.emit('end'); + bigtableStream2.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.throws(() => { + // we make sure that the column name from the first preparedStatement is not present. + responses[0].get('f1'); + }); + assert.equal(responses[0].get('f2'), 1); + done(); + }, + ], + 1, + ); + }); + + it('query plan expired error after data was recieved', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const expiredError = createExpiredQueryError(); + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', expiredError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainAndRefreshQueryPlan', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 2); + preparedStatement.callbacks[1]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows('token1', 111, undefined, {intValue: 2}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream2.emit('end'); + bigtableStream2.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + // assert we only got the second response, the first one was discarded + assert.equal(responses.length, 1); + assert.equal(responses[0].get(0), 2); + done(); + }, + ], + 1, + ); + }); + + it('query plan expired error after token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const expiredError = createExpiredQueryError(); + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + const rowsEmitted = 0; + const responses: QueryResultRow[] = []; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', (row: any) => { + responses.push(row); + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', 111, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.write(expiredError); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(responses.length, 1); + assert.equal(responses[0].get(0), 1); + assert.equal(errorEmitted, true); + done(); + }, + ], + 1, + ); + }); + }); + + describe('streamEnding', () => { + it('empty stream', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + let streamEnded = false; + resultStream.on('end', () => { + streamEnded = true; + }); + resultStream.on('error', () => { + errorEmitted = true; + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.emit('end'); + bigtableStream.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(errorEmitted, false); + assert.equal(responses.length, 0); + assert.equal(streamEnded, true); + done(); + }, + ], + 1, + ); + }); + + it('unexpected end after some data before token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + let rowsEmitted = 0; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', () => { + rowsEmitted += 1; + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.emit('end'); + bigtableStream.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(rowsEmitted, 0); + done(); + }, + ], + 1, + ); + }); + + it('unexpected end before a token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + resultStream.on('error', () => { + errorEmitted = true; + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', 111, undefined, {intValue: 2}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 3}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.emit('end'); + bigtableStream.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(responses.length, 2); + assert.equal(responses[0].get(0), 1); + assert.equal(responses[1].get(0), 2); + done(); + }, + ], + 1, + ); + }); + + it('empty response - query returned no rows', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + let rowsEmitted = 0; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', () => { + rowsEmitted += 1; + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', undefined, undefined), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream.emit('end'); + bigtableStream.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(errorEmitted, false); + assert.equal(rowsEmitted, 0); + done(); + }, + ], + 1, + ); + }); + }); + + describe('streamErrors', () => { + it('retryable error before anything', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const retryableError = { + code: grpc.status.DEADLINE_EXCEEDED, + message: 'retryable error', + }; + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', retryableError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainingBeforeResumeToken', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(preparedStatement.callbacks.length, 1); // query plan was not refreshed + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows('token', 111, undefined, {intValue: 2}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream2.emit('end'); + bigtableStream2.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(responses.length, 2); + assert.equal(responses[0].get(0), 1); + assert.equal(responses[1].get(0), 2); + done(); + }, + ], + 1, + ); + }); + + it('retryable error before token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const retryableError = { + code: grpc.status.DEADLINE_EXCEEDED, + message: 'retryable error', + }; + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', retryableError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainingBeforeResumeToken', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(preparedStatement.callbacks.length, 1); // query plan was not refreshed + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows(undefined, undefined, undefined, {intValue: 2}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows('token', 111, undefined, {intValue: 3}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream2.emit('end'); + bigtableStream2.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(responses.length, 2); + // the first message before the retry should have been discarded + assert.equal(responses[0].get(0), 2); + assert.equal(responses[1].get(0), 3); + done(); + }, + ], + 1, + ); + }); + + it('retryable error before token, byteBuffer keeps emitting data', done => { + // in this test we simulate a situation where even though the + // error was emitted, a data event emitted after it. This can + // happen if an event is buffered in the readable part of the byteBuffer + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const retryableError = { + code: grpc.status.DEADLINE_EXCEEDED, + message: 'retryable error', + }; + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + let valuesStream: any = null; + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + valuesStream = resultStream._stateMachine.valuesStream; + bigtableStream.emit('error', retryableError); + }, + () => { + // emit data after the error was emitted + valuesStream.emit('data', [ + [ + createProtoRows(undefined, undefined, undefined, {intValue: 2}) + .results?.protoRowsBatch?.batchData, + ], + 'unreachableToken', + ]); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainingBeforeResumeToken', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(preparedStatement.callbacks.length, 1); // query plan was not refreshed + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows(undefined, undefined, undefined, {intValue: 3}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream2.write( + createProtoRows('token', 111, undefined, {intValue: 4}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream2.emit('end'); + bigtableStream2.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(responses.length, 2); + // the first message before the retry should have been discarded + assert.equal(responses[0].get(0), 3); + assert.equal(responses[1].get(0), 4); + done(); + }, + ], + 1, + ); + }); + + it('retryable error before token then expire', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream2.abort = () => {}; + + const bigtableStream3 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream3.abort = () => {}; + + const retryableError = { + code: grpc.status.DEADLINE_EXCEEDED, + message: 'retryable error', + }; + + const expiredError = createExpiredQueryError(); + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', retryableError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainingBeforeResumeToken', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(preparedStatement.callbacks.length, 1); // query plan was not refreshed + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream3 as any; + }; + bigtableStream2.emit('error', expiredError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainAndRefreshQueryPlan', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 2); + preparedStatement.callbacks[1]( + undefined, + 'bytes', + createResultSetMetadata(['f2', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream3.write( + createProtoRows('token1', 111, undefined, {intValue: 2}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + bigtableStream3.emit('end'); + bigtableStream3.emit('close'); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Finished'); + assert.equal(responses.length, 1); + // the first message before the retry should have been discarded + assert.equal(responses[0].get('f2'), 2); + assert.throws(() => { + // we make sure that the column name from the first preparedStatement is not present. + responses[0].get('f1'); + }); + done(); + }, + ], + 1, + ); + }); + + it('retryable error after token then expire', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream2.abort = () => {}; + + const bigtableStream3 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream3.abort = () => {}; + + const retryableError = { + code: grpc.status.DEADLINE_EXCEEDED, + message: 'retryable error', + }; + + const expiredError = createExpiredQueryError(); + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + resultStream.on('error', () => { + errorEmitted = true; + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', 111, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', retryableError); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'DrainingAfterResumeToken', + ); + assert.equal(resultStream._stateMachine.retryTimer !== null, true); + // speed up the retry timer + clearTimeout(resultStream._stateMachine.retryTimer); + resultStream._stateMachine.startNextAttempt(); + }, + () => { + assert.equal(preparedStatement.callbacks.length, 1); // query plan was not refreshed + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream3 as any; + }; + bigtableStream2.emit('error', expiredError); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(responses.length, 1); + assert.equal(responses[0].get(0), 1); + done(); + }, + ], + 1, + ); + }); + + it('non-retryable error before token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const nonretryableError = { + code: grpc.status.PERMISSION_DENIED, + message: 'non-retryable error', + }; + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + let rowsEmitted = 0; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', () => { + rowsEmitted += 1; + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', nonretryableError); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(rowsEmitted, 0); + done(); + }, + ], + 1, + ); + }); + + it('non-retryable error after token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + + const bigtableStream2 = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream2.abort = () => {}; + + const nonretryableError = { + code: grpc.status.PERMISSION_DENIED, + message: 'non-retryable error', + }; + + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + resultStream.on('error', () => { + errorEmitted = true; + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', 111, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + BIGTABLE.request = () => { + return bigtableStream2 as any; + }; + bigtableStream.emit('error', nonretryableError); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(responses.length, 1); + assert.equal(responses[0].get(0), 1); + done(); + }, + ], + 1, + ); + }); + }); + + describe('timeouts', () => { + it('timeout immediately', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + let rowsEmitted = 0; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', () => { + rowsEmitted += 1; + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + }, + () => { + resultStream._stateMachine.handleTotalTimeout(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(rowsEmitted, 0); + done(); + }, + ], + 1, + ); + }); + + it('timeout after PQ', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + let rowsEmitted = 0; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', () => { + rowsEmitted += 1; + }); + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0](new Error('fetching QP failed!')); + }, + () => { + assert.equal(errorEmitted, false); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 2); + preparedStatement.callbacks[1]( + new Error('fetching QP failed again!'), + ); + }, + () => { + assert.equal(errorEmitted, false); + resultStream._stateMachine.handleTotalTimeout(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(rowsEmitted, 0); + done(); + }, + ], + 1, + ); + }); + + it('timeout before token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + let rowsEmitted = 0; + resultStream + .on('error', () => { + errorEmitted = true; + }) + .on('data', () => { + rowsEmitted += 1; + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + resultStream._stateMachine.handleTotalTimeout(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(rowsEmitted, 0); + done(); + }, + ], + 1, + ); + }); + + it('timeout after token', done => { + const bigtableStream = new PassThrough({ + objectMode: true, + highWaterMark: 0, + }) as any; + bigtableStream.abort = () => {}; + BIGTABLE.request = () => bigtableStream as any; + const preparedStatement = new MockPreparedStatement(); + const resultStream = instance.createExecuteQueryStream({ + preparedStatement, + } as any) as any; + let errorEmitted = false; + const responses: QueryResultRow[] = []; + resultStream.on('data', (row: any) => { + responses.push(row); + }); + resultStream.on('error', () => { + errorEmitted = true; + }); + + performCallbacks( + [ + () => { + clearTimeout(resultStream._stateMachine.timeoutTimer); + assert.equal(resultStream._stateMachine.state, 'AwaitingQueryPlan'); + assert.equal(preparedStatement.callbacks.length, 1); + }, + () => { + preparedStatement.callbacks[0]( + undefined, + 'bytes', + createResultSetMetadata(['f1', pbType({int64Type: {}})]), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'BeforeFirstResumeToken', + ); + bigtableStream.write( + createProtoRows('token1', 111, undefined, {intValue: 1}), + ); + }, + () => { + assert.equal( + resultStream._stateMachine.state, + 'AfterFirstResumeToken', + ); + resultStream._stateMachine.handleTotalTimeout(); + }, + () => { + assert.equal(resultStream._stateMachine.state, 'Failed'); + assert.equal(errorEmitted, true); + assert.equal(responses.length, 1); + assert.equal(responses[0].get(0), 1); + done(); + }, + ], + 1, + ); + }); + }); +}); + +describe('Bigtable/ExecuteQueryPreparedStatementObject', () => { + const INSTANCE_ID = 'my-instance'; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const BIGTABLE = { + projectName: 'projects/my-project', + projectId: 'my-project', + request: () => {}, + } as Bigtable; + + let clock: sinon.SinonFakeTimers; + + beforeEach(() => { + clock = sinon.useFakeTimers({ + toFake: ['setTimeout', 'clearTimeout', 'Date'], + }); + }); + + afterEach(() => { + clock.restore(); + sandbox.restore(); + }); + + describe('happy_path', () => { + it('getting prepared query plan', done => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Int64()}, + ); + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f').type, 'int64'); + done(); + }, 1000); + clock.runAll(); + }); + + it('getting prepared query plan close to validUntil', done => { + const resp = createPrepareQueryResponse(['f', pbType({int64Type: {}})]); + let pqRequestCb = null; + let requestCounter = 0; + (BIGTABLE as any).request = (req: any, cb: any) => { + requestCounter += 1; + pqRequestCb = cb; + }; + const someTimestamp = 1740000000; + resp.validUntil = google.protobuf.Timestamp.create({ + seconds: someTimestamp / 1000, + nanos: 0, + }); + const preparedStatement = new PreparedStatement( + BIGTABLE, + resp, + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + // Set the time to 100 ms after the "should-refresh" point in time + clock.setSystemTime(someTimestamp - SHOULD_REFRESH_SOON_PERIOD_MS + 100); + let getDataCalls = 0; + const doneAfterGetData = () => { + getDataCalls += 1; + if (getDataCalls > 1) { + // assert only one request was made. + assert.equal(requestCounter, 1); + done(); + } + }; + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f').type, 'int64'); + doneAfterGetData(); + }, 1000); + clock.runAll(); + + // refresh is scheduled + assert.equal(pqRequestCb !== null, true); + + // both getData calls should get the old value before the refresh finishes + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f').type, 'int64'); + doneAfterGetData(); + }, 1000); + clock.runAll(); + }); + + it('getting prepared query plan past validUntil', done => { + const resp = createPrepareQueryResponse(['f', pbType({int64Type: {}})]); + const someTimestamp = 1740000000; + resp.validUntil = google.protobuf.Timestamp.create({ + seconds: someTimestamp / 1000, + nanos: 0, + }); + const preparedStatement = new PreparedStatement( + BIGTABLE, + resp, + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + // Set the time to 100 ms after the "validUntil" point in time + clock.setSystemTime(someTimestamp + 100); + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f').type, 'int64'); + done(); + }, 1000); + clock.runAll(); + }); + + it('multiple getData calls result in only one request', done => { + const originalResp = createPrepareQueryResponse([ + 'f1', + pbType({int64Type: {}}), + ]); + const secondResp = createPrepareQueryResponse([ + 'f2', + pbType({int64Type: {}}), + ]); + let pqRequestCb = null; + let requestCounter = 0; + (BIGTABLE as any).request = (req: any, cb: any) => { + requestCounter += 1; + pqRequestCb = cb; + }; + const someTimestamp = 1740000000; + originalResp.validUntil = google.protobuf.Timestamp.create({ + seconds: someTimestamp / 1000, + nanos: 0, + }); + const preparedStatement = new PreparedStatement( + BIGTABLE, + originalResp, + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + // Set the time to 100 ms after the "should-refresh" point in time + clock.setSystemTime(someTimestamp - SHOULD_REFRESH_SOON_PERIOD_MS + 100); + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f1').type, 'int64'); + }, 1000); + clock.runAll(); + + // refresh is scheduled + assert.equal(pqRequestCb !== null, true); + + // second getData call + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f1').type, 'int64'); + }, 1000); + clock.runAll(); + + // assert only one request was made even though getData was called twice + assert.equal(requestCounter, 1); + + // Bigtable returns the prepareQuery response + pqRequestCb!(null, secondResp); + + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f2').type, 'int64'); + assert.throws(() => { + // we make sure that the column name from the first preparedStatement is not present. + metadata?.get('f1'); + }); + }, 1000); + clock.runAll(); + done(); + }); + }); + + describe('other_cases', () => { + it('getting data after expiration hangs', done => { + const resp = createPrepareQueryResponse(['f', pbType({int64Type: {}})]); + const preparedStatement = new PreparedStatement( + BIGTABLE, + resp, + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + (BIGTABLE as any).request = (req: any, cb: any) => cb(null, resp); + + preparedStatement.markAsExpired(); + assert.equal(preparedStatement.isExpired(), true); + assert.equal((preparedStatement as any).isRefreshing, false); + assert.equal((preparedStatement as any).timer, null); + + let callbackCalled = false; + preparedStatement.getData((err, pqBytes, metadata) => { + callbackCalled = true; + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f').type, 'int64'); + }, 1000); + + // getData scheduled getting the query plan immediately after + assert.equal((preparedStatement as any).timer !== null, true); + assert.equal(callbackCalled, false); + clock.tick(1); + assert.equal((preparedStatement as any).timer, null); + assert.equal(callbackCalled, true); + done(); + }); + + it('plan expired during getData callback', done => { + const resp = createPrepareQueryResponse(['f', pbType({int64Type: {}})]); + const preparedStatement = new PreparedStatement( + BIGTABLE, + resp, + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + (BIGTABLE as any).request = (req: any, cb: any) => cb(null, resp); + + preparedStatement.markAsExpired(); + assert.equal(preparedStatement.isExpired(), true); + assert.equal((preparedStatement as any).isRefreshing, false); + assert.equal((preparedStatement as any).timer, null); + + let callbackCalled = false; + preparedStatement.getData((err, pqBytes, metadata) => { + callbackCalled = true; + assert.equal(err, undefined); + assert.equal(pqBytes, 'xd'); + assert.equal(metadata?.get('f').type, 'int64'); + preparedStatement.markAsExpired(); + }, 1000); + + // this callback gets served second. It will get an error + // because the query got expired between the last refresh and serving of this callback + preparedStatement.getData((err, pqBytes, metadata) => { + assert.equal(callbackCalled, true); + assert.equal(pqBytes, undefined); + assert.equal(metadata, undefined); + assert.equal(err?.message, 'Getting a fresh query plan failed.'); + }, 1000); + + // getData scheduled getting the query plan immediately after + assert.equal((preparedStatement as any).timer !== null, true); + assert.equal(callbackCalled, false); + clock.tick(1); + assert.equal((preparedStatement as any).timer, null); + assert.equal(callbackCalled, true); + done(); + }); + + it('plan refresh failed', done => { + const resp = createPrepareQueryResponse(['f', pbType({int64Type: {}})]); + const preparedStatement = new PreparedStatement( + BIGTABLE, + resp, + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + (BIGTABLE as any).request = (req: any, cb: any) => + cb(new Error('Problem')); + + preparedStatement.markAsExpired(); + assert.equal(preparedStatement.isExpired(), true); + assert.equal((preparedStatement as any).isRefreshing, false); + assert.equal((preparedStatement as any).timer, null); + + let callbackCalled = false; + preparedStatement.getData((err, pqBytes, metadata) => { + callbackCalled = true; + assert.equal(pqBytes, undefined); + assert.equal(metadata, undefined); + assert.equal(err?.message, 'Problem'); + }, 1000); + + assert.equal(callbackCalled, false); + clock.tick(1); + assert.equal((preparedStatement as any).timer, null); + assert.equal(callbackCalled, true); + done(); + }); + }); +}); diff --git a/test/instance.ts b/test/instance.ts index e26b7f9e1..ed3c6bcdd 100644 --- a/test/instance.ts +++ b/test/instance.ts @@ -19,7 +19,7 @@ import * as sinon from 'sinon'; import * as proxyquire from 'proxyquire'; import {ServiceError} from 'google-gax'; import * as snapshot from 'snap-shot-it'; - +import {Readable} from 'stream'; import * as inst from '../src/instance'; import {AppProfile, AppProfileOptions} from '../src/app-profile'; import {Cluster, CreateClusterOptions} from '../src/cluster'; @@ -34,9 +34,31 @@ import {Bigtable, RequestOptions} from '../src'; import {PassThrough} from 'stream'; import * as pumpify from 'pumpify'; import {FakeCluster} from '../system-test/common'; +import { + BigtableDate, + BigtableMap, + QueryResultRow, + SqlValue, + Struct, +} from '../src/execute-query/values'; +import * as SqlTypes from '../src/execute-query/types'; import {RestoreTableConfig} from '../src/backup'; import {Options} from './cluster'; import {createClusterOptionsList} from './constants/cluster'; +import {google} from '../protos/protos'; +import {PreciseDate} from '@google-cloud/precise-date'; +import Long = require('long'); +import { + createMetadata, + createPreparedStatement, + createPrepareQueryResponse, + createProtoRows, + pbType, +} from './utils/proto-bytes'; +import {PreparedStatement} from '../src/execute-query/preparedstatement'; +import * as SqlValues from '../src/execute-query/values'; + +const concat = require('concat-stream'); const sandbox = sinon.createSandbox(); @@ -87,6 +109,25 @@ class FakeTable extends Table { } } +// convenience function for ExecuteQuery tests +function executeQueryResultWithMetadata( + instance: any, + preparedStatement: PreparedStatement | null, + callback: (...args: any[]) => void, +): void { + const stream = instance.createExecuteQueryStream({preparedStatement}); + stream.on('error', callback!).pipe( + concat((rows: QueryResultRow[]) => { + const metadata = stream.getMetadata(); + if (metadata === null) { + callback!(new Error('Server error - did not receive metadata.')); + } else { + callback!(null, rows, metadata); + } + }), + ); +} + describe('Bigtable/Instance', () => { const INSTANCE_ID = 'my-instance'; // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -1165,7 +1206,6 @@ describe('Bigtable/Instance', () => { .getAppProfilesStream() .on('error', err => { assert.strictEqual(appProfiles.length, counter); - console.log(err.message); assert.deepStrictEqual( err, new Error( @@ -1966,3 +2006,1389 @@ describe('Bigtable/Instance', () => { }); }); }); + +describe('Bigtable/ExecuteQueryInstance', () => { + // Create an array of Response objects + + const responsesRef = { + responses: [] as google.bigtable.v2.ExecuteQueryResponse[], + + setResponses(values: google.bigtable.v2.ExecuteQueryResponse[]) { + responsesRef.responses = values; + }, + }; + + let requests: any[] = []; + + const INSTANCE_ID = 'my-instance'; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const BIGTABLE = { + projectName: 'projects/my-project', + projectId: 'my-project', + request: (config?: any) => { + requests.push(config); + const result: any = Readable.from(responsesRef.responses); + result.abort = () => {}; + return result; + }, + } as Bigtable; + let Instance: typeof inst.Instance; + let instance: inst.Instance; + let checksumValidStub: any; + let checksumIsValid = true; + + before(() => { + Instance = proxyquire('../src/instance.js', { + '@google-cloud/promisify': fakePromisify, + './app-profile.js': {AppProfile: FakeAppProfile}, + './backup.js': {Backup: FakeBackup}, + './cluster.js': {Cluster: FakeCluster}, + './family.js': {Family: FakeFamily}, + './table.js': {Table: FakeTable}, + pumpify, + }).Instance; + }); + + beforeEach(() => { + responsesRef.responses = []; + requests = []; + instance = new Instance(BIGTABLE, INSTANCE_ID); + checksumIsValid = true; + checksumValidStub = sinon + .stub(SqlValues, 'checksumValid') + .callsFake(() => checksumIsValid); + }); + + afterEach(() => { + sandbox.restore(); + checksumValidStub.restore(); + }); + + describe('execute', () => { + it('parses non-composite types', done => { + const preparedStatement = createPreparedStatement( + ['int64', pbType({int64Type: {}})], + ['float64', pbType({float64Type: {}})], + ['string', pbType({stringType: {}})], + ['bytes', pbType({bytesType: {}})], + ['date', pbType({dateType: {}})], + ['timestamp', pbType({timestampType: {}})], + ['bool', pbType({boolType: {}})], + ); + + responsesRef.setResponses([ + createProtoRows( + 'token1', + 111, + undefined, + {intValue: 1}, + {floatValue: 2.5}, + {stringValue: '3'}, + {bytesValue: new Uint8Array([4, 5, 6])}, + {dateValue: new google.type.Date({year: 2024, month: 0, day: 1})}, + { + timestampValue: new google.protobuf.Timestamp({ + seconds: 1234, + nanos: 5678, + }), + }, + {boolValue: true}, + ), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(metadata!.get(0).type, 'int64'); + assert.strictEqual(metadata!.get(1).type, 'float64'); + assert.strictEqual(metadata!.get(2).type, 'string'); + assert.strictEqual(metadata!.get(3).type, 'bytes'); + assert.strictEqual(metadata!.get(4).type, 'date'); + assert.strictEqual(metadata!.get(5).type, 'timestamp'); + assert.strictEqual(metadata!.get(6).type, 'bool'); + + assert.strictEqual(result![0].get(0), BigInt(1)); + assert.strictEqual(result![0].get(1), 2.5); + assert.strictEqual(result![0].get(2), '3'); + assert.deepEqual(result![0].get(3), new Uint8Array([4, 5, 6])); + assert.deepEqual(result![0].get(4), new BigtableDate(2024, 0, 1)); + assert.deepEqual(result![0].get(5), new PreciseDate([1234, 5678])); + assert.strictEqual(result![0].get(6), true); + done(); + }, + ); + }); + + it('parses multiple rows', done => { + const preparedStatement = createPreparedStatement([ + 'f1', + pbType({int64Type: {}}), + ]); + + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, {intValue: 1}), + createProtoRows('token2', 111, undefined, {intValue: 2}), + createProtoRows('token3', 111, undefined, {intValue: 3}), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(metadata!.get(0), metadata!.get('f1')); + assert.strictEqual(metadata!.get(0).type, 'int64'); + + assert.strictEqual(result![0].get(0), BigInt(1)); + assert.strictEqual(result![1].get(0), BigInt(2)); + assert.strictEqual(result![2].get(0), BigInt(3)); + done(); + }, + ); + }); + + it('handles nulls properly', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + + responsesRef.setResponses([ + createProtoRows(undefined, undefined, undefined, {intValue: 1}), + createProtoRows('token1', 111, undefined, {}), + createProtoRows(undefined, undefined, undefined, {}), + createProtoRows('token2', 111, undefined, {intValue: 2}), + createProtoRows(undefined, undefined, undefined, {}), + createProtoRows(undefined, undefined, undefined, {intValue: 3}), + createProtoRows('token3', 111, undefined), + createProtoRows(undefined, undefined, undefined, {}), + createProtoRows(undefined, undefined, undefined, {}), + createProtoRows('token4', 111, undefined), + createProtoRows(undefined, undefined, undefined, {intValue: 4}), + createProtoRows(undefined, undefined, undefined, {intValue: 5}), + createProtoRows('token5', 111, undefined), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(result?.length, 5); + + assert.strictEqual(result![0].get(0), BigInt(1)); + assert.strictEqual(result![0].get(1), null); + + assert.strictEqual(result![1].get(0), null); + assert.strictEqual(result![1].get(1), BigInt(2)); + + assert.strictEqual(result![2].get(0), null); + assert.strictEqual(result![2].get(1), BigInt(3)); + + assert.strictEqual(result![3].get(0), null); + assert.strictEqual(result![3].get(1), null); + + assert.strictEqual(result![4].get(0), BigInt(4)); + assert.strictEqual(result![4].get(1), BigInt(5)); + + done(); + }, + ); + }); + + it('handles nulls for all types', done => { + const preparedStatement = createPreparedStatement( + ['int64', pbType({int64Type: {}})], + ['float64', pbType({float64Type: {}})], + ['string', pbType({stringType: {}})], + ['bytes', pbType({bytesType: {}})], + ['date', pbType({dateType: {}})], + ['timestamp', pbType({timestampType: {}})], + ['bool', pbType({boolType: {}})], + ['array', pbType({arrayType: {elementType: pbType({int64Type: {}})}})], + [ + 'map', + pbType({ + mapType: { + keyType: pbType({int64Type: {}}), + valueType: pbType({int64Type: {}}), + }, + }), + ], + [ + 'struct', + pbType({ + structType: { + fields: [{fieldName: 'f1', type: pbType({int64Type: {}})}], + }, + }), + ], + [ + 'arrayWithNulls', + pbType({arrayType: {elementType: pbType({int64Type: {}})}}), + ], + [ + 'mapWithNulls', + pbType({ + mapType: { + keyType: pbType({int64Type: {}}), + valueType: pbType({stringType: {}}), + }, + }), + ], + [ + 'structWithNulls', + pbType({ + structType: { + fields: [ + {fieldName: 'f1', type: pbType({int64Type: {}})}, + {fieldName: null, type: pbType({float64Type: {}})}, + {fieldName: 'f3', type: pbType({stringType: {}})}, + ], + }, + }), + ], + ); + responsesRef.setResponses([ + createProtoRows( + 'token1', + 111, + undefined, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + // arrayWithNulls + { + arrayValue: { + values: [{intValue: 1}, {}, {intValue: 3}], + }, + }, + // mapWithNulls + { + arrayValue: { + values: [ + { + arrayValue: { + values: [{intValue: 1}, {}], + }, + }, + { + arrayValue: { + values: [{intValue: 2}, {}], + }, + }, + { + arrayValue: { + values: [{intValue: 3}, {stringValue: 'c'}], + }, + }, + ], + }, + }, + //structWithNulls + { + arrayValue: { + values: [{intValue: 1}, {}, {}], + }, + }, + ), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(result![0].get(0), null); + assert.strictEqual(result![0].get(1), null); + assert.strictEqual(result![0].get(2), null); + assert.strictEqual(result![0].get(3), null); + assert.strictEqual(result![0].get(4), null); + assert.strictEqual(result![0].get(5), null); + assert.strictEqual(result![0].get(6), null); + assert.strictEqual(result![0].get(7), null); + assert.strictEqual(result![0].get(8), null); + assert.strictEqual(result![0].get(9), null); + + const arrayWithNulls = result![0].get(10) as SqlValue[]; + assert.strictEqual(arrayWithNulls[0], BigInt(1)); + assert.strictEqual(arrayWithNulls[1], null); + assert.strictEqual(arrayWithNulls[2], BigInt(3)); + + const mapWithNulls = result![0].get(11) as BigtableMap; + assert.strictEqual(mapWithNulls.size, 3); + assert.strictEqual(mapWithNulls.get(BigInt(1)), null); + assert.strictEqual(mapWithNulls.get(BigInt(2)), null); + assert.strictEqual(mapWithNulls.get(BigInt(3)), 'c'); + + const structWithNulls = result![0].get(12) as Struct; + assert.strictEqual(structWithNulls.get('f1'), BigInt(1)); + assert.strictEqual(structWithNulls.get(1), null); + assert.strictEqual(structWithNulls.get('f3'), null); + + done(); + }, + ); + }); + + it('parses multiple rows in one batch', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows( + undefined, + undefined, + undefined, + {intValue: 1}, + {intValue: 2}, + ), + createProtoRows( + undefined, + undefined, + undefined, + {intValue: 3}, + {intValue: 4}, + ), + createProtoRows('token1', 111, undefined, {intValue: 5}, {intValue: 6}), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(metadata!.get(0), metadata!.get('f1')); + assert.strictEqual(metadata!.get(1), metadata!.get('f2')); + assert.strictEqual(metadata!.get(0).type, 'int64'); + assert.strictEqual(metadata!.get(1).type, 'int64'); + + assert.strictEqual(result![0].get(0), BigInt(1)); + assert.strictEqual(result![0].get('f1'), BigInt(1)); + assert.strictEqual(result![0].get(1), BigInt(2)); + assert.strictEqual(result![0].get('f2'), BigInt(2)); + + assert.strictEqual(result![1].get(0), BigInt(3)); + assert.strictEqual(result![1].get('f1'), BigInt(3)); + assert.strictEqual(result![1].get(1), BigInt(4)); + assert.strictEqual(result![1].get('f2'), BigInt(4)); + + assert.strictEqual(result![2].get(0), BigInt(5)); + assert.strictEqual(result![2].get('f1'), BigInt(5)); + assert.strictEqual(result![2].get(1), BigInt(6)); + assert.strictEqual(result![2].get('f2'), BigInt(6)); + done(); + }, + ); + }); + + it('parses an array of ints', done => { + const preparedStatement = createPreparedStatement([ + 'f1', + pbType({arrayType: {elementType: pbType({int64Type: {}})}}), + ]); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, { + arrayValue: { + values: [{intValue: 1}, {intValue: 2}, {intValue: 3}], + }, + }), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(metadata!.get(0), metadata!.get('f1')); + assert.strictEqual(metadata!.get(0).type, 'array'); + const arrayType = metadata!.get(0); + assert(arrayType.type === 'array'); + assert.strictEqual(arrayType.elementType.type, 'int64'); + + const structResult = result![0].get('f1') as SqlValue[]; + assert.strictEqual(structResult[0], BigInt(1)); + assert.strictEqual(structResult[1], BigInt(2)); + assert.strictEqual(structResult[2], BigInt(3)); + done(); + }, + ); + }); + + it('parses a struct', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + [ + 'f2', + pbType({ + structType: { + fields: [ + {fieldName: 'f1', type: pbType({int64Type: {}})}, + {fieldName: null, type: pbType({float64Type: {}})}, + {fieldName: 'f3', type: pbType({stringType: {}})}, + ], + }, + }), + ], + ); + responsesRef.setResponses([ + createProtoRows( + 'token1', + 111, + undefined, + {intValue: 1}, + { + arrayValue: { + values: [{intValue: 1}, {floatValue: 2.5}, {stringValue: '3'}], + }, + }, + ), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(metadata!.get(0), metadata!.get('f1')); + assert.strictEqual(metadata!.get(1), metadata!.get('f2')); + assert.strictEqual(metadata!.get(0).type, 'int64'); + const structType = metadata!.get(1); + assert.strictEqual(structType.type, 'struct'); + assert.strictEqual(structType.get('f1').type, 'int64'); + assert.strictEqual(structType.get(1).type, 'float64'); + assert.strictEqual(structType.get('f3').type, 'string'); + + assert.strictEqual(result![0].get(0), BigInt(1)); + assert.strictEqual(result![0].get('f1'), BigInt(1)); + const structResult = result![0].get(1) as Struct; + assert.strictEqual(structResult.get('f1'), structResult.get(0)); + assert.strictEqual(structResult.get('f3'), structResult.get(2)); + + assert.strictEqual(structResult.get(0), BigInt(1)); + assert.strictEqual(structResult.get(1), 2.5); + assert.strictEqual(structResult.get(2), '3'); + done(); + }, + ); + }); + + it('parses a map', done => { + const preparedStatement = createPreparedStatement([ + 'f1', + pbType({ + mapType: { + keyType: pbType({int64Type: {}}), + valueType: pbType({stringType: {}}), + }, + }), + ]); + responsesRef.setResponses([ + createProtoRows(undefined, undefined, undefined, { + arrayValue: { + values: [ + { + arrayValue: { + values: [{intValue: 1}, {stringValue: 'a'}], + }, + }, + { + arrayValue: { + values: [{intValue: 2}, {stringValue: 'b'}], + }, + }, + { + arrayValue: { + values: [{intValue: 3}, {stringValue: 'c'}], + }, + }, + ], + }, + }), + createProtoRows('token2', 111, undefined, { + arrayValue: { + values: [ + { + arrayValue: { + values: [{intValue: 4}, {stringValue: 'd'}], + }, + }, + { + arrayValue: { + values: [{intValue: 5}, {stringValue: 'e'}], + }, + }, + { + arrayValue: { + values: [{intValue: 6}, {stringValue: 'f'}], + }, + }, + ], + }, + }), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + const mapType = metadata!.get(0); + assert.strictEqual(mapType.type, 'map'); + assert.strictEqual(mapType.keyType.type, 'int64'); + assert.strictEqual(mapType.valueType.type, 'string'); + + assert.strictEqual(result?.length, 2); + + const mapResult0 = result![0].get('f1') as BigtableMap; + assert.strictEqual(mapResult0.size, 3); + assert.strictEqual(mapResult0.get(BigInt(1)), 'a'); + assert.strictEqual(mapResult0.get(BigInt(2)), 'b'); + assert.strictEqual(mapResult0.get(BigInt(3)), 'c'); + + const mapResult1 = result![1].get('f1') as BigtableMap; + assert.strictEqual(mapResult1.size, 3); + assert.strictEqual(mapResult1.get(BigInt(4)), 'd'); + assert.strictEqual(mapResult1.get(BigInt(5)), 'e'); + assert.strictEqual(mapResult1.get(BigInt(6)), 'f'); + done(); + }, + ); + }); + + it('map retains last encountered value for duplicate key', done => { + const preparedStatement = createPreparedStatement([ + 'f1', + pbType({ + mapType: { + keyType: pbType({int64Type: {}}), + valueType: pbType({stringType: {}}), + }, + }), + ]); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, { + arrayValue: { + values: [ + { + arrayValue: { + values: [{intValue: 1}, {stringValue: 'a'}], + }, + }, + { + arrayValue: { + values: [{intValue: 2}, {stringValue: 'b'}], + }, + }, + { + arrayValue: { + values: [{intValue: 1}, {stringValue: 'c'}], + }, + }, + ], + }, + }), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + const mapType = metadata!.get(0); + assert.strictEqual(mapType.type, 'map'); + assert.strictEqual(mapType.keyType.type, 'int64'); + assert.strictEqual(mapType.valueType.type, 'string'); + + assert.strictEqual(result?.length, 1); + + const mapResult0 = result![0].get('f1') as BigtableMap; + assert.strictEqual(mapResult0.size, 2); + assert.strictEqual(mapResult0.get(BigInt(1)), 'c'); + assert.strictEqual(mapResult0.get(BigInt(2)), 'b'); + done(); + }, + ); + }); + + it('accessing duplicated struct field throws', done => { + const preparedStatement = createPreparedStatement([ + 'structColumn', + pbType({ + structType: { + fields: [ + {fieldName: 'f1', type: pbType({int64Type: {}})}, + {fieldName: null, type: pbType({float64Type: {}})}, + {fieldName: 'f1', type: pbType({stringType: {}})}, + ], + }, + }), + ]); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, { + arrayValue: { + values: [{intValue: 1}, {floatValue: 2.5}, {stringValue: '3'}], + }, + }), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, result, metadata) => { + assert.strictEqual(metadata!.get(0).type, 'struct'); + + const struct = result![0].get(0) as Struct; + assert.strictEqual(struct.get(0), BigInt(1)); + assert.strictEqual(struct.get(1), 2.5); + assert.strictEqual(struct.get(2), '3'); + + assert.throws(() => { + result![0].get('f1'); + }, Error); + done(); + }, + ); + }); + + it('unsupported kind in metadata is detected', done => { + const type = {kind: 'unknown-type'}; + const BIGTABLE2 = { + projectName: 'projects/my-project2', + projectId: 'my-project2', + request: (req, cb: any) => { + cb!( + null, + createPrepareQueryResponse( + ['f1', pbType({int64Type: {}})], + ['f2', type as any], + ), + ); + }, + } as Bigtable; + const instance2 = new Instance(BIGTABLE2, INSTANCE_ID); + + instance2.prepareStatement('query', (err, result) => { + assert.notStrictEqual(err, null); + assert.ok(err instanceof Error); + done(); + }); + }); + + it('unsupported map key type throws', done => { + const BIGTABLE2 = { + projectName: 'projects/my-project2', + projectId: 'my-project2', + request: (req, cb: any) => { + cb!( + null, + createPrepareQueryResponse([ + 'map', + pbType({ + mapType: { + keyType: pbType({dateType: {}}), + valueType: pbType({int64Type: {}}), + }, + }), + ]), + ); + }, + } as Bigtable; + const instance2 = new Instance(BIGTABLE2, INSTANCE_ID); + instance2.prepareStatement('query', (err, result) => { + assert.notStrictEqual(err, null); + done(); + }); + }); + + it('map with null key is rejected', done => { + const preparedStatement = createPreparedStatement([ + 'map', + pbType({ + mapType: { + keyType: pbType({int64Type: {}}), + valueType: pbType({int64Type: {}}), + }, + }), + ]); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, { + arrayValue: { + values: [ + { + arrayValue: { + values: [{}, {intValue: 1}], + }, + }, + ], + }, + }), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.strictEqual(result?.length, 1); + done(); + }); + }); + + it('map with null value is ok', done => { + const preparedStatement = createPreparedStatement([ + 'map', + pbType({ + mapType: { + keyType: pbType({int64Type: {}}), + valueType: pbType({int64Type: {}}), + }, + }), + ]); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, { + arrayValue: { + values: [ + { + arrayValue: { + values: [{intValue: 1}, {}], + }, + }, + ], + }, + }), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.strictEqual(result?.length, 1); + done(); + }); + }); + + it('bigints are correctly converted to longs', done => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + Object.fromEntries( + Array.from({length: 11}, (_, i) => [ + String.fromCharCode(97 + i), + SqlTypes.Int64(), + ]), + ), // parameter types: {a:INT64, b:INT64, ... } + ); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, {intValue: 1}), + ]); + instance.executeQuery( + { + preparedStatement, + parameters: { + a: BigInt(1), + b: BigInt(-1), + c: BigInt(0), + d: BigInt(Number.MAX_SAFE_INTEGER), + e: BigInt(Number.MIN_SAFE_INTEGER), + f: BigInt('9007199254740992'), // MAX_SAFE_INTEGER + 1 + g: BigInt('-9007199254740992'), // MIN_SAFE_INTEGER - 1 + h: BigInt('1152921504606846976'), // 2^60 + i: BigInt('-1152921504606846976'), // - 2^60 + j: BigInt('9223372036854775807'), // 2^63 - 1 + k: BigInt('-9223372036854775808'), // - 2^63 + }, + } as any, + err => { + assert.equal(err, null); + assert.strictEqual(requests.length, 1); + const reqOpts = requests[0] + .reqOpts as google.bigtable.v2.IExecuteQueryRequest; + + assert.deepEqual(reqOpts.params!['a'].intValue, Long.fromInt(1)); + assert.deepEqual(reqOpts.params!['b'].intValue, Long.fromInt(-1)); + assert.deepEqual(reqOpts.params!['c'].intValue, Long.fromInt(0)); + assert.deepEqual( + reqOpts.params!['d'].intValue, + Long.fromNumber(Number.MAX_SAFE_INTEGER), + ); + assert.deepEqual( + reqOpts.params!['e'].intValue, + Long.fromNumber(Number.MIN_SAFE_INTEGER), + ); + assert.deepEqual( + reqOpts.params!['f'].intValue, + Long.fromString('9007199254740992'), + ); + assert.deepEqual( + reqOpts.params!['g'].intValue, + Long.fromString('-9007199254740992'), + ); + assert.deepEqual( + reqOpts.params!['h'].intValue, + Long.fromString('1152921504606846976'), + ); + assert.deepEqual( + reqOpts.params!['i'].intValue, + Long.fromString('-1152921504606846976'), + ); + assert.deepEqual( + reqOpts.params!['j'].intValue, + Long.fromString('9223372036854775807'), + ); + assert.deepEqual( + reqOpts.params!['k'].intValue, + Long.fromString('-9223372036854775808'), + ); + done(); + }, + ); + }); + + it('value not matching provided type is rejected', () => { + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Int64()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 'a'}, + } as any, + () => {}, + ); + }, + {message: 'Value a cannot be converted to int64.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Float64()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: BigInt(1)}, + } as any, + () => {}, + ); + }, + {message: 'Value 1 cannot be converted to float64.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.String()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 1}, + } as any, + () => {}, + ); + }, + {message: 'Value 1 cannot be converted to string.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Bytes()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 1}, + } as any, + () => {}, + ); + }, + {message: 'Value 1 cannot be converted to bytes.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Bool()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 1}, + } as any, + () => {}, + ); + }, + {message: 'Value 1 cannot be converted to boolean.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Timestamp()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 1}, + } as any, + () => {}, + ); + }, + { + message: + 'Value 1 cannot be converted to timestamp, please use PreciseDate instead.', + }, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Date()}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 1}, + } as any, + () => {}, + ); + }, + {message: 'Value 1 cannot be converted to date.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Array(SqlTypes.Int64())}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: 1}, + } as any, + () => {}, + ); + }, + {message: 'Value 1 cannot be converted to an array.'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Array(SqlTypes.Int64())}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: {a: [1, 'a']}, + } as any, + () => {}, + ); + }, + { + message: + 'Error while converting element 0 of an array: Value 1 cannot be converted to int64 - argument of type INT64 should by passed as BigInt.', + }, + ); + // TS does not permit passing a Struct or a Map as parameters, + // but we want to check it throws an error + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Map(SqlTypes.Int64(), SqlTypes.Int64())}, + ); + instance.executeQuery( + { + preparedStatement, + parameters: { + a: new Map([ + [BigInt(1), 2], + [BigInt(3), 'a'] as any as [bigint, number], + ]), + }, + } as any, + () => {}, + ); + }, + {message: 'Map is not a supported query param type'}, + ); + assert.throws( + () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + { + a: SqlTypes.Struct({ + name: 'f1', + type: SqlTypes.Int64(), + }), + }, + ); + instance.executeQuery( + { + preparedStatement, + parameters: { + a: SqlTypes.Struct({ + name: 'f1', + type: SqlTypes.Int64(), + }), + }, + } as any, + () => {}, + ); + }, + {message: 'Struct is not a supported query param type'}, + ); + }); + + it('null value is accepted', done => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + { + a: SqlTypes.Int64(), + }, + ); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, {intValue: 1}), + ]); + instance.executeQuery( + { + preparedStatement, + parameters: {a: null}, + } as any, + () => { + assert.strictEqual(requests.length, 1); + const reqOpts = requests[0] + .reqOpts as google.bigtable.v2.IExecuteQueryRequest; + + assert.notStrictEqual(reqOpts.params!['a'].type!.int64Type, null); + done(); + }, + ); + }); + + it('parameter type is used for null', done => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + { + a: SqlTypes.Int64(), + b: SqlTypes.Float64(), + c: SqlTypes.Bool(), + d: SqlTypes.Bytes(), + e: SqlTypes.String(), + f: SqlTypes.Date(), + g: SqlTypes.Timestamp(), + h: SqlTypes.Array(SqlTypes.Int64()), + }, + ); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, {intValue: 1}), + ]); + instance.executeQuery( + { + preparedStatement, + parameters: { + a: null, + b: null, + c: null, + d: null, + e: null, + f: null, + g: null, + h: null, + }, + } as any, + () => { + assert.strictEqual(requests.length, 1); + const reqOpts = requests[0] + .reqOpts as google.bigtable.v2.IExecuteQueryRequest; + + assert.notStrictEqual(reqOpts.params!['a'].type!.int64Type, null); + assert.notStrictEqual(reqOpts.params!['b'].type!.float64Type, null); + assert.notStrictEqual(reqOpts.params!['c'].type!.boolType, null); + assert.notStrictEqual(reqOpts.params!['d'].type!.bytesType, null); + assert.notStrictEqual(reqOpts.params!['e'].type!.stringType, null); + assert.notStrictEqual(reqOpts.params!['f'].type!.dateType, null); + assert.notStrictEqual(reqOpts.params!['g'].type!.timestampType, null); + assert.notStrictEqual(reqOpts.params!['h'].type!.arrayType, null); + assert.notStrictEqual( + reqOpts.params!['h'].type!.arrayType?.elementType, + null, + ); + done(); + }, + ); + }); + + it('large bigints are rejected', () => { + const preparedStatement = new PreparedStatement( + BIGTABLE, + createPrepareQueryResponse(['f', pbType({int64Type: {}})]), + {} as any, + {a: SqlTypes.Int64()}, + ); + assert.throws( + () => { + instance.executeQuery( + { + preparedStatement, + parameters: {a: BigInt('-9223372036854775809')}, + } as any, + () => {}, + ); + }, + { + message: + 'Value -9223372036854775809 cannot be converted to int64 - it is out of range.', + }, + ); + assert.throws( + () => { + instance.executeQuery( + { + preparedStatement, + parameters: {a: BigInt('9223372036854775808')}, + } as any, + () => {}, + ); + }, + { + message: + 'Value 9223372036854775808 cannot be converted to int64 - it is out of range.', + }, + ); + }); + + it('duplicate struct field names are not accessible by name', done => { + const preparedStatement = createPreparedStatement([ + 's', + pbType({ + structType: { + fields: [ + {fieldName: 'f1', type: pbType({int64Type: {}})}, + {fieldName: 'f2', type: pbType({int64Type: {}})}, + {fieldName: 'f1', type: pbType({stringType: {}})}, + ], + }, + }), + ]); + responsesRef.setResponses([ + createProtoRows('token1', 111, undefined, { + arrayValue: { + values: [{intValue: 1}, {intValue: 2}, {stringValue: '3'}], + }, + }), + ]); + instance.executeQuery(preparedStatement, (err, rows) => { + const struct = rows![0].get('s')! as Struct; + assert.strictEqual(struct.get(0), BigInt(1)); + assert.strictEqual(struct.get(1), BigInt(2)); + assert.strictEqual(struct.get(2), '3'); + + assert.throws(() => { + struct.get('f1'); + }, Error); + done(); + }); + }); + + it('duplicate row field names are not accessible by name', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ['f1', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows( + 'token1', + 111, + undefined, + {intValue: 1}, + {intValue: 2}, + {intValue: 3}, + ), + ]); + executeQueryResultWithMetadata( + instance, + preparedStatement, + (err, rows, metadata) => { + const row = rows![0]; + assert.strictEqual(row.get(0), BigInt(1)); + assert.strictEqual(row.get(1), BigInt(2)); + assert.strictEqual(row.get(2), BigInt(3)); + assert.strictEqual(row.get('f2'), BigInt(2)); + + assert.throws(() => { + row.get('f1'); + }, Error); + + assert.strictEqual(metadata!.get(0).type, 'int64'); + assert.strictEqual(metadata!.get(1).type, 'int64'); + assert.strictEqual(metadata!.get(2).type, 'int64'); + assert.strictEqual(metadata!.get('f2').type, 'int64'); + + assert.throws(() => { + metadata!.get('f1'); + }, Error); + + done(); + }, + ); + }); + + it('unfinished batch is detected', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows(undefined, undefined, undefined, {intValue: 3}), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.notStrictEqual(err, null); + assert.ok(err instanceof Error); + done(); + }); + }); + + it('token without batch ending detected', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows('token', undefined, undefined, {intValue: 3}), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.notStrictEqual(err, null); + assert.ok(err instanceof Error); + done(); + }); + }); + + it('reset works', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + + const respWithReset1 = createProtoRows(undefined, undefined, undefined, { + intValue: 1, + }); + respWithReset1.results!.reset = true; + + const respWithReset2 = createProtoRows( + undefined, + 111, + undefined, + {intValue: 3}, + {intValue: 4}, + ); + respWithReset2.results!.reset = true; + + responsesRef.setResponses([ + createProtoRows( + undefined, + undefined, + undefined, + {intValue: 1}, + {intValue: 2}, + ), + respWithReset1, + respWithReset2, + createProtoRows('token', 222, undefined, {intValue: 5}, {intValue: 6}), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.equal(err, null); + assert.strictEqual(result![0].get(0), BigInt(3)); + assert.strictEqual(result![0].get(1), BigInt(4)); + assert.strictEqual(result![1].get(0), BigInt(5)); + assert.strictEqual(result![1].get(1), BigInt(6)); + done(); + }); + }); + + it('partial row after token detected', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows( + 'token1', + 111, + undefined, + {intValue: 1}, + {intValue: 2}, + {intValue: 3}, + ), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.notStrictEqual(err, null); + assert.ok(err instanceof Error); + done(); + }); + }); + + it('partial row after batch checksum detected', done => { + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows( + undefined, + 111, + undefined, + {intValue: 1}, + {intValue: 2}, + {intValue: 3}, + ), + createProtoRows('token1', 222, undefined, {intValue: 4}), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.notStrictEqual(err, null); + assert.ok(err instanceof Error); + done(); + }); + }); + + it('cheksum fail detected', done => { + checksumIsValid = false; + const preparedStatement = createPreparedStatement( + ['f1', pbType({int64Type: {}})], + ['f2', pbType({int64Type: {}})], + ); + responsesRef.setResponses([ + createProtoRows( + undefined, + 111, + undefined, + {intValue: 1}, + {intValue: 2}, + ), + createProtoRows('token1', 222, undefined, {intValue: 3}, {intValue: 4}), + ]); + instance.executeQuery(preparedStatement, (err, result) => { + assert.notStrictEqual(err, null); + assert.ok(err instanceof Error); + done(); + }); + }); + }); +}); diff --git a/test/utils/proto-bytes.ts b/test/utils/proto-bytes.ts new file mode 100644 index 000000000..031ade177 --- /dev/null +++ b/test/utils/proto-bytes.ts @@ -0,0 +1,144 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +import {Readable} from 'stream'; +import {google} from '../../protos/protos'; +import {PreparedStatement} from '../../src/execute-query/preparedstatement'; + +export const createMetadata = ( + ...values: [string | null, google.bigtable.v2.Type][] +): google.bigtable.v2.ExecuteQueryResponse => { + return google.bigtable.v2.ExecuteQueryResponse.create({ + metadata: { + protoSchema: { + columns: values.map(v => + google.bigtable.v2.ColumnMetadata.create({ + name: v[0], + type: v[1], + }), + ), + }, + }, + }); +}; + +export const createPreparedStatement = ( + ...values: [string | null, google.bigtable.v2.Type][] +): PreparedStatement => { + const metadataPB = createMetadata(...values).metadata!; + const prepareQueryResponse = google.bigtable.v2.PrepareQueryResponse.create({ + metadata: metadataPB, + preparedQuery: 'xd', + validUntil: null, + }); + return new PreparedStatement( + undefined as any, + prepareQueryResponse, + {} as any, + {}, + ); +}; + +export const createPrepareQueryResponse = ( + ...values: [string | null, google.bigtable.v2.Type][] +): google.bigtable.v2.PrepareQueryResponse => { + const metadataPB = createMetadata(...values).metadata!; + return google.bigtable.v2.PrepareQueryResponse.create({ + metadata: metadataPB, + preparedQuery: 'xd', + validUntil: null, + }); +}; + +export const pbType = ( + value: google.bigtable.v2.IType, +): google.bigtable.v2.Type => { + return google.bigtable.v2.Type.create(value); +}; + +export const createProtoRows = ( + resumeToken?: string, + batchChecksum?: number, + reset?: boolean, + ...values: google.bigtable.v2.IValue[] +): google.bigtable.v2.ExecuteQueryResponse => { + const bytes = google.bigtable.v2.ProtoRows.encode( + google.bigtable.v2.ProtoRows.create({ + values: values.map(v => google.bigtable.v2.Value.create(v)), + }), + ).finish(); + + return { + response: 'results', + results: { + protoRowsBatch: values.length > 0 ? {batchData: bytes} : undefined, + resumeToken: resumeToken ? Buffer.from(resumeToken) : undefined, + batchChecksum: batchChecksum, + reset: reset || false, + }, + } as google.bigtable.v2.ExecuteQueryResponse; +}; + +interface BigtableError { + status: any; + message: any; + code: any; +} + +interface CustomCallback { + callback: any; +} + +type DataObject = + | BigtableError + | google.bigtable.v2.ExecuteQueryResponse + | CustomCallback; + +export class ArrayReadableStream extends Readable { + private data: DataObject[]; + private index: number; + private aborted: boolean; + + constructor(data: DataObject[]) { + super({objectMode: true, highWaterMark: 0}); + this.data = data; + this.index = 0; + this.aborted = false; + } + + _read() { + if (!this.aborted && this.index < this.data.length) { + const item = this.data[this.index]; + this.index++; + if ((item as BigtableError).code !== undefined) { + this.emit('error', item as BigtableError); + } else if ((item as CustomCallback).callback !== undefined) { + (item as CustomCallback).callback(); + this._read(); + } else { + this.push(item as google.bigtable.v2.ExecuteQueryResponse); + } + } else { + this.push(null); // No more data + } + } + + abort() { + this.aborted = true; + super.destroy(); + } + + end() { + super.destroy(); + } +} diff --git a/testproxy/known_failures.txt b/testproxy/known_failures.txt index 64c5f23af..f8210a570 100644 --- a/testproxy/known_failures.txt +++ b/testproxy/known_failures.txt @@ -15,6 +15,4 @@ TestSampleRowKeys_Retry_WithRoutingCookie\| TestSampleRowKeys_Generic_CloseClient\| TestSampleRowKeys_Generic_Headers\| TestSampleRowKeys_NoRetry_NoEmptyKey\| -TestSampleRowKeys_Retry_WithRetryInfo\| -TestExecuteQuery_EmptyResponse\| -TestExecuteQuery_SingleSimpleRow +TestSampleRowKeys_Retry_WithRetryInfo diff --git a/testproxy/services/execute-query.js b/testproxy/services/execute-query.js new file mode 100644 index 000000000..efb73ac2d --- /dev/null +++ b/testproxy/services/execute-query.js @@ -0,0 +1,64 @@ +// Copyright 2022 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +'use strict'; + +const grpc = require('@grpc/grpc-js'); +const { + parseMetadata, + parseRows, + parseParameters, +} = require('../../build/testproxy/services/utils/request/createExecuteQueryResponse.js'); +const normalizeCallback = require('./utils/normalize-callback.js'); + +const executeQuery = ({clientMap}) => + normalizeCallback(async rawRequest => { + const {request, clientId} = rawRequest.request; + + const {instanceName} = request; + const bigtable = clientMap.get(clientId); + const instance = bigtable.instance(instanceName); + + try { + const [parameters, parameterTypes] = await parseParameters( + request.params, + ); + const [preparedStatement] = await instance.prepareStatement({ + query: request.query, + parameterTypes: parameterTypes, + }); + const [rows] = await instance.executeQuery({ + preparedStatement, + parameters: parameters, + retryOptions: {}, + }); + + const parsedMetadata = await parseMetadata(preparedStatement); + const parsedRows = await parseRows(preparedStatement, rows); + return { + status: {code: grpc.status.OK, details: []}, + rows: parsedRows, + metadata: {columns: parsedMetadata}, + }; + } catch (e) { + return { + status: { + code: e.code || grpc.status.INTERNAL, + details: [], // e.details must be in an empty array for the test runner to return the status. This is tracked in https://b.corp.google.com/issues/383096533. + message: e.message, + }, + }; + } + }); + +module.exports = executeQuery; diff --git a/testproxy/services/index.js b/testproxy/services/index.js index 8e539fec3..09b169109 100644 --- a/testproxy/services/index.js +++ b/testproxy/services/index.js @@ -24,6 +24,7 @@ const readRow = require('./read-row.js'); const readRows = require('./read-rows.js'); const removeClient = require('./remove-client.js'); const sampleRowKeys = require('./sample-row-keys.js'); +const executeQuery = require('./execute-query.js'); /* * Starts the client pool map and retrieves the object that @@ -45,6 +46,7 @@ function getServicesImplementation() { readRows: readRows({clientMap}), removeClient: removeClient({clientMap}), sampleRowKeys: sampleRowKeys({clientMap}), + executeQuery: executeQuery({clientMap}), }; } diff --git a/testproxy/services/utils/request/createExecuteQueryResponse.ts b/testproxy/services/utils/request/createExecuteQueryResponse.ts new file mode 100644 index 000000000..da4f4eeb6 --- /dev/null +++ b/testproxy/services/utils/request/createExecuteQueryResponse.ts @@ -0,0 +1,223 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import * as protos from '../../../../protos/protos'; +import {SqlTypes} from '../../../../src'; +import {MetadataConsumer} from '../../../../src/execute-query/metadataconsumer'; +import { + convertJsValueToValue, + executeQueryTypeToPBType, +} from '../../../../src/execute-query/parameterparsing'; +import {PreparedStatement} from '../../../../src/execute-query/preparedstatement'; +import {ExecuteQueryStreamTransformWithMetadata} from '../../../../src/execute-query/queryresultrowtransformer'; +import { + QueryResultRow, + SqlValue, + EncodedKeyMap, + Struct, + BigtableMap, + ExecuteQueryParameterValue, +} from '../../../../src/execute-query/values'; +import * as is from 'is'; + +async function getMetadataFromPreparedStatement( + preparedStatement: PreparedStatement, +): Promise { + return await new Promise((resolve, reject) => { + preparedStatement.getData( + ( + err?: Error, + preparedQueryBytes?: Uint8Array | string, + metadata?: SqlTypes.ResultSetMetadata, + ) => { + if (err) { + reject(err); + } else { + resolve(metadata!); + } + }, + 100, + ); + }); +} + +export async function parseMetadata(preparedStatement: PreparedStatement) { + const metadata = await getMetadataFromPreparedStatement(preparedStatement); + const values = metadata.columns.map((v, i) => { + return [metadata.getFieldNameAtIndex(i), executeQueryTypeToPBType(v)]; + }); + return values.map(v => + protos.google.bigtable.v2.ColumnMetadata.create({ + name: v[0] as any, + type: v[1] as any, + }), + ); +} + +function convertToArray( + value: SqlValue, + type: SqlTypes.ArrayType, +): protos.google.bigtable.v2.IValue { + if (!is.array(value)) { + throw new Error(`Value ${value} cannot be converted to an array.`); + } + const arrayValue = value as Array; + return { + arrayValue: { + values: arrayValue.map((element, index) => { + try { + return convertAnyValueToPb(element, type.elementType); + // eslint-disable-next-line + } catch (conversionError: any) { + if (conversionError instanceof Error) { + throw new Error( + `Error while converting element ${index} of an array: ${conversionError.message}`, + ); + } else { + throw conversionError; + } + } + }), + }, + }; +} + +function convertToStruct( + value: SqlValue, + type: SqlTypes.StructType, +): protos.google.bigtable.v2.IValue { + if (!(typeof value === 'object' && value instanceof Struct)) { + throw new Error(`Value ${value} cannot be converted to an array.`); + } + const arrayValue = value as Struct; + return { + arrayValue: { + values: arrayValue.values.map((element, index) => { + try { + return convertAnyValueToPb(element, type.get(index)); + // eslint-disable-next-line + } catch (convertionError: any) { + if (convertionError instanceof Error) { + throw new Error( + `Error while converting element ${index} of a Struct to field of type ${ + type.get(index).type + }: ${convertionError.message}`, + ); + } else { + throw convertionError; + } + } + }), + }, + }; +} + +function convertToMap( + value: SqlValue, + type: SqlTypes.MapType, +): protos.google.bigtable.v2.IValue { + if (!(value instanceof EncodedKeyMap)) { + throw new Error(`Value ${value} cannot be converted to a map.`); + } + const arrayValue = value as BigtableMap; + return { + arrayValue: { + values: Array.from(arrayValue.entries()).flatMap(([key, value]) => { + return { + arrayValue: { + values: [ + convertMapEntry(key, type.keyType, key, 'key'), + convertMapEntry(value, type.valueType, key, 'value'), + ], + }, + }; + }), + }, + }; +} + +function convertMapEntry( + value: SqlValue, + type: SqlTypes.Type, + keyName: string | bigint | Uint8Array | null, + keyOrValue: 'key' | 'value', +): protos.google.bigtable.v2.IValue { + try { + return convertAnyValueToPb(value, type); + // eslint-disable-next-line + } catch (convertionError: any) { + if (convertionError instanceof Error) { + throw new Error( + `Error while converting element ${keyName} of a Map to map ${keyOrValue} of type ${type.type}: ${convertionError.message}`, + ); + } else { + throw convertionError; + } + } +} + +function convertAnyValueToPb( + value: SqlValue, + type: SqlTypes.Type, +): protos.google.bigtable.v2.IValue { + if (value === null || value === undefined) { + return protos.google.bigtable.v2.Value.create({}); + } + if (type.type === 'array') { + return convertToArray(value, type); + } else if (type.type === 'struct') { + return convertToStruct(value, type); + } else if (type.type === 'map') { + return convertToMap(value, type); + } else { + return convertJsValueToValue(value as ExecuteQueryParameterValue, type); + } +} + +export async function parseRows( + preparedStatement: PreparedStatement, + rows: QueryResultRow[], +) { + const metadata = await getMetadataFromPreparedStatement(preparedStatement); + const parsedRows = rows.map(row => { + const rowValues = metadata.columns.map((type, i) => { + const value = row.get(i); + return convertAnyValueToPb(value, type); + }); + return {values: rowValues}; + }); + return parsedRows; +} + +export async function parseParameters(params: { + [param: string]: protos.google.bigtable.v2.Value; +}) { + const parameters: {[param: string]: SqlValue} = {}; + const parameterTypes: {[param: string]: SqlTypes.Type} = {}; + const transfomer = new ExecuteQueryStreamTransformWithMetadata( + null as any, + () => false, + 'utf-8', + {}, + ); + for (const [paramName, pbValue] of Object.entries(params)) { + const type = MetadataConsumer.parsePBType( + pbValue.type as protos.google.bigtable.v2.Type, + ); + const value = transfomer.valueToJsType(pbValue, type); + parameters[paramName] = value; + parameterTypes[paramName] = type; + } + return [parameters, parameterTypes]; +} From 78d6bfaf108cabb0bd3ab143289b4f92335f12dc Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Fri, 6 Jun 2025 09:41:59 +0200 Subject: [PATCH 02/14] fix license-headers --- src/execute-query/bytebuffertransformer.ts | 2 +- src/execute-query/namedlist.ts | 2 +- src/execute-query/parameterparsing.ts | 2 +- src/execute-query/protobufreadertransformer.ts | 2 +- src/execute-query/queryresultrowtransformer.ts | 2 +- src/execute-query/types.ts | 2 +- src/execute-query/values.ts | 2 +- src/utils/retry-options.ts | 2 +- test/base64keymap.ts | 4 ++-- test/bytebuffertransformer.ts | 3 ++- testproxy/services/execute-query.js | 3 ++- 11 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/execute-query/bytebuffertransformer.ts b/src/execute-query/bytebuffertransformer.ts index 223328a20..c2222fa6e 100644 --- a/src/execute-query/bytebuffertransformer.ts +++ b/src/execute-query/bytebuffertransformer.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/execute-query/namedlist.ts b/src/execute-query/namedlist.ts index b745c537a..6dd5dfa3c 100644 --- a/src/execute-query/namedlist.ts +++ b/src/execute-query/namedlist.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/execute-query/parameterparsing.ts b/src/execute-query/parameterparsing.ts index d0fc35be2..658c0b50c 100644 --- a/src/execute-query/parameterparsing.ts +++ b/src/execute-query/parameterparsing.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/execute-query/protobufreadertransformer.ts b/src/execute-query/protobufreadertransformer.ts index 7349b095a..f98517d83 100644 --- a/src/execute-query/protobufreadertransformer.ts +++ b/src/execute-query/protobufreadertransformer.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/execute-query/queryresultrowtransformer.ts b/src/execute-query/queryresultrowtransformer.ts index 3cd277ccb..f2b3cdb25 100644 --- a/src/execute-query/queryresultrowtransformer.ts +++ b/src/execute-query/queryresultrowtransformer.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/execute-query/types.ts b/src/execute-query/types.ts index e374bd4bb..012f7713d 100644 --- a/src/execute-query/types.ts +++ b/src/execute-query/types.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/execute-query/values.ts b/src/execute-query/values.ts index 1186c984a..9325bb526 100644 --- a/src/execute-query/values.ts +++ b/src/execute-query/values.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/src/utils/retry-options.ts b/src/utils/retry-options.ts index d50c82fe1..9690c7849 100644 --- a/src/utils/retry-options.ts +++ b/src/utils/retry-options.ts @@ -1,4 +1,4 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. diff --git a/test/base64keymap.ts b/test/base64keymap.ts index 48cb55662..9691e0bbc 100644 --- a/test/base64keymap.ts +++ b/test/base64keymap.ts @@ -1,10 +1,10 @@ -// Copyright 2024 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // -// http://www.apache.org/licenses/LICENSE-2.0 +// https://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, diff --git a/test/bytebuffertransformer.ts b/test/bytebuffertransformer.ts index c5f541ad7..d96f49b90 100644 --- a/test/bytebuffertransformer.ts +++ b/test/bytebuffertransformer.ts @@ -4,13 +4,14 @@ // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // -// http://www.apache.org/licenses/LICENSE-2.0 +// https://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. + import * as assert from 'assert'; import {describe, it} from 'mocha'; import * as sinon from 'sinon'; diff --git a/testproxy/services/execute-query.js b/testproxy/services/execute-query.js index efb73ac2d..1c6ba26ef 100644 --- a/testproxy/services/execute-query.js +++ b/testproxy/services/execute-query.js @@ -1,4 +1,4 @@ -// Copyright 2022 Google LLC +// Copyright 2025 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -11,6 +11,7 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. + 'use strict'; const grpc = require('@grpc/grpc-js'); From f67221c50787c16432dcfe0b6845847e3fb9ec8d Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Thu, 12 Jun 2025 14:21:45 +0200 Subject: [PATCH 03/14] fix typo in ExecuteQuery request --- src/execute-query/executequerystatemachine.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/execute-query/executequerystatemachine.ts b/src/execute-query/executequerystatemachine.ts index 5f458a47d..d63eec9e6 100644 --- a/src/execute-query/executequerystatemachine.ts +++ b/src/execute-query/executequerystatemachine.ts @@ -272,7 +272,7 @@ export class ExecuteQueryStateMachine { private createValuesStream = (): AbortableDuplex => { const reqOpts: google.bigtable.v2.IExecuteQueryRequest = { ...this.requestParams, - preparedStatement: this.lastPreparedStatementBytes, + preparedQuery: this.lastPreparedStatementBytes, resumeToken: this.callerStream.getLatestResumeToken(), }; From d0aa12564d6af119d7301cc9f123f2f4cdead228 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Thu, 12 Jun 2025 14:50:52 +0200 Subject: [PATCH 04/14] remove obsolete protobuf field --- src/instance.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/instance.ts b/src/instance.ts index 1acceaf6e..93984058c 100644 --- a/src/instance.ts +++ b/src/instance.ts @@ -1573,7 +1573,6 @@ Please use the format 'my-instance' or '${bigtable.projectName}/instances/my-ins instanceName: this.name, appProfileId: this.bigtable.appProfileId, query: opts.query, - protoFormat: google.bigtable.v2.ProtoFormat.create(), paramTypes: protoParamTypes, }, gaxOpts: opts.retryOptions, @@ -1708,7 +1707,6 @@ Please use the format 'my-instance' or '${bigtable.projectName}/instances/my-ins const reqOpts: google.bigtable.v2.IExecuteQueryRequest = { instanceName: this.name, appProfileId: this.bigtable.appProfileId, - protoFormat: google.bigtable.v2.ProtoFormat.create(), params: protoParams, }; From 9c0d019d7574c18b2d415f0098b90992c0c6bd99 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Thu, 12 Jun 2025 14:52:00 +0200 Subject: [PATCH 05/14] add sample --- .../api-reference-doc-snippets/instance.js | 62 +++++++++++-------- 1 file changed, 35 insertions(+), 27 deletions(-) diff --git a/samples/api-reference-doc-snippets/instance.js b/samples/api-reference-doc-snippets/instance.js index 67f116d97..a4be02fc2 100644 --- a/samples/api-reference-doc-snippets/instance.js +++ b/samples/api-reference-doc-snippets/instance.js @@ -369,7 +369,7 @@ const snippets = { const bigtable = new Bigtable(); const instance = bigtable.instance(instanceId); - const query = `SELECT + const query = `SELECT _key from \`${tableId}\` WHERE _key=@row_key`; const parameters = { @@ -377,56 +377,64 @@ const snippets = { }; // if query parameter types are ambiguous, you can pass types explicitly - const parameter_types = { + const parameterTypes = { row_key: Bigtable.ExecuteQueryTypes.String(), }; - const options = { + const prepareStatementOptions = { query, - parameters, - parameter_types, // optional + parameterTypes, }; - instance - .executeQuery(options) - .then(result => { - const rows = result[0]; + instance.prepareStatement(prepareStatementOptions).then( + preparedStatement => instance.executeQuery({ + preparedStatement, + parameters, }) - .catch(err => { - // Handle the error. - }); + ).then(result => { + const rows = result[0]; + }).catch(err => { + // Handle errors + }); // [END bigtable_api_execute_query] }, createExecuteQueryStream: (instanceId, tableId) => { // [START bigtable_api_create_query_stream] - const {Bigtable} = require('@google-cloud/bigtable'); + const { Bigtable } = require('@google-cloud/bigtable'); const bigtable = new Bigtable(); const instance = bigtable.instance(instanceId); - const query = `SELECT + const query = `SELECT _key from \`${tableId}\` WHERE _key=@row_key`; const parameters = { row_key: 'alincoln', }; - const options = { + const prepareStatementOptions = { query, - parameters, + parameterTypes, }; - instance - .createExecuteQueryStream(options) - .on('error', err => { - // Handle the error. - }) - .on('data', row => { - // `row` is a QueryResultRow object. - }) - .on('end', () => { - // All rows retrieved. - }); + instance.prepareStatement(prepareStatementOptions).then( + preparedStatement => { + instance + .createExecuteQueryStream({ + preparedStatement, + parameters, + }) + .on('error', err => { + // Handle the error. + }) + .on('data', row => { + // `row` is a QueryResultRow object. + }) + .on('end', () => { + // All rows retrieved. + }); + } + ) // If you anticipate many results, you can end a stream early to prevent // unnecessary processing. From da5764a01366c2db544e6b84b8a40bfbff461a75 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Thu, 12 Jun 2025 16:10:38 +0200 Subject: [PATCH 06/14] linter fixes --- .../api-reference-doc-snippets/instance.js | 36 +++++++++++-------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/samples/api-reference-doc-snippets/instance.js b/samples/api-reference-doc-snippets/instance.js index a4be02fc2..73820bbe6 100644 --- a/samples/api-reference-doc-snippets/instance.js +++ b/samples/api-reference-doc-snippets/instance.js @@ -376,7 +376,6 @@ const snippets = { row_key: 'alincoln', }; - // if query parameter types are ambiguous, you can pass types explicitly const parameterTypes = { row_key: Bigtable.ExecuteQueryTypes.String(), }; @@ -386,23 +385,27 @@ const snippets = { parameterTypes, }; - instance.prepareStatement(prepareStatementOptions).then( - preparedStatement => instance.executeQuery({ - preparedStatement, - parameters, + instance + .prepareStatement(prepareStatementOptions) + .then(preparedStatement => + instance.executeQuery({ + preparedStatement, + parameters, + }), + ) + .then(result => { + const rows = result[0]; }) - ).then(result => { - const rows = result[0]; - }).catch(err => { - // Handle errors - }); + .catch(err => { + // Handle errors + }); // [END bigtable_api_execute_query] }, createExecuteQueryStream: (instanceId, tableId) => { // [START bigtable_api_create_query_stream] - const { Bigtable } = require('@google-cloud/bigtable'); + const {Bigtable} = require('@google-cloud/bigtable'); const bigtable = new Bigtable(); const instance = bigtable.instance(instanceId); @@ -412,13 +415,17 @@ const snippets = { const parameters = { row_key: 'alincoln', }; + const parameterTypes = { + row_key: Bigtable.ExecuteQueryTypes.String(), + }; const prepareStatementOptions = { query, parameterTypes, }; - instance.prepareStatement(prepareStatementOptions).then( - preparedStatement => { + instance + .prepareStatement(prepareStatementOptions) + .then(preparedStatement => { instance .createExecuteQueryStream({ preparedStatement, @@ -433,8 +440,7 @@ const snippets = { .on('end', () => { // All rows retrieved. }); - } - ) + }); // If you anticipate many results, you can end a stream early to prevent // unnecessary processing. From d550c390ab93b4be3a55529dd52d10b6aeedd041 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Thu, 12 Jun 2025 19:45:16 +0200 Subject: [PATCH 07/14] fix ci --- .github/workflows/ci.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 883082c0b..bb39ad3d9 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -72,3 +72,4 @@ jobs: - uses: JustinBeckwith/linkinator-action@v1 with: paths: docs/ + concurrency: 1 From 5ac729daa79dcbe1044cd278b109bcea443162bd Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Thu, 12 Jun 2025 20:42:48 +0200 Subject: [PATCH 08/14] fix ci - retry the 429 code --- .github/workflows/ci.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index bb39ad3d9..6e119b319 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -73,3 +73,4 @@ jobs: with: paths: docs/ concurrency: 1 + retry: true From e0c568b6f7261fd076a3afa99703dbff5e36fd90 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Tue, 24 Jun 2025 15:05:25 +0200 Subject: [PATCH 09/14] fix sample --- samples/api-reference-doc-snippets/instance.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/samples/api-reference-doc-snippets/instance.js b/samples/api-reference-doc-snippets/instance.js index 73820bbe6..5683ea94c 100644 --- a/samples/api-reference-doc-snippets/instance.js +++ b/samples/api-reference-doc-snippets/instance.js @@ -377,7 +377,7 @@ const snippets = { }; const parameterTypes = { - row_key: Bigtable.ExecuteQueryTypes.String(), + row_key: Bigtable.SqlTypes.String(), }; const prepareStatementOptions = { @@ -387,7 +387,7 @@ const snippets = { instance .prepareStatement(prepareStatementOptions) - .then(preparedStatement => + .then(([preparedStatement]) => instance.executeQuery({ preparedStatement, parameters, From 42cb7b2f9f8ba1e9409cc70fd60555c26b559562 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Tue, 24 Jun 2025 16:36:26 +0200 Subject: [PATCH 10/14] revert the CI changes --- .github/workflows/ci.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 6e119b319..883082c0b 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -72,5 +72,3 @@ jobs: - uses: JustinBeckwith/linkinator-action@v1 with: paths: docs/ - concurrency: 1 - retry: true From e5ffa5f5321c1ae7051d741f571f806141afe4fc Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Fri, 27 Jun 2025 16:27:44 +0200 Subject: [PATCH 11/14] fixed isRstStreamError location after rebase --- src/tabular-api-surface.ts | 12 ------------ src/utils/createReadStreamInternal.ts | 26 +++++++++++++------------- 2 files changed, 13 insertions(+), 25 deletions(-) diff --git a/src/tabular-api-surface.ts b/src/tabular-api-surface.ts index 72cfe01ef..462658e92 100644 --- a/src/tabular-api-surface.ts +++ b/src/tabular-api-surface.ts @@ -691,15 +691,3 @@ export class PartialFailureError extends Error { } } } -// Retry on "received rst stream" errors -export function isRstStreamError(error: ServiceError): boolean { - if (error.code === 13 && error.message) { - const error_message = (error.message || '').toLowerCase(); - return ( - error.code === 13 && - (error_message.includes('rst_stream') || - error_message.includes('rst stream')) - ); - } - return false; -} diff --git a/src/utils/createReadStreamInternal.ts b/src/utils/createReadStreamInternal.ts index e1561715a..9158bdc03 100644 --- a/src/utils/createReadStreamInternal.ts +++ b/src/utils/createReadStreamInternal.ts @@ -364,19 +364,6 @@ export function createReadStreamInternal( rowStream = pumpify.obj([requestStream, chunkTransformer, toRowStream]); - // Retry on "received rst stream" errors - const isRstStreamError = (error: ServiceError): boolean => { - if (error.code === 13 && error.message) { - const error_message = (error.message || '').toLowerCase(); - return ( - error.code === 13 && - (error_message.includes('rst_stream') || - error_message.includes('rst stream')) - ); - } - return false; - }; - metricsCollector.handleStatusAndMetadata(requestStream); rowStream .on('error', (error: ServiceError) => { @@ -438,3 +425,16 @@ export function createReadStreamInternal( makeNewRequest(); return userStream; } + +// Retry on "received rst stream" errors +export function isRstStreamError(error: ServiceError): boolean { + if (error.code === 13 && error.message) { + const error_message = (error.message || '').toLowerCase(); + return ( + error.code === 13 && + (error_message.includes('rst_stream') || + error_message.includes('rst stream')) + ); + } + return false; +} From 2e781725a299bc16029091dcf832f71290631bab Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Fri, 27 Jun 2025 18:22:47 +0200 Subject: [PATCH 12/14] fix isRstStreamError reference in executequerystatemachine.ts --- src/execute-query/executequerystatemachine.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/execute-query/executequerystatemachine.ts b/src/execute-query/executequerystatemachine.ts index d63eec9e6..917bc8426 100644 --- a/src/execute-query/executequerystatemachine.ts +++ b/src/execute-query/executequerystatemachine.ts @@ -31,8 +31,8 @@ import {ProtobufReaderTransformer} from './protobufreadertransformer'; import {MetadataConsumer} from './metadataconsumer'; import { DEFAULT_BACKOFF_SETTINGS, - isRstStreamError, } from '../tabular-api-surface'; +import { isRstStreamError } from '../utils/createReadStreamInternal'; const pumpify = require('pumpify'); /** From 43ea4941ab60572481dcaa2955550fda3166d9b2 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Fri, 27 Jun 2025 19:33:04 +0200 Subject: [PATCH 13/14] linter --- src/execute-query/executequerystatemachine.ts | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/execute-query/executequerystatemachine.ts b/src/execute-query/executequerystatemachine.ts index 917bc8426..1f2b89a1b 100644 --- a/src/execute-query/executequerystatemachine.ts +++ b/src/execute-query/executequerystatemachine.ts @@ -29,10 +29,8 @@ import { import {ExecuteQueryStreamWithMetadata} from './values'; import {ProtobufReaderTransformer} from './protobufreadertransformer'; import {MetadataConsumer} from './metadataconsumer'; -import { - DEFAULT_BACKOFF_SETTINGS, -} from '../tabular-api-surface'; -import { isRstStreamError } from '../utils/createReadStreamInternal'; +import {DEFAULT_BACKOFF_SETTINGS} from '../tabular-api-surface'; +import {isRstStreamError} from '../utils/createReadStreamInternal'; const pumpify = require('pumpify'); /** From 9f6cdd8fd1d1ea1f8b1da0cc6a30c975032d6392 Mon Sep 17 00:00:00 2001 From: Kajetan Boroszko Date: Tue, 8 Jul 2025 15:10:39 +0200 Subject: [PATCH 14/14] fix EncodedKeyMap --- src/execute-query/values.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/execute-query/values.ts b/src/execute-query/values.ts index 9325bb526..66eca4aa9 100644 --- a/src/execute-query/values.ts +++ b/src/execute-query/values.ts @@ -171,11 +171,11 @@ export class EncodedKeyMap return this.map_impl.size; } - entries(): IterableIterator<[string | bigint | Uint8Array | null, SqlValue]> { + entries(): MapIterator<[string | bigint | Uint8Array | null, SqlValue]> { return this.map_impl.values(); } - keys(): IterableIterator { + keys(): MapIterator { const iterator = this.map_impl.values(); return { next: () => { @@ -189,7 +189,7 @@ export class EncodedKeyMap }; } - values(): IterableIterator { + values(): MapIterator { const iterator = this.map_impl.values(); return { next: () => { @@ -203,7 +203,7 @@ export class EncodedKeyMap }; } - [Symbol.iterator](): IterableIterator< + [Symbol.iterator](): MapIterator< [string | bigint | Uint8Array | null, SqlValue] > { return this.entries();