From 96c5d27972724ef0aae8789a238308f439b9cc99 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Fri, 31 Jan 2025 15:00:04 -0500 Subject: [PATCH 01/11] Create basic wiring for PrepareQuery This doesn't do anything yet. I will follow with: - Refactoring of type converion code to populate param types - Connect PreparedStatement to ExecuteQuery - PreparedStatement refresh implementation Change-Id: I98f32cd65fee3dd43605f6353f0a4d6ad52630c8 --- .../bigtable/data/v2/BigtableDataClient.java | 14 +++ .../internal/AbstractProtoStructReader.java | 11 +- .../data/v2/internal/PrepareQueryRequest.java | 52 +++++++++ .../data/v2/internal/PrepareResponse.java | 47 ++++++++ .../v2/internal/PreparedStatementImpl.java | 39 +++++++ .../data/v2/internal/TimestampUtil.java | 28 +++++ .../data/v2/models/sql/PreparedStatement.java | 22 ++++ .../data/v2/stub/EnhancedBigtableStub.java | 22 ++++ .../v2/stub/EnhancedBigtableStubSettings.java | 58 ++++++++++ .../data/v2/BigtableDataClientTests.java | 16 +++ .../v2/internal/PrepareQueryRequestTest.java | 41 +++++++ .../data/v2/internal/TimestampUtilTest.java | 41 +++++++ .../data/v2/stub/CookiesHolderTest.java | 69 ++++++++++- .../EnhancedBigtableStubSettingsTest.java | 53 +++++++++ .../v2/stub/EnhancedBigtableStubTest.java | 109 ++++++++++++++++++ .../bigtable/data/v2/stub/HeadersTest.java | 23 ++++ .../bigtable/data/v2/stub/RetryInfoTest.java | 62 +++++++++- 17 files changed, 695 insertions(+), 12 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareResponse.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtil.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtilTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 8572947e77..9777eaee33 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -30,6 +30,9 @@ import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ResultSetImpl; import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord; @@ -48,7 +51,9 @@ import com.google.cloud.bigtable.data.v2.models.SampleRowKeysRequest; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream; @@ -56,6 +61,7 @@ import com.google.protobuf.ByteString; import java.io.IOException; import java.util.List; +import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -2729,6 +2735,14 @@ public ResultSet executeQuery(Statement statement) { return ResultSetImpl.create(stream); } + // TODO document once BoundStatement exists and executeQuery is updated so usage is clear + @BetaApi + public PreparedStatement prepareStatement(String query, Map> paramTypes) { + PrepareQueryRequest request = PrepareQueryRequest.create(query, paramTypes); + PrepareResponse response = stub.prepareQueryCallable().call(request); + return PreparedStatementImpl.create(response); + } + /** Close the clients and releases all associated resources. */ @Override public void close() { diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReader.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReader.java index 2a74fccd22..953db55182 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReader.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReader.java @@ -24,7 +24,6 @@ import com.google.cloud.bigtable.data.v2.models.sql.StructReader; import com.google.common.base.Preconditions; import com.google.protobuf.ByteString; -import com.google.protobuf.Timestamp; import java.time.Instant; import java.util.ArrayList; import java.util.Collections; @@ -169,7 +168,7 @@ public boolean getBoolean(String columnName) { public Instant getTimestamp(int columnIndex) { checkNonNullOfType(columnIndex, SqlType.timestamp(), columnIndex); Value value = values().get(columnIndex); - return toInstant(value.getTimestampValue()); + return TimestampUtil.toInstant(value.getTimestampValue()); } @Override @@ -177,7 +176,7 @@ public Instant getTimestamp(String columnName) { int columnIndex = getColumnIndex(columnName); checkNonNullOfType(columnIndex, SqlType.timestamp(), columnName); Value value = values().get(columnIndex); - return toInstant(value.getTimestampValue()); + return TimestampUtil.toInstant(value.getTimestampValue()); } @Override @@ -275,7 +274,7 @@ Object decodeValue(Value value, SqlType type) { case BOOL: return value.getBoolValue(); case TIMESTAMP: - return toInstant(value.getTimestampValue()); + return TimestampUtil.toInstant(value.getTimestampValue()); case DATE: return fromProto(value.getDateValue()); case STRUCT: @@ -329,10 +328,6 @@ private void checkNonNullOfType( } } - private Instant toInstant(Timestamp timestamp) { - return Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()); - } - private Date fromProto(com.google.type.Date proto) { return Date.fromYearMonthDay(proto.getYear(), proto.getMonth(), proto.getDay()); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java new file mode 100644 index 0000000000..cd6bfdb770 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java @@ -0,0 +1,52 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import com.google.api.core.InternalApi; +import com.google.auto.value.AutoValue; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import java.util.Map; + +/** + * Internal representation of PrepareQueryRequest that handles conversion from user-facing types to + * proto. + * + *

This is considered an internal implementation detail and should not be used by applications. + */ +@InternalApi("For internal use only") +@AutoValue +public abstract class PrepareQueryRequest { + + public abstract String query(); + + public abstract Map> paramTypes(); + + public static PrepareQueryRequest create(String query, Map> paramTypes) { + return new AutoValue_PrepareQueryRequest(query, paramTypes); + } + + public com.google.bigtable.v2.PrepareQueryRequest toProto(RequestContext requestContext) { + return com.google.bigtable.v2.PrepareQueryRequest.newBuilder() + .setInstanceName( + NameUtil.formatInstanceName( + requestContext.getProjectId(), requestContext.getInstanceId())) + .setAppProfileId(requestContext.getAppProfileId()) + .setQuery(query()) + // TODO validate and convert params to protobuf types + // .putAllParamTypes() + .build(); + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareResponse.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareResponse.java new file mode 100644 index 0000000000..35247e2dc9 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareResponse.java @@ -0,0 +1,47 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import com.google.api.core.InternalApi; +import com.google.auto.value.AutoValue; +import com.google.bigtable.v2.PrepareQueryResponse; +import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.protobuf.ByteString; +import java.time.Instant; + +/** + * Wrapper for results of a PrepareQuery call. + * + *

This should only be managed by {@link + * com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement}, and never used directly by users + * + *

This is considered an internal implementation detail and should not be used by applications. + */ +@InternalApi("For internal use only") +@AutoValue +public abstract class PrepareResponse { + public abstract ResultSetMetadata resultSetMetadata(); + + public abstract ByteString preparedQuery(); + + public abstract Instant validUntil(); + + public static PrepareResponse fromProto(PrepareQueryResponse proto) { + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(proto.getMetadata()); + Instant validUntil = TimestampUtil.toInstant(proto.getValidUntil()); + return new AutoValue_PrepareResponse(metadata, proto.getPreparedQuery(), validUntil); + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java new file mode 100644 index 0000000000..67eb41a3a3 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java @@ -0,0 +1,39 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import com.google.api.core.InternalApi; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; + +/** + * Implementation of PreparedStatement that handles PreparedQuery refresh + * + *

This is considered an internal implementation detail and should not be used by applications. + */ +// TODO Add ability to create BoundStatement +// TODO implement plan refresh +@InternalApi("For internal use only") +public class PreparedStatementImpl implements PreparedStatement { + private PrepareResponse response; + + public PreparedStatementImpl(PrepareResponse response) { + this.response = response; + } + + public static PreparedStatement create(PrepareResponse response) { + return new PreparedStatementImpl(response); + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtil.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtil.java new file mode 100644 index 0000000000..d659e03c2c --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtil.java @@ -0,0 +1,28 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import com.google.api.core.InternalApi; +import com.google.protobuf.Timestamp; +import java.time.Instant; + +/** For internal use only. Utility for converting proto timestamps to appropriate Java types. */ +@InternalApi("For internal use only") +public class TimestampUtil { + public static Instant toInstant(Timestamp timestamp) { + return Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()); + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java new file mode 100644 index 0000000000..95f0c6ed0f --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java @@ -0,0 +1,22 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.models.sql; + +import com.google.api.core.BetaApi; + +// TODO document once you can do something with this +@BetaApi +public interface PreparedStatement {} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 2290d412e3..f8573697a7 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -69,6 +69,8 @@ import com.google.bigtable.v2.SampleRowKeysResponse; import com.google.cloud.bigtable.Version; import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.BulkMutation; @@ -198,6 +200,7 @@ public class EnhancedBigtableStub implements AutoCloseable { readChangeStreamCallable; private final ExecuteQueryCallable executeQueryCallable; + private final UnaryCallable prepareQueryCallable; public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings) throws IOException { @@ -324,6 +327,7 @@ public EnhancedBigtableStub( readChangeStreamCallable = createReadChangeStreamCallable(new DefaultChangeStreamRecordAdapter()); executeQueryCallable = createExecuteQueryCallable(); + prepareQueryCallable = createPrepareQueryCallable(); } // @@ -1221,6 +1225,15 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { requestContext); } + private UnaryCallable createPrepareQueryCallable() { + return createUnaryCallable( + BigtableGrpc.getPrepareQueryMethod(), + req -> composeInstanceLevelRequestParams(req.getInstanceName(), req.getAppProfileId()), + settings.prepareQuerySettings(), + req -> req.toProto(requestContext), + PrepareResponse::fromProto); + } + /** * Wraps a callable chain in a user presentable callable that will inject the default call context * and trace the call. @@ -1242,6 +1255,11 @@ private Map composeRequestParams( return ImmutableMap.of("table_name", tableName, "app_profile_id", appProfileId); } + private Map composeInstanceLevelRequestParams( + String instanceName, String appProfileId) { + return ImmutableMap.of("name", instanceName, "app_profile_id", appProfileId); + } + private UnaryCallable createUnaryCallable( MethodDescriptor methodDescriptor, RequestParamsExtractor headerParamsFn, @@ -1472,6 +1490,10 @@ public ExecuteQueryCallable executeQueryCallable() { return executeQueryCallable; } + public UnaryCallable prepareQueryCallable() { + return prepareQueryCallable; + } + // private SpanName getSpanName(String methodName) { diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index 5e9e2cfe08..f1b8221fc8 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -35,6 +35,8 @@ import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.PingAndWarmRequest; import com.google.cloud.bigtable.Version; +import com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; @@ -211,6 +213,20 @@ public class EnhancedBigtableStubSettings extends StubSettings pingAndWarmSettings; private final ServerStreamingCallSettings executeQuerySettings; + private final UnaryCallSettings prepareQuerySettings; private final FeatureFlags featureFlags; @@ -306,6 +323,7 @@ private EnhancedBigtableStubSettings(Builder builder) { readChangeStreamSettings = builder.readChangeStreamSettings.build(); pingAndWarmSettings = builder.pingAndWarmSettings.build(); executeQuerySettings = builder.executeQuerySettings.build(); + prepareQuerySettings = builder.prepareQuerySettings.build(); featureFlags = builder.featureFlags.build(); } @@ -664,6 +682,31 @@ public ServerStreamingCallSettings executeQuerySettings() { return executeQuerySettings; } + /** + * Returns the object with the settings used for a PrepareQuery request. This is used by + * PreparedStatement to manage PreparedQueries. + * + *

This is an idempotent and non-streaming operation. + * + *

Default retry and timeout settings: + * + *

+ * + * @see RetrySettings for more explanation. + */ + public UnaryCallSettings prepareQuerySettings() { + return prepareQuerySettings; + } /** * Returns the object with the settings used for calls to PingAndWarm. * @@ -706,6 +749,8 @@ public static class Builder extends StubSettings.Builder pingAndWarmSettings; private final ServerStreamingCallSettings.Builder executeQuerySettings; + private final UnaryCallSettings.Builder + prepareQuerySettings; private FeatureFlags.Builder featureFlags; @@ -844,6 +889,11 @@ private Builder() { .setIdleTimeout(Duration.ofMinutes(5)) .setWaitTimeout(Duration.ofMinutes(5)); + prepareQuerySettings = UnaryCallSettings.newUnaryCallSettingsBuilder(); + prepareQuerySettings + .setRetryableCodes(IDEMPOTENT_RETRY_CODES) + .setRetrySettings(PREPARE_QUERY_RETRY_SETTINGS); + featureFlags = FeatureFlags.newBuilder() .setReverseScans(true) @@ -879,6 +929,7 @@ private Builder(EnhancedBigtableStubSettings settings) { readChangeStreamSettings = settings.readChangeStreamSettings.toBuilder(); pingAndWarmSettings = settings.pingAndWarmSettings.toBuilder(); executeQuerySettings = settings.executeQuerySettings().toBuilder(); + prepareQuerySettings = settings.prepareQuerySettings().toBuilder(); featureFlags = settings.featureFlags.toBuilder(); } // @@ -1168,6 +1219,12 @@ public ServerStreamingCallSettings.Builder executeQuerySettin return executeQuerySettings; } + /** Returns the builder with the settings used for calls to PrepareQuery */ + @BetaApi + public UnaryCallSettings.Builder prepareQuerySettings() { + return prepareQuerySettings; + } + @SuppressWarnings("unchecked") public EnhancedBigtableStubSettings build() { Preconditions.checkState(projectId != null, "Project id must be set"); @@ -1240,6 +1297,7 @@ public String toString() { .add("readChangeStreamSettings", readChangeStreamSettings) .add("pingAndWarmSettings", pingAndWarmSettings) .add("executeQuerySettings", executeQuerySettings) + .add("prepareQuerySettings", prepareQuerySettings) .add("metricsProvider", metricsProvider) .add("metricsEndpoint", metricsEndpoint) .add("parent", super.toString()) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java index 880744bc18..f6d50e1c39 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java @@ -24,6 +24,8 @@ import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecord; @@ -42,12 +44,15 @@ import com.google.cloud.bigtable.data.v2.models.SampleRowKeysRequest; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import com.google.protobuf.Empty; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -89,6 +94,7 @@ public class BigtableDataClientTests { @Mock private UnaryCallable mockBulkMutateRowsCallable; @Mock private Batcher mockBulkMutationBatcher; @Mock private Batcher mockBulkReadRowsBatcher; + @Mock private UnaryCallable mockPrepareQueryCallable; @Mock(answer = Answers.RETURNS_DEEP_STUBS) private ServerStreamingCallable @@ -1059,4 +1065,14 @@ public void proxyReadModifyWriterRowCallableTest() { assertThat(bigtableDataClient.readModifyWriteRowCallable()) .isSameInstanceAs(mockReadModifyWriteRowCallable); } + + @Test + public void prepareQueryTest() { + Mockito.when(mockStub.prepareQueryCallable()).thenReturn(mockPrepareQueryCallable); + + String query = "SELECT * FROM table"; + Map> paramTypes = new HashMap<>(); + bigtableDataClient.prepareStatement(query, paramTypes); + Mockito.verify(mockPrepareQueryCallable).call(PrepareQueryRequest.create(query, paramTypes)); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java new file mode 100644 index 0000000000..2db0d8dfea --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java @@ -0,0 +1,41 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import static com.google.common.truth.Truth.assertThat; + +import java.util.HashMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class PrepareQueryRequestTest { + + @Test + public void testProtoConversion() { + // TODO test param conversion when supported + PrepareQueryRequest request = + PrepareQueryRequest.create("SELECT * FROM table", new HashMap<>()); + RequestContext rc = RequestContext.create("project", "instance", "profile"); + com.google.bigtable.v2.PrepareQueryRequest proto = request.toProto(rc); + + assertThat(proto.getQuery()).isEqualTo("SELECT * FROM table"); + assertThat(proto.getAppProfileId()).isEqualTo("profile"); + assertThat(proto.getInstanceName()) + .isEqualTo(NameUtil.formatInstanceName("project", "instance")); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtilTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtilTest.java new file mode 100644 index 0000000000..5c6dac7632 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/TimestampUtilTest.java @@ -0,0 +1,41 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.protobuf.Timestamp; +import java.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class TimestampUtilTest { + + @Test + public void testToInstant() { + assertThat(TimestampUtil.toInstant(Timestamp.getDefaultInstance())) + .isEqualTo(Instant.ofEpochSecond(0)); + assertThat(TimestampUtil.toInstant(Timestamp.newBuilder().setSeconds(1000).build())) + .isEqualTo(Instant.ofEpochSecond(1000)); + assertThat( + TimestampUtil.toInstant(Timestamp.newBuilder().setSeconds(2000).setNanos(3000).build())) + .isEqualTo(Instant.ofEpochSecond(2000, 3000)); + assertThat(TimestampUtil.toInstant(Timestamp.newBuilder().setNanos(3000).build())) + .isEqualTo(Instant.ofEpochSecond(0, 3000)); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java index a58dbe2c77..258819e24b 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java @@ -16,6 +16,9 @@ package com.google.cloud.bigtable.data.v2.stub; import static com.google.cloud.bigtable.data.v2.MetadataSubject.assertThat; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; import com.google.api.gax.retrying.RetrySettings; @@ -29,6 +32,8 @@ import com.google.bigtable.v2.MutateRowResponse; import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.MutateRowsResponse; +import com.google.bigtable.v2.PrepareQueryRequest; +import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.ReadChangeStreamRequest; import com.google.bigtable.v2.ReadChangeStreamResponse; import com.google.bigtable.v2.ReadModifyWriteRowRequest; @@ -64,6 +69,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -336,6 +342,28 @@ public void testGenerateInitialChangeStreamPartition() { serverMetadata.clear(); } + @Test + public void testPrepareQuery() { + client.prepareStatement("SELECT * FROM table", new HashMap<>()); + + assertThat(fakeService.count.get()).isGreaterThan(1); + assertThat(serverMetadata).hasSize(fakeService.count.get()); + + Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); + + assertThat(lastMetadata) + .containsAtLeast( + ROUTING_COOKIE_1.name(), + "prepareQuery", + ROUTING_COOKIE_2.name(), + testCookie, + ROUTING_COOKIE_HEADER.name(), + testHeaderCookie); + assertThat(lastMetadata).doesNotContainKeys(BAD_KEY.name()); + + serverMetadata.clear(); + } + @Test public void testNoCookieSucceedReadRows() { fakeService.returnCookie = false; @@ -461,6 +489,23 @@ public void testNoCookieSucceedGenerateInitialChangeStreamParition() { serverMetadata.clear(); } + @Test + public void testNoCookieSucceedPrepareQuery() { + fakeService.returnCookie = false; + + client.prepareStatement("SELECT * FROM table", new HashMap<>()); + + assertThat(fakeService.count.get()).isGreaterThan(1); + assertThat(serverMetadata).hasSize(fakeService.count.get()); + + Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); + + assertThat(lastMetadata).doesNotContainKeys(ROUTING_COOKIE_2.name(), BAD_KEY.name()); + assertThat(lastMetadata).containsAtLeast(ROUTING_COOKIE_1.name(), routingCookie1Header); + + serverMetadata.clear(); + } + @Test public void testCookiesInHeaders() throws Exception { // Send 2 cookies in the headers, with routingCookieKey and ROUTING_COOKIE_2. ROUTING_COOKIE_2 @@ -551,6 +596,9 @@ public void testAllMethodsAreCalled() { for (ChangeStreamRecord record : client.readChangeStream(ReadChangeStreamQuery.create("fake-table"))) {} + fakeService.count.set(0); + client.prepareStatement("SELECT * FROM table", new HashMap<>()); + Set expected = BigtableGrpc.getServiceDescriptor().getMethods().stream() .map(MethodDescriptor::getBareMethodName) @@ -558,8 +606,7 @@ public void testAllMethodsAreCalled() { // Exclude methods that are not supported by routing cookie methods.add("PingAndWarm"); - methods.add("ExecuteQuery"); - methods.add("PrepareQuery"); + methods.add("ExecuteQuery"); // TODO remove when retries are implemented assertThat(methods).containsExactlyElementsIn(expected); } @@ -796,6 +843,24 @@ public void generateInitialChangeStreamPartitions( responseObserver.onCompleted(); } + @Override + public void prepareQuery( + PrepareQueryRequest request, StreamObserver responseObserver) { + if (count.getAndIncrement() < 1) { + Metadata trailers = new Metadata(); + maybePopulateCookie(trailers, "prepareQuery"); + StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); + responseObserver.onError(exception); + return; + } + responseObserver.onNext( + // Need to set metadata for response to parse + PrepareQueryResponse.newBuilder() + .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .build()); + responseObserver.onCompleted(); + } + private void maybePopulateCookie(Metadata trailers, String label) { if (returnCookie) { trailers.put(ROUTING_COOKIE_1, label); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java index fdc6b5717e..472c74e511 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java @@ -28,6 +28,8 @@ import com.google.api.gax.rpc.WatchdogProvider; import com.google.auth.Credentials; import com.google.bigtable.v2.PingAndWarmRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.KeyOffset; @@ -846,6 +848,56 @@ public void executeQueryRetriesAreDisabled() { assertThat(builder.getRetrySettings().getInitialRpcTimeout()).isAtMost(Duration.ofSeconds(30)); } + @Test + public void prepareQuerySettingsAreNotLost() { + String dummyProjectId = "my-project"; + String dummyInstanceId = "my-instance"; + + EnhancedBigtableStubSettings.Builder builder = + EnhancedBigtableStubSettings.newBuilder() + .setProjectId(dummyProjectId) + .setInstanceId(dummyInstanceId) + // Here and everywhere in this test, disable channel priming so we won't need + // authentication for sending the prime request since we're only testing the settings. + .setRefreshingChannel(false); + + RetrySettings retrySettings = + RetrySettings.newBuilder() + .setMaxAttempts(10) + .setTotalTimeout(Duration.ofHours(1)) + .setInitialRpcTimeout(Duration.ofSeconds(10)) + .setRpcTimeoutMultiplier(1) + .setMaxRpcTimeout(Duration.ofSeconds(10)) + .setJittered(true) + .build(); + + builder + .prepareQuerySettings() + .setRetryableCodes(Code.ABORTED, Code.DEADLINE_EXCEEDED) + .setRetrySettings(retrySettings) + .build(); + + assertThat(builder.prepareQuerySettings().getRetryableCodes()) + .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED); + assertThat(builder.prepareQuerySettings().getRetrySettings()).isEqualTo(retrySettings); + + assertThat(builder.build().prepareQuerySettings().getRetryableCodes()) + .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED); + assertThat(builder.build().prepareQuerySettings().getRetrySettings()).isEqualTo(retrySettings); + + assertThat(builder.build().toBuilder().prepareQuerySettings().getRetryableCodes()) + .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED); + assertThat(builder.build().toBuilder().prepareQuerySettings().getRetrySettings()) + .isEqualTo(retrySettings); + } + + @Test + public void prepareQueryHasSaneDefaults() { + UnaryCallSettings.Builder builder = + EnhancedBigtableStubSettings.newBuilder().prepareQuerySettings(); + verifyRetrySettingAreSane(builder.getRetryableCodes(), builder.getRetrySettings()); + } + private void verifyRetrySettingAreSane(Set retryCodes, RetrySettings retrySettings) { assertThat(retryCodes).containsAtLeast(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE); @@ -974,6 +1026,7 @@ public void enableRetryInfoFalseValueTest() throws IOException { "readChangeStreamSettings", "pingAndWarmSettings", "executeQuerySettings", + "prepareQuerySettings", "metricsProvider", "metricsEndpoint", }; diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index 099c034d14..88ab52fb7d 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -56,6 +56,8 @@ import com.google.bigtable.v2.MutateRowsResponse; import com.google.bigtable.v2.PingAndWarmRequest; import com.google.bigtable.v2.PingAndWarmResponse; +import com.google.bigtable.v2.PrepareQueryRequest; +import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.ReadChangeStreamRequest; import com.google.bigtable.v2.ReadChangeStreamResponse; import com.google.bigtable.v2.ReadModifyWriteRowRequest; @@ -67,6 +69,8 @@ import com.google.cloud.bigtable.admin.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.BulkMutation; @@ -94,6 +98,7 @@ import com.google.protobuf.ByteString; import com.google.protobuf.BytesValue; import com.google.protobuf.StringValue; +import com.google.protobuf.Timestamp; import com.google.rpc.Code; import com.google.rpc.Status; import io.grpc.CallOptions; @@ -120,8 +125,10 @@ import java.security.KeyPair; import java.security.KeyPairGenerator; import java.security.NoSuchAlgorithmException; +import java.time.Instant; import java.util.Base64; import java.util.Collection; +import java.util.HashMap; import java.util.Iterator; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.BlockingQueue; @@ -141,6 +148,7 @@ public class EnhancedBigtableStubTest { private static final String PROJECT_ID = "fake-project"; private static final String INSTANCE_ID = "fake-instance"; + private static final String INSTANCE_NAME = NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID); private static final String TABLE_ID = "fake-table"; private static final String TABLE_NAME = NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID); @@ -404,6 +412,86 @@ public void testMutateRowErrorPropagation() { assertThat(invocationCount.get()).isEqualTo(2); } + @Test + public void testPrepareQueryRequestResponseConversion() + throws ExecutionException, InterruptedException { + com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest req = + com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest.create( + "SELECT * FROM TABLE", new HashMap<>()); + CallOptions.Key testKey = CallOptions.Key.create("test-key"); + + GrpcCallContext ctx = + GrpcCallContext.createDefault() + .withCallOptions(CallOptions.DEFAULT.withOption(testKey, "callopt-value")); + ApiFuture f = enhancedBigtableStub.prepareQueryCallable().futureCall(req, ctx); + f.get(); + + PrepareQueryRequest protoReq = fakeDataService.prepareRequests.poll(1, TimeUnit.SECONDS); + assertThat(protoReq) + .isEqualTo(req.toProto(RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID))); + assertThat(f.get().resultSetMetadata()) + .isEqualTo( + ProtoResultSetMetadata.fromProto( + metadata(columnMetadata("foo", stringType())).getMetadata())); + assertThat(f.get().preparedQuery()).isEqualTo(ByteString.copyFromUtf8("foo")); + assertThat(f.get().validUntil()).isEqualTo(Instant.ofEpochSecond(1000, 1000)); + } + + @Test + public void testPrepareQueryRequestParams() throws ExecutionException, InterruptedException { + com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest req = + com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest.create( + "SELECT * FROM TABLE", new HashMap<>()); + + ApiFuture f = + enhancedBigtableStub.prepareQueryCallable().futureCall(req, null); + f.get(); + + Metadata reqMetadata = metadataInterceptor.headers.poll(1, TimeUnit.SECONDS); + + // RequestParamsExtractor + String reqParams = + reqMetadata.get(Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER)); + assertThat(reqParams).contains("name=" + INSTANCE_NAME.replace("/", "%2F")); + assertThat(reqParams).contains(String.format("app_profile_id=%s", APP_PROFILE_ID)); + + // StatsHeadersUnaryCallable + assertThat(reqMetadata.keys()).contains("bigtable-client-attempt-epoch-usec"); + + assertThat(f.get().resultSetMetadata()) + .isEqualTo( + ProtoResultSetMetadata.fromProto( + metadata(columnMetadata("foo", stringType())).getMetadata())); + assertThat(f.get().preparedQuery()).isEqualTo(ByteString.copyFromUtf8("foo")); + assertThat(f.get().validUntil()).isEqualTo(Instant.ofEpochSecond(1000, 1000)); + } + + @Test + public void testPrepareQueryErrorPropagation() { + AtomicInteger invocationCount = new AtomicInteger(); + Mockito.doAnswer( + invocationOnMock -> { + StreamObserver observer = invocationOnMock.getArgument(1); + if (invocationCount.getAndIncrement() == 0) { + observer.onError(io.grpc.Status.UNAVAILABLE.asRuntimeException()); + } else { + observer.onError(io.grpc.Status.FAILED_PRECONDITION.asRuntimeException()); + } + return null; + }) + .when(fakeDataService) + .prepareQuery(Mockito.any(), Mockito.any(StreamObserver.class)); + com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest req = + com.google.cloud.bigtable.data.v2.internal.PrepareQueryRequest.create( + "SELECT * FROM TABLE", new HashMap<>()); + ApiFuture f = + enhancedBigtableStub.prepareQueryCallable().futureCall(req, null); + + ExecutionException e = assertThrows(ExecutionException.class, f::get); + assertThat(e.getCause()).isInstanceOf(FailedPreconditionException.class); + assertThat(invocationCount.get()).isEqualTo(2); + } + @Test public void testCreateReadRowsCallable() throws InterruptedException { ServerStreamingCallable streamingCallable = @@ -891,6 +979,7 @@ private static class FakeDataService extends BigtableGrpc.BigtableImplBase { final BlockingQueue checkAndMutateRowRequests = Queues.newLinkedBlockingDeque(); final BlockingQueue rmwRequests = Queues.newLinkedBlockingDeque(); + final BlockingQueue prepareRequests = Queues.newLinkedBlockingDeque(); @SuppressWarnings("unchecked") ReadRowsRequest popLastRequest() throws InterruptedException { @@ -1007,5 +1096,25 @@ public void executeQuery( responseObserver.onNext(metadata(columnMetadata("foo", stringType()))); responseObserver.onNext(partialResultSetWithToken(stringValue("test"))); } + + @Override + public void prepareQuery( + PrepareQueryRequest request, StreamObserver responseObserver) { + if (request.getQuery().contains(WAIT_TIME_QUERY)) { + try { + Thread.sleep(WATCHDOG_CHECK_DURATION.toMillis() * 2); + } catch (Exception e) { + + } + } + prepareRequests.add(request); + responseObserver.onNext( + PrepareQueryResponse.newBuilder() + .setPreparedQuery(ByteString.copyFromUtf8("foo")) + .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .setValidUntil(Timestamp.newBuilder().setSeconds(1000).setNanos(1000).build()) + .build()); + responseObserver.onCompleted(); + } } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java index 16e886f9b7..bee8374d9b 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java @@ -15,6 +15,9 @@ */ package com.google.cloud.bigtable.data.v2.stub; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; import com.google.api.gax.batching.Batcher; @@ -27,6 +30,8 @@ import com.google.bigtable.v2.MutateRowResponse; import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.MutateRowsResponse; +import com.google.bigtable.v2.PrepareQueryRequest; +import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.ReadModifyWriteRowRequest; import com.google.bigtable.v2.ReadModifyWriteRowResponse; import com.google.bigtable.v2.ReadRowsRequest; @@ -50,6 +55,7 @@ import io.grpc.ServerCallHandler; import io.grpc.ServerInterceptor; import io.grpc.stub.StreamObserver; +import java.util.HashMap; import java.util.concurrent.ArrayBlockingQueue; import java.util.concurrent.BlockingQueue; import org.junit.After; @@ -169,6 +175,12 @@ public void executeQueryTest() { verifyHeaderSent(true); } + @Test + public void prepareQueryTest() { + client.prepareStatement("SELECT * FROM table", new HashMap<>()); + verifyHeaderSent(true); + } + private void verifyHeaderSent() { verifyHeaderSent(false); } @@ -259,5 +271,16 @@ public void readModifyWriteRow( responseObserver.onNext(ReadModifyWriteRowResponse.getDefaultInstance()); responseObserver.onCompleted(); } + + @Override + public void prepareQuery( + PrepareQueryRequest request, StreamObserver responseObserver) { + responseObserver.onNext( + // Need to set metadata for response to parse + PrepareQueryResponse.newBuilder() + .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .build()); + responseObserver.onCompleted(); + } } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java index 8f97518232..024d36e018 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java @@ -15,6 +15,9 @@ */ package com.google.cloud.bigtable.data.v2.stub; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; @@ -32,6 +35,8 @@ import com.google.bigtable.v2.MutateRowResponse; import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.MutateRowsResponse; +import com.google.bigtable.v2.PrepareQueryRequest; +import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.ReadChangeStreamRequest; import com.google.bigtable.v2.ReadChangeStreamResponse; import com.google.bigtable.v2.ReadModifyWriteRowRequest; @@ -71,6 +76,7 @@ import io.grpc.stub.StreamObserver; import java.io.IOException; import java.time.Duration; +import java.util.HashMap; import java.util.HashSet; import java.util.Queue; import java.util.Set; @@ -195,6 +201,9 @@ public void testAllMethods() { verifyRetryInfoIsUsed( () -> client.generateInitialChangeStreamPartitions("table").iterator().hasNext(), true); + attemptCounter.set(0); + verifyRetryInfoIsUsed( + () -> client.prepareStatement("SELECT * FROM table", new HashMap<>()), true); // Verify that the new data API methods are tested or excluded. This is enforced by // introspecting grpc // method descriptors. @@ -205,8 +214,7 @@ public void testAllMethods() { // Exclude methods that don't support retry info methods.add("PingAndWarm"); - methods.add("ExecuteQuery"); - methods.add("PrepareQuery"); + methods.add("ExecuteQuery"); // TODO remove when retries are implemented assertThat(methods).containsExactlyElementsIn(expected); } @@ -538,6 +546,39 @@ public void testGenerateInitialChangeStreamServerNotReturningRetryInfoClientDisa } } + @Test + public void testPrepareQueryNonRetryableErrorWithRetryInfo() { + verifyRetryInfoIsUsed( + () -> client.prepareStatement("SELECT * FROM table", new HashMap<>()), false); + } + + @Test + public void testPrepareQueryDisableRetryInfo() throws IOException { + settings.stubSettings().setEnableRetryInfo(false); + + try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) { + + verifyRetryInfoCanBeDisabled( + () -> newClient.prepareStatement("SELECT * FROM table", new HashMap<>())); + } + } + + @Test + public void testPrepareQueryServerNotReturningRetryInfo() { + verifyNoRetryInfo(() -> client.prepareStatement("SELECT * FROM table", new HashMap<>()), true); + } + + @Test + public void testPrepareQueryServerNotReturningRetryInfoClientDisabledHandling() + throws IOException { + settings.stubSettings().setEnableRetryInfo(false); + + try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) { + verifyNoRetryInfo( + () -> newClient.prepareStatement("SELECT * FROM table", new HashMap<>()), true); + } + } + // Test the case where server returns retry info and client enables handling of retry info private void verifyRetryInfoIsUsed(Runnable runnable, boolean retryableError) { if (retryableError) { @@ -803,5 +844,22 @@ public void readChangeStream( responseObserver.onError(expectedRpc); } } + + @Override + public void prepareQuery( + PrepareQueryRequest request, StreamObserver responseObserver) { + attemptCounter.incrementAndGet(); + if (expectations.isEmpty()) { + responseObserver.onNext( + // Need to set metadata for response to parse + PrepareQueryResponse.newBuilder() + .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .build()); + responseObserver.onCompleted(); + } else { + Exception expectedRpc = expectations.poll(); + responseObserver.onError(expectedRpc); + } + } } } From ec0e1d9f6bd46350e682b8fea29df82c457954e5 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Fri, 31 Jan 2025 16:06:03 -0500 Subject: [PATCH 02/11] Refactor query param code to share w Prepare & Execute Use it to populate paramTypes in PrepareQueryRequest Change-Id: I8340741fd7377a9e2b87972366ad74ed0cb3354a --- .../data/v2/internal/PrepareQueryRequest.java | 11 ++- .../data/v2/internal/QueryParamUtil.java | 98 +++++++++++++++++++ .../data/v2/models/sql/Statement.java | 72 ++++---------- .../v2/internal/PrepareQueryRequestTest.java | 53 +++++++++- .../data/v2/internal/QueryParamUtilTest.java | 97 ++++++++++++++++++ 5 files changed, 272 insertions(+), 59 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtil.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtilTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java index cd6bfdb770..0a330d32c6 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequest.java @@ -17,7 +17,9 @@ import com.google.api.core.InternalApi; import com.google.auto.value.AutoValue; +import com.google.bigtable.v2.Type; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import java.util.HashMap; import java.util.Map; /** @@ -39,14 +41,19 @@ public static PrepareQueryRequest create(String query, Map> p } public com.google.bigtable.v2.PrepareQueryRequest toProto(RequestContext requestContext) { + HashMap protoParamTypes = new HashMap<>(paramTypes().size()); + for (Map.Entry> entry : paramTypes().entrySet()) { + Type proto = QueryParamUtil.convertToQueryParamProto(entry.getValue()); + protoParamTypes.put(entry.getKey(), proto); + } + return com.google.bigtable.v2.PrepareQueryRequest.newBuilder() .setInstanceName( NameUtil.formatInstanceName( requestContext.getProjectId(), requestContext.getInstanceId())) .setAppProfileId(requestContext.getAppProfileId()) .setQuery(query()) - // TODO validate and convert params to protobuf types - // .putAllParamTypes() + .putAllParamTypes(protoParamTypes) .build(); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtil.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtil.java new file mode 100644 index 0000000000..439f8f7205 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtil.java @@ -0,0 +1,98 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import com.google.api.core.InternalApi; +import com.google.bigtable.v2.Type; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType.Array; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType.Code; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; + +/** + * Helper to convert SqlTypes to protobuf query parameter representation + * + *

This is considered an internal implementation detail and should not be used by applications. + */ +@InternalApi("For internal use only") +public class QueryParamUtil { + private static final Type STRING_TYPE = + Type.newBuilder().setStringType(Type.String.getDefaultInstance()).build(); + private static final Type BYTES_TYPE = + Type.newBuilder().setBytesType(Type.Bytes.getDefaultInstance()).build(); + private static final Type INT64_TYPE = + Type.newBuilder().setInt64Type(Type.Int64.getDefaultInstance()).build(); + private static final Type FLOAT32_TYPE = + Type.newBuilder().setFloat32Type(Type.Float32.getDefaultInstance()).build(); + private static final Type FLOAT64_TYPE = + Type.newBuilder().setFloat64Type(Type.Float64.getDefaultInstance()).build(); + private static final Type BOOL_TYPE = + Type.newBuilder().setBoolType(Type.Bool.getDefaultInstance()).build(); + private static final Type TIMESTAMP_TYPE = + Type.newBuilder().setTimestampType(Type.Timestamp.getDefaultInstance()).build(); + private static final Type DATE_TYPE = + Type.newBuilder().setDateType(Type.Date.getDefaultInstance()).build(); + + private static final Set VALID_ARRAY_ELEMENT_TYPES = + new HashSet<>( + Arrays.asList( + Code.STRING, + Code.BYTES, + Code.INT64, + Code.FLOAT64, + Code.FLOAT32, + Code.BOOL, + Code.TIMESTAMP, + Code.DATE)); + + public static Type convertToQueryParamProto(SqlType sqlType) { + switch (sqlType.getCode()) { + case BYTES: + return BYTES_TYPE; + case STRING: + return STRING_TYPE; + case INT64: + return INT64_TYPE; + case FLOAT64: + return FLOAT64_TYPE; + case FLOAT32: + return FLOAT32_TYPE; + case BOOL: + return BOOL_TYPE; + case TIMESTAMP: + return TIMESTAMP_TYPE; + case DATE: + return DATE_TYPE; + case STRUCT: + throw new IllegalArgumentException("STRUCT is not a supported query parameter type"); + case MAP: + throw new IllegalArgumentException("MAP is not a supported query parameter type"); + case ARRAY: + SqlType.Array arrayType = (Array) sqlType; + if (!VALID_ARRAY_ELEMENT_TYPES.contains(arrayType.getElementType().getCode())) { + throw new IllegalArgumentException( + "Unsupported query parameter Array element type: " + arrayType.getElementType()); + } + Type elementType = convertToQueryParamProto(arrayType.getElementType()); + Type.Array arrayProto = Type.Array.newBuilder().setElementType(elementType).build(); + return Type.newBuilder().setArrayType(arrayProto).build(); + default: + throw new IllegalArgumentException("Unsupported Query parameter type: " + sqlType); + } + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java index c1831219a6..a2f3639691 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java @@ -23,6 +23,7 @@ import com.google.bigtable.v2.Value; import com.google.cloud.Date; import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.QueryParamUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; @@ -61,23 +62,6 @@ @BetaApi public class Statement { - private static final Type STRING_TYPE = - Type.newBuilder().setStringType(Type.String.getDefaultInstance()).build(); - private static final Type BYTES_TYPE = - Type.newBuilder().setBytesType(Type.Bytes.getDefaultInstance()).build(); - private static final Type INT64_TYPE = - Type.newBuilder().setInt64Type(Type.Int64.getDefaultInstance()).build(); - private static final Type FLOAT32_TYPE = - Type.newBuilder().setFloat32Type(Type.Float32.getDefaultInstance()).build(); - private static final Type FLOAT64_TYPE = - Type.newBuilder().setFloat64Type(Type.Float64.getDefaultInstance()).build(); - private static final Type BOOL_TYPE = - Type.newBuilder().setBoolType(Type.Bool.getDefaultInstance()).build(); - private static final Type TIMESTAMP_TYPE = - Type.newBuilder().setTimestampType(Type.Timestamp.getDefaultInstance()).build(); - private static final Type DATE_TYPE = - Type.newBuilder().setDateType(Type.Date.getDefaultInstance()).build(); - private final String sql; private final Map params; @@ -192,7 +176,8 @@ public Builder setListParam( } private static Value stringParamOf(@Nullable String value) { - Value.Builder builder = nullValueWithType(STRING_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.string()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setStringValue(value); } @@ -200,7 +185,8 @@ private static Value stringParamOf(@Nullable String value) { } private static Value bytesParamOf(@Nullable ByteString value) { - Value.Builder builder = nullValueWithType(BYTES_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.bytes()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setBytesValue(value); } @@ -208,7 +194,8 @@ private static Value bytesParamOf(@Nullable ByteString value) { } private static Value int64ParamOf(@Nullable Long value) { - Value.Builder builder = nullValueWithType(INT64_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.int64()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setIntValue(value); } @@ -216,7 +203,8 @@ private static Value int64ParamOf(@Nullable Long value) { } private static Value float32ParamOf(@Nullable Float value) { - Value.Builder builder = nullValueWithType(FLOAT32_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.float32()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setFloatValue(value); } @@ -224,7 +212,8 @@ private static Value float32ParamOf(@Nullable Float value) { } private static Value float64ParamOf(@Nullable Double value) { - Value.Builder builder = nullValueWithType(FLOAT64_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.float64()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setFloatValue(value); } @@ -232,7 +221,8 @@ private static Value float64ParamOf(@Nullable Double value) { } private static Value booleanParamOf(@Nullable Boolean value) { - Value.Builder builder = nullValueWithType(BOOL_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.bool()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setBoolValue(value); } @@ -240,7 +230,8 @@ private static Value booleanParamOf(@Nullable Boolean value) { } private static Value timestampParamOf(@Nullable Instant value) { - Value.Builder builder = nullValueWithType(TIMESTAMP_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.timestamp()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setTimestampValue(toTimestamp(value)); } @@ -248,7 +239,8 @@ private static Value timestampParamOf(@Nullable Instant value) { } private static Value dateParamOf(@Nullable Date value) { - Value.Builder builder = nullValueWithType(DATE_TYPE); + Type type = QueryParamUtil.convertToQueryParamProto(SqlType.date()); + Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setDateValue(toProtoDate(value)); } @@ -256,11 +248,7 @@ private static Value dateParamOf(@Nullable Date value) { } private static Value arrayParamOf(@Nullable List value, SqlType.Array arrayType) { - Type type = - Type.newBuilder() - .setArrayType( - Type.Array.newBuilder().setElementType(getElementType(arrayType)).build()) - .build(); + Type type = QueryParamUtil.convertToQueryParamProto(arrayType); Value.Builder builder = nullValueWithType(type); if (value != null) { builder.setArrayValue(arrayValueOf(value, arrayType)); @@ -268,30 +256,6 @@ private static Value arrayParamOf(@Nullable List value, SqlType.Array return builder.build(); } - private static Type getElementType(SqlType.Array arrayType) { - switch (arrayType.getElementType().getCode()) { - case BYTES: - return BYTES_TYPE; - case STRING: - return STRING_TYPE; - case INT64: - return INT64_TYPE; - case FLOAT32: - return FLOAT32_TYPE; - case FLOAT64: - return FLOAT64_TYPE; - case BOOL: - return BOOL_TYPE; - case TIMESTAMP: - return TIMESTAMP_TYPE; - case DATE: - return DATE_TYPE; - default: - throw new IllegalArgumentException( - "Unsupported query parameter Array element type: " + arrayType.getElementType()); - } - } - private static ArrayValue arrayValueOf(List value, SqlType.Array arrayType) { ArrayValue.Builder valueBuilder = ArrayValue.newBuilder(); for (Object element : value) { diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java index 2db0d8dfea..983bc8521c 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PrepareQueryRequestTest.java @@ -15,9 +15,21 @@ */ package com.google.cloud.bigtable.data.v2.internal; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.arrayType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.boolType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.dateType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.float32Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.float64Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.timestampType; import static com.google.common.truth.Truth.assertThat; +import com.google.bigtable.v2.Type; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import java.util.HashMap; +import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -27,9 +39,25 @@ public class PrepareQueryRequestTest { @Test public void testProtoConversion() { - // TODO test param conversion when supported - PrepareQueryRequest request = - PrepareQueryRequest.create("SELECT * FROM table", new HashMap<>()); + Map> paramTypes = new HashMap<>(); + paramTypes.put("strParam", SqlType.string()); + paramTypes.put("bytesParam", SqlType.bytes()); + paramTypes.put("intParam", SqlType.int64()); + paramTypes.put("float64Param", SqlType.float64()); + paramTypes.put("float32Param", SqlType.float32()); + paramTypes.put("boolParam", SqlType.bool()); + paramTypes.put("timestampParam", SqlType.timestamp()); + paramTypes.put("dateParam", SqlType.date()); + paramTypes.put("strArrayParam", SqlType.arrayOf(SqlType.string())); + paramTypes.put("byteArrayParam", SqlType.arrayOf(SqlType.bytes())); + paramTypes.put("intArrayParam", SqlType.arrayOf(SqlType.int64())); + paramTypes.put("float32ArrayParam", SqlType.arrayOf(SqlType.float32())); + paramTypes.put("float64ArrayParam", SqlType.arrayOf(SqlType.float64())); + paramTypes.put("boolArrayParam", SqlType.arrayOf(SqlType.bool())); + paramTypes.put("tsArrayParam", SqlType.arrayOf(SqlType.timestamp())); + paramTypes.put("dateArrayParam", SqlType.arrayOf(SqlType.date())); + + PrepareQueryRequest request = PrepareQueryRequest.create("SELECT * FROM table", paramTypes); RequestContext rc = RequestContext.create("project", "instance", "profile"); com.google.bigtable.v2.PrepareQueryRequest proto = request.toProto(rc); @@ -37,5 +65,24 @@ public void testProtoConversion() { assertThat(proto.getAppProfileId()).isEqualTo("profile"); assertThat(proto.getInstanceName()) .isEqualTo(NameUtil.formatInstanceName("project", "instance")); + + Map protoParamTypes = new HashMap<>(); + protoParamTypes.put("strParam", stringType()); + protoParamTypes.put("bytesParam", bytesType()); + protoParamTypes.put("intParam", int64Type()); + protoParamTypes.put("float64Param", float64Type()); + protoParamTypes.put("float32Param", float32Type()); + protoParamTypes.put("boolParam", boolType()); + protoParamTypes.put("timestampParam", timestampType()); + protoParamTypes.put("dateParam", dateType()); + protoParamTypes.put("strArrayParam", arrayType(stringType())); + protoParamTypes.put("byteArrayParam", arrayType(bytesType())); + protoParamTypes.put("intArrayParam", arrayType(int64Type())); + protoParamTypes.put("float32ArrayParam", arrayType(float32Type())); + protoParamTypes.put("float64ArrayParam", arrayType(float64Type())); + protoParamTypes.put("boolArrayParam", arrayType(boolType())); + protoParamTypes.put("tsArrayParam", arrayType(timestampType())); + protoParamTypes.put("dateArrayParam", arrayType(dateType())); + assertThat(proto.getParamTypesMap()).isEqualTo(protoParamTypes); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtilTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtilTest.java new file mode 100644 index 0000000000..b0f1d64a9d --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/QueryParamUtilTest.java @@ -0,0 +1,97 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import static com.google.cloud.bigtable.data.v2.internal.QueryParamUtil.convertToQueryParamProto; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.arrayType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.boolType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.dateType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.float32Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.float64Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.timestampType; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class QueryParamUtilTest { + + @Test + public void convertsSimpleTypes() { + assertThat(convertToQueryParamProto(SqlType.string())).isEqualTo(stringType()); + assertThat(convertToQueryParamProto(SqlType.bytes())).isEqualTo(bytesType()); + assertThat(convertToQueryParamProto(SqlType.int64())).isEqualTo(int64Type()); + assertThat(convertToQueryParamProto(SqlType.float64())).isEqualTo(float64Type()); + assertThat(convertToQueryParamProto(SqlType.float32())).isEqualTo(float32Type()); + assertThat(convertToQueryParamProto(SqlType.bool())).isEqualTo(boolType()); + assertThat(convertToQueryParamProto(SqlType.timestamp())).isEqualTo(timestampType()); + assertThat(convertToQueryParamProto(SqlType.date())).isEqualTo(dateType()); + } + + @Test + public void convertsValidArrayTypes() { + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.string()))) + .isEqualTo(arrayType(stringType())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.bytes()))) + .isEqualTo(arrayType(bytesType())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.int64()))) + .isEqualTo(arrayType(int64Type())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.float64()))) + .isEqualTo(arrayType(float64Type())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.float32()))) + .isEqualTo(arrayType(float32Type())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.bool()))) + .isEqualTo(arrayType(boolType())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.timestamp()))) + .isEqualTo(arrayType(timestampType())); + assertThat(convertToQueryParamProto(SqlType.arrayOf(SqlType.date()))) + .isEqualTo(arrayType(dateType())); + } + + @Test + public void failsForInvalidArrayElementTypes() { + assertThrows( + IllegalArgumentException.class, + () -> convertToQueryParamProto(SqlType.arrayOf(SqlType.struct()))); + assertThrows( + IllegalArgumentException.class, + () -> convertToQueryParamProto(SqlType.arrayOf(SqlType.arrayOf(SqlType.string())))); + assertThrows( + IllegalArgumentException.class, + () -> + convertToQueryParamProto( + SqlType.arrayOf(SqlType.mapOf(SqlType.string(), SqlType.string())))); + } + + @Test + public void failsForMap() { + assertThrows( + IllegalArgumentException.class, + () -> convertToQueryParamProto(SqlType.mapOf(SqlType.string(), SqlType.string()))); + } + + @Test + public void failsForStruct() { + assertThrows(IllegalArgumentException.class, () -> convertToQueryParamProto(SqlType.struct())); + } +} From 692aefa7afa604e6b184ee76ccf4cf72f8ae86f1 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Fri, 31 Jan 2025 17:07:45 -0500 Subject: [PATCH 03/11] Connect execute and prepare This makes a series of changes to use the result of PrepareQuery for execute: - Replace Statement with a BoundStatement based on PreparedStatement - set preparedQuery in request, rather than 'query' - Use the metadata response from Prepare, this will no longer be present in the execute stream - Update Metadata future handling to work well with PrepareQuery refresh in the future Change-Id: Ie3249d0c0bd422b4aaf1c8878660d27288b4b6b0 --- .../bigtable/data/v2/BigtableDataClient.java | 40 ++- .../v2/internal/PreparedStatementImpl.java | 19 +- .../data/v2/internal/SqlRowMergerUtil.java | 5 +- .../{Statement.java => BoundStatement.java} | 51 ++-- .../data/v2/models/sql/PreparedStatement.java | 24 +- .../data/v2/stub/EnhancedBigtableStub.java | 8 +- .../v2/stub/EnhancedBigtableStubSettings.java | 10 +- .../v2/stub/sql/ExecuteQueryCallContext.java | 53 +++- .../v2/stub/sql/ExecuteQueryCallable.java | 7 +- .../stub/sql/MetadataResolvingCallable.java | 47 ++-- .../data/v2/stub/sql/SqlRowMerger.java | 66 ++--- .../v2/stub/sql/SqlRowMergingCallable.java | 6 +- .../AbstractProtoStructReaderTest.java | 35 ++- .../data/v2/internal/ResultSetImplTest.java | 112 ++++---- .../v2/internal/SqlRowMergerUtilTest.java | 57 ++-- .../bigtable/data/v2/it/ExecuteQueryIT.java | 200 ++++++++------ ...ementTest.java => BoundStatementTest.java} | 248 ++++++++---------- .../data/v2/stub/CookiesHolderTest.java | 2 +- .../EnhancedBigtableStubSettingsTest.java | 6 +- .../v2/stub/EnhancedBigtableStubTest.java | 46 ++-- .../bigtable/data/v2/stub/HeadersTest.java | 14 +- .../bigtable/data/v2/stub/RetryInfoTest.java | 2 +- .../stub/sql/ExecuteQueryCallContextTest.java | 89 +++++++ .../v2/stub/sql/ExecuteQueryCallableTest.java | 40 ++- .../sql/MetadataResolvingCallableTest.java | 96 +++---- .../sql/ProtoRowsMergingStateMachineTest.java | 50 ++-- .../data/v2/stub/sql/SqlProtoFactory.java | 20 +- .../data/v2/stub/sql/SqlRowMergerTest.java | 149 ++++++----- .../stub/sql/SqlRowMergingCallableTest.java | 77 ++++-- 29 files changed, 905 insertions(+), 674 deletions(-) rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/{Statement.java => BoundStatement.java} (88%) rename google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/{StatementTest.java => BoundStatementTest.java} (78%) create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 9777eaee33..215cf0305b 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -51,10 +51,10 @@ import com.google.cloud.bigtable.data.v2.models.SampleRowKeysRequest; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream; import com.google.common.util.concurrent.MoreExecutors; @@ -2716,26 +2716,40 @@ public void readChangeStreamAsync( * 

{@code
    * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
    *   String query = "SELECT CAST(cf['stringCol'] AS STRING) FROM [TABLE]";
-   *
-   *   try (ResultSet resultSet = bigtableDataClient.executeQuery(Statement.of(query))) {
-   *     while (resultSet.next()) {
-   *        String s = resultSet.getString("stringCol");
-   *        // do something with data
-   *     }
-   *   } catch (RuntimeException e) {
-   *     e.printStackTrace();
+   *   Map> paramTypes = new HashMap<>();
+   *   try (PreparedStatement preparedStatement = bigtableDataClient.prepareStatement(query, paramTypes)) {
+   *       // Ideally one PreparedStatement should be reused across requests
+   *       BoundStatement boundStatement = preparedStatement.bind()
+   *          // set any query params before calling build
+   *          .build();
+   *       try (ResultSet resultSet = bigtableDataClient.executeQuery(boundStatement)) {
+   *           while (resultSet.next()) {
+   *               String s = resultSet.getString("stringCol");
+   *                // do something with data
+   *           }
+   *        } catch (RuntimeException e) {
+   *            e.printStackTrace();
+   *        }
    *   }
    * }
* - * @see Statement For query options. + * @see {@link PreparedStatement} & {@link BoundStatement} For query options. */ @BetaApi - public ResultSet executeQuery(Statement statement) { - SqlServerStream stream = stub.createExecuteQueryCallable().call(statement); + public ResultSet executeQuery(BoundStatement boundStatement) { + SqlServerStream stream = stub.createExecuteQueryCallable().call(boundStatement); return ResultSetImpl.create(stream); } - // TODO document once BoundStatement exists and executeQuery is updated so usage is clear + /** + * Prepares a query for execution. If possible this should be called once and reused across + * requests. This will amortize the cost of query preparation. + * + * @param query sql query string to prepare + * @param paramTypes a Map of the parameter names and the corresponding {@link SqlType} for all + * query parameters in 'query' + * @return {@link PreparedStatement} which is used to create {@link BoundStatement}s to execute + */ @BetaApi public PreparedStatement prepareStatement(String query, Map> paramTypes) { PrepareQueryRequest request = PrepareQueryRequest.create(query, paramTypes); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java index 67eb41a3a3..46dbad2c3c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java @@ -16,6 +16,8 @@ package com.google.cloud.bigtable.data.v2.internal; import com.google.api.core.InternalApi; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement.Builder; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; /** @@ -23,7 +25,6 @@ * *

This is considered an internal implementation detail and should not be used by applications. */ -// TODO Add ability to create BoundStatement // TODO implement plan refresh @InternalApi("For internal use only") public class PreparedStatementImpl implements PreparedStatement { @@ -36,4 +37,20 @@ public PreparedStatementImpl(PrepareResponse response) { public static PreparedStatement create(PrepareResponse response) { return new PreparedStatementImpl(response); } + + @Override + public BoundStatement.Builder bind() { + return new Builder(this); + } + + // TODO update when plan refresh is implement + @Override + public PrepareResponse getPrepareResponse() { + return response; + } + + @Override + public void close() throws Exception { + // TODO cancel any background refresh + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtil.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtil.java index edb8cf6dcf..90631f3bbd 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtil.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtil.java @@ -18,6 +18,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.bigtable.v2.ExecuteQueryResponse; +import com.google.bigtable.v2.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.stub.sql.SqlRowMerger; import com.google.common.collect.ImmutableList; import java.util.List; @@ -33,8 +34,8 @@ public class SqlRowMergerUtil implements AutoCloseable { private final SqlRowMerger merger; - public SqlRowMergerUtil() { - merger = new SqlRowMerger(); + public SqlRowMergerUtil(ResultSetMetadata metadata) { + merger = new SqlRowMerger(() -> ProtoResultSetMetadata.fromProto(metadata)); } @Override diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java similarity index 88% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java index a2f3639691..ffb8fa9c90 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Statement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java @@ -23,6 +23,7 @@ import com.google.bigtable.v2.Value; import com.google.cloud.Date; import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.QueryParamUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.common.collect.ImmutableMap; @@ -34,9 +35,10 @@ import java.util.Map; import javax.annotation.Nullable; +// TODO update doc /** * A SQL statement that can be executed by calling {@link - * com.google.cloud.bigtable.data.v2.BigtableDataClient#executeQuery(Statement)}. + * com.google.cloud.bigtable.data.v2.BigtableDataClient#executeQuery(BoundStatement)}. * *

A statement contains a SQL string and optional parameters. A parameterized query should * contain placeholders in the form of {@literal @} followed by the parameter name. Parameter names @@ -60,38 +62,47 @@ * } */ @BetaApi -public class Statement { +public class BoundStatement { - private final String sql; + private final PreparedStatement preparedStatement; private final Map params; - private Statement(String sql, Map params) { - this.sql = sql; + private BoundStatement(PreparedStatement preparedStatement, Map params) { + this.preparedStatement = preparedStatement; this.params = params; } - /** Creates a {@code Statement} with the given SQL query and no query parameters. */ - public static Statement of(String sql) { - return newBuilder(sql).build(); - } - - /** Creates a new {@code Builder} with the given SQL query */ - public static Builder newBuilder(String sql) { - return new Builder(sql); + // TODO return a future when plan refresh is implemented + /** + * Get's the most recent version of the PrepareResponse associated with this query. + * + *

This is considered an internal implementation detail and should not be used by applications. + */ + @InternalApi("For internal use only") + public PrepareResponse getLatestPrepareResponse() { + return preparedStatement.getPrepareResponse(); } + // TODO pass in paramTypes, to support validation public static class Builder { - private final String sql; + private final PreparedStatement preparedStatement; private final Map params; - private Builder(String sql) { - this.sql = sql; + /** + * Creates a builder from a {@link PreparedStatement} + * + *

This is considered an internal implementation detail and should not be used by + * applications. + */ + @InternalApi("For internal use only") + public Builder(PreparedStatement preparedStatement) { + this.preparedStatement = preparedStatement; this.params = new HashMap<>(); } /** Builds a {@code Statement} from the builder */ - public Statement build() { - return new Statement(sql, ImmutableMap.copyOf(params)); + public BoundStatement build() { + return new BoundStatement(preparedStatement, ImmutableMap.copyOf(params)); } /** @@ -330,13 +341,13 @@ private static Value.Builder nullValueWithType(Type type) { * not meant to be used by applications. */ @InternalApi("For internal use only") - public ExecuteQueryRequest toProto(RequestContext requestContext) { + public ExecuteQueryRequest toProto(ByteString preparedQuery, RequestContext requestContext) { return ExecuteQueryRequest.newBuilder() .setInstanceName( NameUtil.formatInstanceName( requestContext.getProjectId(), requestContext.getInstanceId())) .setAppProfileId(requestContext.getAppProfileId()) - .setQuery(sql) + .setPreparedQuery(preparedQuery) .putAllParams(params) .build(); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java index 95f0c6ed0f..cbcaba829c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java @@ -16,7 +16,27 @@ package com.google.cloud.bigtable.data.v2.models.sql; import com.google.api.core.BetaApi; +import com.google.api.core.InternalApi; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -// TODO document once you can do something with this +/** + * The results of query preparation that can be used to create {@link BoundStatement}s to execute + * queries. + * + *

Whenever possible this should be shared across different instances the same query, in order to + * amortize query preparation costs. + */ @BetaApi -public interface PreparedStatement {} +public interface PreparedStatement extends AutoCloseable { + + /** + * @return {@link BoundStatement.Builder} to bind query params to and pass to {@link + * com.google.cloud.bigtable.data.v2.BigtableDataClient#executeQuery(BoundStatement)} + */ + BoundStatement.Builder bind(); + + // TODO once refresh is implemented this will return a future so we can + // wait while plan is refreshing + @InternalApi("For internal use only") + PrepareResponse getPrepareResponse(); +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index f8573697a7..966263c1ac 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -92,7 +92,7 @@ import com.google.cloud.bigtable.data.v2.models.SampleRowKeysRequest; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMergingCallable; import com.google.cloud.bigtable.data.v2.stub.changestream.GenerateInitialChangeStreamPartitionsUserCallable; import com.google.cloud.bigtable.data.v2.stub.changestream.ReadChangeStreamResumptionStrategy; @@ -1139,8 +1139,8 @@ private UnaryCallable createReadModifyWriteRowCallable( * Creates a callable chain to handle streaming ExecuteQuery RPCs. The chain will: * *

    - *
  • Convert a {@link Statement} into a {@link ExecuteQueryCallContext}, which passes the - * {@link Statement} & a future for the {@link + *
  • Convert a {@link BoundStatement} into a {@link ExecuteQueryCallContext}, which passes the + * {@link BoundStatement} & a future for the {@link * com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata} up the call chain. *
  • Upon receiving the response stream, it will set the metadata future and translate the * {@link com.google.bigtable.v2.PartialResultSet}s into {@link SqlRow}s @@ -1185,7 +1185,7 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { Callables.watched(withStatsHeaders, watchdogSettings, clientContext); ServerStreamingCallable withMetadataObserver = - new MetadataResolvingCallable(watched); + new MetadataResolvingCallable(watched, requestContext); ServerStreamingCallable merging = new SqlRowMergingCallable(withMetadataObserver); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index f1b8221fc8..218fb682ad 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -47,7 +47,7 @@ import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowMutation; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.stub.metrics.DefaultMetricsProvider; import com.google.cloud.bigtable.data.v2.stub.metrics.MetricsProvider; import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor; @@ -271,7 +271,7 @@ public class EnhancedBigtableStubSettings extends StubSettings readChangeStreamSettings; private final UnaryCallSettings pingAndWarmSettings; - private final ServerStreamingCallSettings executeQuerySettings; + private final ServerStreamingCallSettings executeQuerySettings; private final UnaryCallSettings prepareQuerySettings; private final FeatureFlags featureFlags; @@ -678,7 +678,7 @@ public UnaryCallSettings readModifyWriteRowSettings() { return readChangeStreamSettings; } - public ServerStreamingCallSettings executeQuerySettings() { + public ServerStreamingCallSettings executeQuerySettings() { return executeQuerySettings; } @@ -748,7 +748,7 @@ public static class Builder extends StubSettings.Builder readChangeStreamSettings; private final UnaryCallSettings.Builder pingAndWarmSettings; - private final ServerStreamingCallSettings.Builder executeQuerySettings; + private final ServerStreamingCallSettings.Builder executeQuerySettings; private final UnaryCallSettings.Builder prepareQuerySettings; @@ -1215,7 +1215,7 @@ public UnaryCallSettings.Builder pingAndWarmSettings() * requests will not be retried currently. */ @BetaApi - public ServerStreamingCallSettings.Builder executeQuerySettings() { + public ServerStreamingCallSettings.Builder executeQuerySettings() { return executeQuerySettings; } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java index 8d0e6b81d0..c777a13fc1 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java @@ -17,29 +17,62 @@ import com.google.api.core.InternalApi; import com.google.api.core.SettableApiFuture; -import com.google.auto.value.AutoValue; import com.google.bigtable.v2.ExecuteQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; /** - * POJO used to provide a future to the ExecuteQuery callable chain in order to return metadata to - * users outside of the stream of rows. + * Used to provide a future to the ExecuteQuery callable chain in order to return metadata to users + * outside of the stream of rows. * *

    This should only be constructed by {@link ExecuteQueryCallable} not directly by users. * *

    This is considered an internal implementation detail and should not be used by applications. */ @InternalApi("For internal use only") -@AutoValue -public abstract class ExecuteQueryCallContext { +public class ExecuteQueryCallContext { + + private final BoundStatement boundStatement; + private final SettableApiFuture metadataFuture; + private final PrepareResponse latestPrepareResponse; + // TODO this will be used to track latest resume token here + + private ExecuteQueryCallContext( + BoundStatement boundStatement, SettableApiFuture metadataFuture) { + this.boundStatement = boundStatement; + this.metadataFuture = metadataFuture; + this.latestPrepareResponse = boundStatement.getLatestPrepareResponse(); + } - @InternalApi("For internal use only") public static ExecuteQueryCallContext create( - ExecuteQueryRequest request, SettableApiFuture metadataFuture) { - return new AutoValue_ExecuteQueryCallContext(request, metadataFuture); + BoundStatement boundStatement, SettableApiFuture metadataFuture) { + return new ExecuteQueryCallContext(boundStatement, metadataFuture); } - abstract ExecuteQueryRequest request(); + ExecuteQueryRequest toRequest(RequestContext requestContext) { + return boundStatement.toProto(latestPrepareResponse.preparedQuery(), requestContext); + } - abstract SettableApiFuture resultSetMetadataFuture(); + /** + * Metadata can change as the plan is refreshed. Once a response or complete has been received + * from the stream we know that the {@link com.google.bigtable.v2.PrepareQueryResponse} can no + * longer change, so we can set the metadata. + */ + void firstResponseReceived() { + metadataFuture.set(latestPrepareResponse.resultSetMetadata()); + } + + /** + * If the stream receives an error before receiving any response it needs to be passed through to + * the metadata future + */ + void setMetadataException(Throwable t) { + metadataFuture.setException(t); + } + + SettableApiFuture resultSetMetadataFuture() { + return this.metadataFuture; + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java index 9563b6c6f9..534964a97c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java @@ -24,8 +24,8 @@ import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; /** * Callable that creates {@link SqlServerStream}s from {@link ExecuteQueryRequest}s. @@ -48,11 +48,10 @@ public ExecuteQueryCallable( this.requestContext = requestContext; } - public SqlServerStream call(Statement statement) { - ExecuteQueryRequest request = statement.toProto(requestContext); + public SqlServerStream call(BoundStatement boundStatement) { SettableApiFuture metadataFuture = SettableApiFuture.create(); ServerStream rowStream = - this.call(ExecuteQueryCallContext.create(request, metadataFuture)); + this.call(ExecuteQueryCallContext.create(boundStatement, metadataFuture)); return SqlServerStreamImpl.create(metadataFuture, rowStream); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java index 6b2f2b171f..0ceacdecf0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java @@ -16,14 +16,13 @@ package com.google.cloud.bigtable.data.v2.stub.sql; import com.google.api.core.InternalApi; -import com.google.api.core.SettableApiFuture; import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.StreamController; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; -import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.stub.SafeResponseObserver; @@ -37,10 +36,13 @@ public class MetadataResolvingCallable extends ServerStreamingCallable { private final ServerStreamingCallable inner; + private final RequestContext requestContext; public MetadataResolvingCallable( - ServerStreamingCallable inner) { + ServerStreamingCallable inner, + RequestContext requestContext) { this.inner = inner; + this.requestContext = requestContext; } @Override @@ -48,25 +50,23 @@ public void call( ExecuteQueryCallContext callContext, ResponseObserver responseObserver, ApiCallContext apiCallContext) { - MetadataObserver observer = - new MetadataObserver(responseObserver, callContext.resultSetMetadataFuture()); - inner.call(callContext.request(), observer, apiCallContext); + MetadataObserver observer = new MetadataObserver(responseObserver, callContext); + inner.call(callContext.toRequest(requestContext), observer, apiCallContext); } static final class MetadataObserver extends SafeResponseObserver { - private final SettableApiFuture metadataFuture; + private final ExecuteQueryCallContext callContext; private final ResponseObserver outerObserver; // This doesn't need to be synchronized because this is called above the reframer // so onResponse will be called sequentially private boolean isFirstResponse; MetadataObserver( - ResponseObserver outerObserver, - SettableApiFuture metadataFuture) { + ResponseObserver outerObserver, ExecuteQueryCallContext callContext) { super(outerObserver); this.outerObserver = outerObserver; - this.metadataFuture = metadataFuture; + this.callContext = callContext; this.isFirstResponse = true; } @@ -77,22 +77,10 @@ protected void onStartImpl(StreamController streamController) { @Override protected void onResponseImpl(ExecuteQueryResponse response) { - if (isFirstResponse && !response.hasMetadata()) { - IllegalStateException e = - new IllegalStateException("First response must always contain metadata"); - metadataFuture.setException(e); - throw e; + if (isFirstResponse) { + callContext.firstResponseReceived(); } isFirstResponse = false; - if (response.hasMetadata()) { - try { - ResultSetMetadata md = ProtoResultSetMetadata.fromProto(response.getMetadata()); - metadataFuture.set(md); - } catch (Throwable t) { - metadataFuture.setException(t); - throw t; - } - } outerObserver.onResponse(response); } @@ -100,17 +88,16 @@ protected void onResponseImpl(ExecuteQueryResponse response) { protected void onErrorImpl(Throwable throwable) { // When we support retries this will have to move after the retrying callable in a separate // observer. - metadataFuture.setException(throwable); + callContext.setMetadataException(throwable); outerObserver.onError(throwable); } + // TODO this becomes a valid state @Override protected void onCompleteImpl() { - if (!metadataFuture.isDone()) { - IllegalStateException missingMetadataException = - new IllegalStateException("Unexpected Stream complete without receiving metadata"); - metadataFuture.setException(missingMetadataException); - throw missingMetadataException; + if (isFirstResponse) { + // If the stream completes successfully we know we used the current metadata + callContext.firstResponseReceived(); } outerObserver.onComplete(); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java index 6178a1efcd..c44ead34f3 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java @@ -18,13 +18,13 @@ import com.google.api.core.InternalApi; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.PartialResultSet; -import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.gaxx.reframing.Reframer; import com.google.common.base.Preconditions; import java.util.ArrayDeque; import java.util.Queue; +import java.util.function.Supplier; /** * Used to transform a stream of ExecuteQueryResponse objects into rows. This class is not thread @@ -33,18 +33,20 @@ @InternalApi public final class SqlRowMerger implements Reframer { - enum State { - AWAITING_METADATA, - PROCESSING_DATA, - } - private final Queue queue; private ProtoRowsMergingStateMachine stateMachine; - private State currentState; + private final Supplier metadataSupplier; + private Boolean isFirstResponse; - public SqlRowMerger() { + /** + * @param metadataSupplier a supplier of {@link ResultSetMetadata}. This is expected to return + * successfully once the first call to push has been made. + *

    This exists to facilitate plan refresh that can happen after creation of the row merger. + */ + public SqlRowMerger(Supplier metadataSupplier) { + this.metadataSupplier = metadataSupplier; queue = new ArrayDeque<>(); - currentState = State.AWAITING_METADATA; + isFirstResponse = true; } /** @@ -52,33 +54,21 @@ public SqlRowMerger() { * * @param response the next response in the stream of query responses */ - // Suppress this because it won't be forced to be exhaustive once it is open-sourced, so we want a - // default. - @SuppressWarnings("UnnecessaryDefaultInEnumSwitch") @Override public void push(ExecuteQueryResponse response) { - switch (currentState) { - case AWAITING_METADATA: - Preconditions.checkState( - response.hasMetadata(), - "Expected metadata response, but received: %s", - response.getResponseCase().name()); - ResultSetMetadata responseMetadata = - ProtoResultSetMetadata.fromProto(response.getMetadata()); - stateMachine = new ProtoRowsMergingStateMachine(responseMetadata); - currentState = State.PROCESSING_DATA; - break; - case PROCESSING_DATA: - Preconditions.checkState( - response.hasResults(), - "Expected results response, but received: %s", - response.getResponseCase().name()); - PartialResultSet results = response.getResults(); - processProtoRows(results); - break; - default: - throw new IllegalStateException("Unknown State: " + currentState.name()); + if (isFirstResponse) { + // Wait until we've received the first response to get the metadata, as a + // PreparedQuery may need to be refreshed based on initial errors. Once we've + // received a response, it will never change, even upon request resumption. + stateMachine = new ProtoRowsMergingStateMachine(metadataSupplier.get()); + isFirstResponse = false; } + Preconditions.checkState( + response.hasResults(), + "Expected results response, but received: %s", + response.getResponseCase().name()); + PartialResultSet results = response.getResults(); + processProtoRows(results); } private void processProtoRows(PartialResultSet results) { @@ -105,14 +95,10 @@ public boolean hasFullFrame() { */ @Override public boolean hasPartialFrame() { - switch (currentState) { - case AWAITING_METADATA: - return false; - case PROCESSING_DATA: - return hasFullFrame() || stateMachine.isBatchInProgress(); - default: - throw new IllegalStateException("Unknown State: " + currentState.name()); + if (isFirstResponse) { + return false; } + return hasFullFrame() || stateMachine.isBatchInProgress(); } /** pops a completed row from the FIFO queue built from the given responses. */ diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallable.java index 6d5d0ea4a4..c788fe9230 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallable.java @@ -17,6 +17,7 @@ import com.google.api.core.InternalApi; import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ApiExceptions; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.bigtable.v2.ExecuteQueryResponse; @@ -38,7 +39,10 @@ public void call( ExecuteQueryCallContext callContext, ResponseObserver responseObserver, ApiCallContext apiCallContext) { - SqlRowMerger merger = new SqlRowMerger(); + SqlRowMerger merger = + new SqlRowMerger( + () -> + ApiExceptions.callAndTranslateApiException(callContext.resultSetMetadataFuture())); ReframingResponseObserver observer = new ReframingResponseObserver<>(responseObserver, merger); inner.call(callContext, observer, apiCallContext); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReaderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReaderTest.java index 8770880983..967c8c2dbb 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReaderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/AbstractProtoStructReaderTest.java @@ -107,8 +107,7 @@ public void simpleMapField_validatesType() { TestProtoStruct structWithMap = TestProtoStruct.create( ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("testField", mapType(bytesType(), stringType()))) - .getMetadata()), + metadata(columnMetadata("testField", mapType(bytesType(), stringType())))), Collections.singletonList( mapValue( mapElement(bytesValue("foo"), stringValue("bar")), @@ -143,15 +142,14 @@ public void nestedMapField_validatesType() { TestProtoStruct.create( ProtoResultSetMetadata.fromProto( metadata( - columnMetadata( - "testField", - mapType( - bytesType(), - arrayType( - structType( - structField("timestamp", timestampType()), - structField("value", bytesType())))))) - .getMetadata()), + columnMetadata( + "testField", + mapType( + bytesType(), + arrayType( + structType( + structField("timestamp", timestampType()), + structField("value", bytesType()))))))), Collections.singletonList( mapValue( mapElement( @@ -205,7 +203,7 @@ public void arrayField_validatesType() { TestProtoStruct structWithList = TestProtoStruct.create( ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("testField", arrayType(stringType()))).getMetadata()), + metadata(columnMetadata("testField", arrayType(stringType())))), Collections.singletonList(arrayValue(stringValue("foo"), stringValue("bar")))); List expectedList = Arrays.asList("foo", "bar"); @@ -229,7 +227,7 @@ public void arrayField_accessingFloat() { TestProtoStruct structWithList = TestProtoStruct.create( ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("testField", arrayType(float32Type()))).getMetadata()), + metadata(columnMetadata("testField", arrayType(float32Type())))), Collections.singletonList(arrayValue(floatValue(1.1f), floatValue(1.2f)))); List floatList = @@ -565,8 +563,7 @@ public static List parameters() { private TestProtoStruct getTestRow() { return TestProtoStruct.create( - ProtoResultSetMetadata.fromProto( - metadata(schema.toArray(new ColumnMetadata[] {})).getMetadata()), + ProtoResultSetMetadata.fromProto(metadata(schema.toArray(new ColumnMetadata[] {}))), values); } @@ -632,7 +629,7 @@ public void getByColumnIndex_throwsExceptionOnWrongType() { TestProtoStruct row = TestProtoStruct.create( ProtoResultSetMetadata.fromProto( - metadata(updatedSchema.toArray(new ColumnMetadata[] {})).getMetadata()), + metadata(updatedSchema.toArray(new ColumnMetadata[] {}))), updatedValues); assertThrows(IllegalStateException.class, () -> getByIndex.apply(row, index)); @@ -654,7 +651,7 @@ public void getByColumnName_throwsExceptionOnWrongType() { TestProtoStruct row = TestProtoStruct.create( ProtoResultSetMetadata.fromProto( - metadata(updatedSchema.toArray(new ColumnMetadata[] {})).getMetadata()), + metadata(updatedSchema.toArray(new ColumnMetadata[] {}))), updatedValues); assertThrows(IllegalStateException.class, () -> getByColumn.apply(row, columnName)); @@ -693,7 +690,7 @@ public void getByColumnName_throwsExceptionForDuplicateColumnName() { duplicatedSchema.addAll(schema); ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( - metadata(duplicatedSchema.toArray(new ColumnMetadata[] {})).getMetadata()); + metadata(duplicatedSchema.toArray(new ColumnMetadata[] {}))); List duplicatedValues = new ArrayList<>(values); duplicatedValues.addAll(values); TestProtoStruct row = TestProtoStruct.create(metadata, duplicatedValues); @@ -708,7 +705,7 @@ public void getByIndex_worksWithDuplicateColumnName() { duplicatedSchema.addAll(schema); ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( - metadata(duplicatedSchema.toArray(new ColumnMetadata[] {})).getMetadata()); + metadata(duplicatedSchema.toArray(new ColumnMetadata[] {}))); List duplicatedValues = new ArrayList<>(values); duplicatedValues.addAll(values); TestProtoStruct row = TestProtoStruct.create(metadata, duplicatedValues); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java index a8c5776a87..b4f1515923 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java @@ -33,6 +33,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.structField; @@ -46,6 +47,8 @@ import com.google.api.core.SettableApiFuture; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.cloud.Date; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; @@ -67,36 +70,39 @@ @RunWith(JUnit4.class) public class ResultSetImplTest { - private static ResultSet resultSetWithFakeStream(ResultSetMetadata metadata, SqlRow... rows) { + private static ResultSet resultSetWithFakeStream( + com.google.bigtable.v2.ResultSetMetadata protoMetadata, SqlRow... rows) { ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>(Arrays.asList(rows)); SettableApiFuture future = SettableApiFuture.create(); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); future.set(metadata); + PrepareResponse response = PrepareResponse.fromProto(prepareResponse(protoMetadata)); + PreparedStatement preparedStatement = PreparedStatementImpl.create(response); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create(ExecuteQueryRequest.newBuilder().build(), future); + ExecuteQueryCallContext.create(preparedStatement.bind().build(), future); return ResultSetImpl.create(SqlServerStreamImpl.create(future, stream.call(fakeCallContext))); } @Test public void testSingleRow() throws ExecutionException, InterruptedException { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata( - columnMetadata("string", stringType()), - columnMetadata("bytes", bytesType()), - columnMetadata("long", int64Type()), - columnMetadata("double", float64Type()), - columnMetadata("float", float32Type()), - columnMetadata("boolean", boolType()), - columnMetadata("timestamp", timestampType()), - columnMetadata("date", dateType()), - columnMetadata("struct", structType(structField("string", stringType()))), - columnMetadata("list", arrayType(stringType())), - columnMetadata("map", mapType(stringType(), stringType()))) - .getMetadata()); + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata( + columnMetadata("string", stringType()), + columnMetadata("bytes", bytesType()), + columnMetadata("long", int64Type()), + columnMetadata("double", float64Type()), + columnMetadata("float", float32Type()), + columnMetadata("boolean", boolType()), + columnMetadata("timestamp", timestampType()), + columnMetadata("date", dateType()), + columnMetadata("struct", structType(structField("string", stringType()))), + columnMetadata("list", arrayType(stringType())), + columnMetadata("map", mapType(stringType(), stringType()))); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); ResultSet resultSet = resultSetWithFakeStream( - metadata, + protoMetadata, ProtoSqlRow.create( metadata, Arrays.asList( @@ -170,12 +176,12 @@ public void testSingleRow() throws ExecutionException, InterruptedException { @Test public void testIteration() { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata(columnMetadata("string", stringType())); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); try (ResultSet resultSet = resultSetWithFakeStream( - metadata, + protoMetadata, ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("foo"))), ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("bar"))), ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("baz"))), @@ -197,11 +203,11 @@ public void testIteration() { } @Test - public void testEmptyResultSet() throws ExecutionException, InterruptedException { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); - try (ResultSet resultSet = resultSetWithFakeStream(metadata)) { + public void testEmptyResultSet() { + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata(columnMetadata("string", stringType())); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); + try (ResultSet resultSet = resultSetWithFakeStream(protoMetadata)) { assertThat(resultSet.next()).isFalse(); assertThat(resultSet.getMetadata()).isEqualTo(metadata); } @@ -209,13 +215,13 @@ public void testEmptyResultSet() throws ExecutionException, InterruptedException @Test public void getCallsPrevented_afterNextReturnsFalse() { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata(columnMetadata("string", stringType())); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); ResultSet resultSet = resultSetWithFakeStream( - metadata, + protoMetadata, ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("foo"))), ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("bar")))); @@ -233,12 +239,13 @@ public void getCallsPrevented_afterNextReturnsFalse() { @Test public void close_preventsGetCalls() { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata(columnMetadata("string", stringType())); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); ResultSet resultSet = resultSetWithFakeStream( - metadata, ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("foo")))); + protoMetadata, + ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("foo")))); assertThat(resultSet.next()).isTrue(); resultSet.close(); @@ -248,8 +255,7 @@ public void close_preventsGetCalls() { @Test public void close_cancelsStreamWhenResultsNotConsumed() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("string", stringType()))); ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>( Collections.singletonList( @@ -267,8 +273,7 @@ public void close_cancelsStreamWhenResultsNotConsumed() { @Test public void close_doesNotCancelStreamWhenResultsConsumed() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("string", stringType()))); ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>( Collections.singletonList( @@ -287,12 +292,12 @@ public void close_doesNotCancelStreamWhenResultsConsumed() { @Test public void getBeforeNext_throwsException() { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("string", stringType())).getMetadata()); + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata(columnMetadata("string", stringType())); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); try (ResultSet resultSet = resultSetWithFakeStream( - metadata, + protoMetadata, ProtoSqlRow.create(metadata, Collections.singletonList(stringValue("foo"))))) { assertThrows(IllegalStateException.class, () -> resultSet.getString(0)); } @@ -300,13 +305,12 @@ public void getBeforeNext_throwsException() { @Test public void getOnColumnWithDuplicateName_throwsException() { - ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("name", stringType()), columnMetadata("name", stringType())) - .getMetadata()); + com.google.bigtable.v2.ResultSetMetadata protoMetadata = + metadata(columnMetadata("string", stringType())); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); try (ResultSet resultSet = resultSetWithFakeStream( - metadata, + protoMetadata, ProtoSqlRow.create(metadata, Arrays.asList(stringValue("foo"), stringValue("bar"))))) { assertThat(resultSet.next()).isTrue(); @@ -319,8 +323,12 @@ public void getMetadata_unwrapsExecutionExceptions() { SettableApiFuture metadataFuture = SettableApiFuture.create(); ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>(Collections.emptyList()); + PrepareResponse prepareResponse = + PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))); + PreparedStatement preparedStatement = PreparedStatementImpl.create(prepareResponse); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create(ExecuteQueryRequest.newBuilder().build(), metadataFuture); + ExecuteQueryCallContext.create( + new BoundStatement.Builder(preparedStatement).build(), metadataFuture); ResultSet rs = ResultSetImpl.create( SqlServerStreamImpl.create(metadataFuture, stream.call(fakeCallContext))); @@ -334,8 +342,12 @@ public void getMetadata_returnsNonRuntimeExecutionExceptionsWrapped() { SettableApiFuture metadataFuture = SettableApiFuture.create(); ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>(Collections.emptyList()); + PrepareResponse prepareResponse = + PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))); + PreparedStatement preparedStatement = PreparedStatementImpl.create(prepareResponse); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create(ExecuteQueryRequest.newBuilder().build(), metadataFuture); + ExecuteQueryCallContext.create( + new BoundStatement.Builder(preparedStatement).build(), metadataFuture); ResultSet rs = ResultSetImpl.create( SqlServerStreamImpl.create(metadataFuture, stream.call(fakeCallContext))); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java index 6ed96ec517..9986476abb 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java @@ -46,40 +46,27 @@ public class SqlRowMergerUtilTest { @Test public void close_succeedsWhenEmpty() { - try (SqlRowMergerUtil util = new SqlRowMergerUtil()) {} + com.google.bigtable.v2.ResultSetMetadata md = metadata(columnMetadata("a", stringType())); + try (SqlRowMergerUtil util = new SqlRowMergerUtil(md)) {} - try (SqlRowMergerUtil util = new SqlRowMergerUtil()) { + try (SqlRowMergerUtil util = new SqlRowMergerUtil(md)) { // Metadata with no rows - List unused = - util.parseExecuteQueryResponses( - ImmutableList.of(metadata(columnMetadata("a", stringType())))); - } - } - - @Test - public void parseExecuteQueryResponses_failsWithoutMetadata_serializedProtoRows() { - try (SqlRowMergerUtil util = new SqlRowMergerUtil()) { - // users must pass metadata, as it should always be returned by the server. - assertThrows( - IllegalStateException.class, - () -> - util.parseExecuteQueryResponses( - ImmutableList.of(partialResultSetWithToken(stringValue("val"))))); + List unused = util.parseExecuteQueryResponses(ImmutableList.of()); } } @Test public void parseExecuteQueryResponses_handlesSingleValue_serializedProtoRows() { - ExecuteQueryResponse metadata = metadata(columnMetadata("str", stringType())); + com.google.bigtable.v2.ResultSetMetadata metadata = + metadata(columnMetadata("str", stringType())); ImmutableList responses = - ImmutableList.of(metadata, partialResultSetWithToken(stringValue("val"))); - try (SqlRowMergerUtil util = new SqlRowMergerUtil()) { + ImmutableList.of(partialResultSetWithToken(stringValue("val"))); + try (SqlRowMergerUtil util = new SqlRowMergerUtil(metadata)) { List rows = util.parseExecuteQueryResponses(responses); assertThat(rows) .containsExactly( ProtoSqlRow.create( - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("str", stringType())).getMetadata()), + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("str", stringType()))), ImmutableList.of(stringValue("val")))); ; } @@ -94,10 +81,10 @@ public void parseExecuteQueryResponses_handlesSingleValue_serializedProtoRows() columnMetadata("strArr", arrayType(stringType())), columnMetadata("map", mapType(stringType(), bytesType())) }; - ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columns).getMetadata()); + com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columns); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); ImmutableList responses = ImmutableList.of( - metadata(columns), partialResultSetWithoutToken( stringValue("str1"), bytesValue("bytes1"), @@ -113,7 +100,7 @@ public void parseExecuteQueryResponses_handlesSingleValue_serializedProtoRows() bytesValue("bytes3"), arrayValue(stringValue("arr3")), mapValue(mapElement(stringValue("key3"), bytesValue("val3"))))); - try (SqlRowMergerUtil util = new SqlRowMergerUtil()) { + try (SqlRowMergerUtil util = new SqlRowMergerUtil(metadataProto)) { List rows = util.parseExecuteQueryResponses(responses); assertThat(rows) .containsExactly( @@ -143,12 +130,12 @@ public void parseExecuteQueryResponses_handlesSingleValue_serializedProtoRows() @Test public void parseExecuteQueryResponses_throwsOnCloseWithPartialBatch_serializedProtoRows() { + com.google.bigtable.v2.ResultSetMetadata metadata = + metadata(columnMetadata("str", stringType())); ImmutableList responses = - ImmutableList.of( - metadata(columnMetadata("str", stringType())), - partialResultSetWithoutToken(stringValue("str1"))); + ImmutableList.of(partialResultSetWithoutToken(stringValue("str1"))); - SqlRowMergerUtil util = new SqlRowMergerUtil(); + SqlRowMergerUtil util = new SqlRowMergerUtil(metadata); List unused = util.parseExecuteQueryResponses(responses); assertThrows(IllegalStateException.class, util::close); } @@ -156,13 +143,14 @@ public void parseExecuteQueryResponses_throwsOnCloseWithPartialBatch_serializedP @Test public void parseExecuteQueryResponses_throwsOnParseWithPartialRowsInCompleteBatch_serializedProtoRows() { + com.google.bigtable.v2.ResultSetMetadata metadata = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); ImmutableList responses = ImmutableList.of( - metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())), partialResultSetWithToken( stringValue("str1"), bytesValue("bytes1"), stringValue("str2"))); - SqlRowMergerUtil util = new SqlRowMergerUtil(); + SqlRowMergerUtil util = new SqlRowMergerUtil(metadata); assertThrows(IllegalStateException.class, () -> util.parseExecuteQueryResponses(responses)); } @@ -174,10 +162,10 @@ public void parseExecuteQueryResponses_worksWithIncrementalSetsOfResponses_seria columnMetadata("strArr", arrayType(stringType())), columnMetadata("map", mapType(stringType(), bytesType())) }; - ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columns).getMetadata()); + com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columns); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); ImmutableList responses = ImmutableList.of( - metadata(columns), partialResultSetWithoutToken( stringValue("str1"), bytesValue("bytes1"), @@ -193,12 +181,11 @@ public void parseExecuteQueryResponses_worksWithIncrementalSetsOfResponses_seria bytesValue("bytes3"), arrayValue(stringValue("arr3")), mapValue(mapElement(stringValue("key3"), bytesValue("val3"))))); - try (SqlRowMergerUtil util = new SqlRowMergerUtil()) { + try (SqlRowMergerUtil util = new SqlRowMergerUtil(metadataProto)) { List rows = new ArrayList<>(); rows.addAll(util.parseExecuteQueryResponses(responses.subList(0, 1))); rows.addAll(util.parseExecuteQueryResponses(responses.subList(1, 2))); rows.addAll(util.parseExecuteQueryResponses(responses.subList(2, 3))); - rows.addAll(util.parseExecuteQueryResponses(responses.subList(3, 4))); assertThat(rows) .containsExactly( diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/ExecuteQueryIT.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/ExecuteQueryIT.java index 34d0952401..a72c397500 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/ExecuteQueryIT.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/ExecuteQueryIT.java @@ -23,9 +23,10 @@ import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.cloud.bigtable.data.v2.models.sql.Struct; import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv; import com.google.cloud.bigtable.test_helpers.env.TestEnvRule; @@ -83,10 +84,12 @@ public static void setUpAll() throws IOException { @Test public void selectStar() { - try (ResultSet rs = - dataClient.executeQuery( - Statement.of( - "SELECT * FROM " + tableId + " WHERE _key LIKE '" + uniquePrefix + "%'"))) { + PreparedStatement preparedStatement = + dataClient.prepareStatement( + "SELECT * FROM " + tableId + " WHERE _key LIKE '" + uniquePrefix + "%'", + new HashMap<>()); + BoundStatement statement = preparedStatement.bind().build(); + try (ResultSet rs = dataClient.executeQuery(statement)) { assertThat(rs.next()).isTrue(); assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "a")); assertThat( @@ -107,15 +110,16 @@ public void selectStar() { @Test public void withHistoryQuery() { - try (ResultSet rs = - dataClient.executeQuery( - Statement.of( - "SELECT * FROM `" - + tableId - + "`(with_history => true) WHERE _key LIKE '" - + uniquePrefix - + "%'"))) { - + PreparedStatement preparedStatement = + dataClient.prepareStatement( + "SELECT * FROM `" + + tableId + + "`(with_history => true) WHERE _key LIKE '" + + uniquePrefix + + "%'", + new HashMap<>()); + BoundStatement statement = preparedStatement.bind().build(); + try (ResultSet rs = dataClient.executeQuery(statement)) { assertThat(rs.next()).isTrue(); assertThat(rs.getBytes("_key")).isEqualTo(ByteString.copyFromUtf8(uniquePrefix + "a")); Map> rowACf = rs.getMap(cf, SqlType.historicalMap()); @@ -143,19 +147,20 @@ public void withHistoryQuery() { @Test public void allTypes() { - try (ResultSet rs = - dataClient.executeQuery( - Statement.of( - "SELECT 'stringVal' AS strCol, b'foo' as bytesCol, 1 AS intCol, CAST(1.2 AS FLOAT32) as f32Col, " - + "CAST(1.3 AS FLOAT64) as f64Col, true as boolCol, TIMESTAMP_FROM_UNIX_MILLIS(1000) AS tsCol, " - + "DATE(2024, 06, 01) as dateCol, STRUCT(1 as a, \"foo\" as b) AS structCol, [1,2,3] AS arrCol, " - + cf - + " as mapCol FROM `" - + tableId - + "` WHERE _key='" - + uniquePrefix - + "a' LIMIT 1"))) { - + PreparedStatement preparedStatement = + dataClient.prepareStatement( + "SELECT 'stringVal' AS strCol, b'foo' as bytesCol, 1 AS intCol, CAST(1.2 AS FLOAT32) as f32Col, " + + "CAST(1.3 AS FLOAT64) as f64Col, true as boolCol, TIMESTAMP_FROM_UNIX_MILLIS(1000) AS tsCol, " + + "DATE(2024, 06, 01) as dateCol, STRUCT(1 as a, \"foo\" as b) AS structCol, [1,2,3] AS arrCol, " + + cf + + " as mapCol FROM `" + + tableId + + "` WHERE _key='" + + uniquePrefix + + "a' LIMIT 1", + new HashMap<>()); + BoundStatement statement = preparedStatement.bind().build(); + try (ResultSet rs = dataClient.executeQuery(statement)) { assertThat(rs.next()).isTrue(); assertThat(rs.getString("strCol")).isEqualTo("stringVal"); assertThat(rs.getString(0)).isEqualTo("stringVal"); @@ -206,59 +211,78 @@ public void allTypes() { @Test public void allQueryParamsTypes() { - ResultSet rs = - dataClient.executeQuery( - Statement.newBuilder( - "SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, " - + "@doubleParam AS doubleCol, @floatParam AS floatCol, @boolParam AS boolCol, " - + "@tsParam AS tsCol, @dateParam AS dateCol, @byteArrayParam AS byteArrayCol, " - + "@stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol, " - + "@floatArrayParam AS floatArrayCol, @doubleArrayParam AS doubleArrayCol, " - + "@boolArrayParam AS boolArrayCol, @tsArrayParam AS tsArrayCol, " - + "@dateArrayParam AS dateArrayCol") - .setStringParam("stringParam", "stringVal") - .setBytesParam("bytesParam", ByteString.copyFromUtf8("foo")) - .setLongParam("int64Param", 1L) - .setDoubleParam("doubleParam", 1.3d) - .setFloatParam("floatParam", 1.4f) - .setBooleanParam("boolParam", true) - .setTimestampParam("tsParam", Instant.ofEpochMilli(1000)) - .setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 1)) - .setListParam( - "byteArrayParam", - Arrays.asList( - ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")), - SqlType.arrayOf(SqlType.bytes())) - .setListParam( - "stringArrayParam", - Arrays.asList("foo", null, "bar"), - SqlType.arrayOf(SqlType.string())) - .setListParam( - "intArrayParam", Arrays.asList(1L, null, 2L), SqlType.arrayOf(SqlType.int64())) - .setListParam( - "floatArrayParam", - Arrays.asList(1.2f, null, 1.3f), - SqlType.arrayOf(SqlType.float32())) - .setListParam( - "doubleArrayParam", - Arrays.asList(1.4d, null, 1.5d), - SqlType.arrayOf(SqlType.float64())) - .setListParam( - "boolArrayParam", - Arrays.asList(true, null, false), - SqlType.arrayOf(SqlType.bool())) - .setListParam( - "tsArrayParam", - Arrays.asList( - Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)), - SqlType.arrayOf(SqlType.timestamp())) - .setListParam( - "dateArrayParam", - Arrays.asList( - Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)), - SqlType.arrayOf(SqlType.date())) - .build()); + Map> paramTypes = new HashMap<>(); + paramTypes.put("stringParam", SqlType.string()); + paramTypes.put("bytesParam", SqlType.bytes()); + paramTypes.put("int64Param", SqlType.int64()); + paramTypes.put("doubleParam", SqlType.float64()); + paramTypes.put("floatParam", SqlType.float32()); + paramTypes.put("boolParam", SqlType.bool()); + paramTypes.put("tsParam", SqlType.timestamp()); + paramTypes.put("dateParam", SqlType.date()); + paramTypes.put("stringArrayParam", SqlType.arrayOf(SqlType.string())); + paramTypes.put("byteArrayParam", SqlType.arrayOf(SqlType.bytes())); + paramTypes.put("intArrayParam", SqlType.arrayOf(SqlType.int64())); + paramTypes.put("doubleArrayParam", SqlType.arrayOf(SqlType.float64())); + paramTypes.put("floatArrayParam", SqlType.arrayOf(SqlType.float32())); + paramTypes.put("boolArrayParam", SqlType.arrayOf(SqlType.bool())); + paramTypes.put("tsArrayParam", SqlType.arrayOf(SqlType.timestamp())); + paramTypes.put("dateArrayParam", SqlType.arrayOf(SqlType.date())); + + PreparedStatement preparedStatement = + dataClient.prepareStatement( + "SELECT @stringParam AS strCol, @bytesParam as bytesCol, @int64Param AS intCol, " + + "@doubleParam AS doubleCol, @floatParam AS floatCol, @boolParam AS boolCol, " + + "@tsParam AS tsCol, @dateParam AS dateCol, @byteArrayParam AS byteArrayCol, " + + "@stringArrayParam AS stringArrayCol, @intArrayParam AS intArrayCol, " + + "@floatArrayParam AS floatArrayCol, @doubleArrayParam AS doubleArrayCol, " + + "@boolArrayParam AS boolArrayCol, @tsArrayParam AS tsArrayCol, " + + "@dateArrayParam AS dateArrayCol", + paramTypes); + BoundStatement boundStatement = + preparedStatement + .bind() + .setStringParam("stringParam", "stringVal") + .setBytesParam("bytesParam", ByteString.copyFromUtf8("foo")) + .setLongParam("int64Param", 1L) + .setDoubleParam("doubleParam", 1.3d) + .setFloatParam("floatParam", 1.4f) + .setBooleanParam("boolParam", true) + .setTimestampParam("tsParam", Instant.ofEpochMilli(1000)) + .setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 1)) + .setListParam( + "byteArrayParam", + Arrays.asList(ByteString.copyFromUtf8("foo"), null, ByteString.copyFromUtf8("bar")), + SqlType.arrayOf(SqlType.bytes())) + .setListParam( + "stringArrayParam", + Arrays.asList("foo", null, "bar"), + SqlType.arrayOf(SqlType.string())) + .setListParam( + "intArrayParam", Arrays.asList(1L, null, 2L), SqlType.arrayOf(SqlType.int64())) + .setListParam( + "floatArrayParam", + Arrays.asList(1.2f, null, 1.3f), + SqlType.arrayOf(SqlType.float32())) + .setListParam( + "doubleArrayParam", + Arrays.asList(1.4d, null, 1.5d), + SqlType.arrayOf(SqlType.float64())) + .setListParam( + "boolArrayParam", Arrays.asList(true, null, false), SqlType.arrayOf(SqlType.bool())) + .setListParam( + "tsArrayParam", + Arrays.asList( + Instant.ofEpochSecond(1000, 1000), null, Instant.ofEpochSecond(2000, 2000)), + SqlType.arrayOf(SqlType.timestamp())) + .setListParam( + "dateArrayParam", + Arrays.asList( + Date.fromYearMonthDay(2024, 8, 1), null, Date.fromYearMonthDay(2024, 8, 2)), + SqlType.arrayOf(SqlType.date())) + .build(); + ResultSet rs = dataClient.executeQuery(boundStatement); assertThat(rs.next()).isTrue(); assertThat(rs.getString("strCol")).isEqualTo("stringVal"); assertThat(rs.getString(0)).isEqualTo("stringVal"); @@ -322,14 +346,16 @@ public void allQueryParamsTypes() { @Test public void testNullColumns() { - try (ResultSet rs = - dataClient.executeQuery( - Statement.of( - "SELECT cf['qual'] AS neverNull, cf['qual3'] AS maybeNull FROM " - + tableId - + " WHERE _key LIKE '" - + uniquePrefix - + "%'"))) { + PreparedStatement preparedStatement = + dataClient.prepareStatement( + "SELECT cf['qual'] AS neverNull, cf['qual3'] AS maybeNull FROM " + + tableId + + " WHERE _key LIKE '" + + uniquePrefix + + "%'", + new HashMap<>()); + BoundStatement statement = preparedStatement.bind().build(); + try (ResultSet rs = dataClient.executeQuery(statement)) { assertThat(rs.next()).isTrue(); assertThat(rs.getBytes("neverNull")).isEqualTo(ByteString.copyFromUtf8("val")); // qual3 is set in row A but not row B diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/StatementTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java similarity index 78% rename from google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/StatementTest.java rename to google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java index 6d4765230e..a27a10ef05 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/StatementTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java @@ -21,6 +21,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.boolValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.dateType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.dateValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.float32Type; @@ -28,6 +29,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.floatValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Value; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.nullValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; @@ -36,11 +38,17 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.bigtable.v2.ColumnMetadata; import com.google.bigtable.v2.ExecuteQueryRequest; +import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.Value; import com.google.cloud.Date; +import com.google.cloud.Timestamp; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.protobuf.ByteString; +import java.time.Duration; import java.time.Instant; import java.util.Arrays; import java.util.Collections; @@ -49,22 +57,43 @@ import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public class StatementTest { +public class BoundStatementTest { private static final String EXPECTED_APP_PROFILE = "test-profile"; private static final RequestContext REQUEST_CONTEXT = RequestContext.create("test-project", "test-instance", EXPECTED_APP_PROFILE); private static final String EXPECTED_INSTANCE_NAME = "projects/test-project/instances/test-instance"; + private static final ByteString EXPECTED_PREPARED_QUERY = ByteString.copyFromUtf8("foo"); + // BoundStatement doesn't validate params against schema right now, so we can use hardcoded + // columns for now + private static final ColumnMetadata[] DEFAULT_COLUMNS = { + columnMetadata("_key", bytesType()), columnMetadata("cf", stringType()) + }; + + public static BoundStatement.Builder boundStatementBuilder() { + // This doesn't impact bound statement, but set it so it looks like a real response + Instant expiry = Instant.now().plus(Duration.ofMinutes(1)); + return PreparedStatementImpl.create( + PrepareResponse.fromProto( + PrepareQueryResponse.newBuilder() + .setPreparedQuery(EXPECTED_PREPARED_QUERY) + .setMetadata(metadata(DEFAULT_COLUMNS)) + .setValidUntil( + Timestamp.ofTimeSecondsAndNanos(expiry.getEpochSecond(), expiry.getNano()) + .toProto()) + .build())) + .bind(); + } @Test public void statementWithoutParameters() { - Statement s = Statement.of("SELECT * FROM table"); + BoundStatement s = boundStatementBuilder().build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) .build()); @@ -72,15 +101,13 @@ public void statementWithoutParameters() { @Test public void statementWithBytesParam() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE _key=@key") - .setBytesParam("key", ByteString.copyFromUtf8("test")) - .build(); + BoundStatement s = + boundStatementBuilder().setBytesParam("key", ByteString.copyFromUtf8("test")).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE _key=@key") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "key", Value.newBuilder() @@ -94,15 +121,12 @@ public void statementWithBytesParam() { @Test public void statementWithNullBytesParam() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE _key=@key") - .setBytesParam("key", null) - .build(); + BoundStatement s = boundStatementBuilder().setBytesParam("key", null).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE _key=@key") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("key", Value.newBuilder().setType(bytesType()).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -111,15 +135,12 @@ public void statementWithNullBytesParam() { @Test public void statementWithStringParam() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE _key=@key") - .setStringParam("key", "test") - .build(); + BoundStatement s = boundStatementBuilder().setStringParam("key", "test").build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE _key=@key") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "key", Value.newBuilder().setType(stringType()).setStringValue("test").build()) .setInstanceName(EXPECTED_INSTANCE_NAME) @@ -129,15 +150,12 @@ public void statementWithStringParam() { @Test public void statementWithNullStringParam() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE _key=@key") - .setStringParam("key", null) - .build(); + BoundStatement s = boundStatementBuilder().setStringParam("key", null).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE _key=@key") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("key", Value.newBuilder().setType(stringType()).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -146,15 +164,12 @@ public void statementWithNullStringParam() { @Test public void statementWithInt64Param() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE 1=@number") - .setLongParam("number", 1L) - .build(); + BoundStatement s = boundStatementBuilder().setLongParam("number", 1L).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE 1=@number") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("number", Value.newBuilder().setType(int64Type()).setIntValue(1).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -163,15 +178,12 @@ public void statementWithInt64Param() { @Test public void statementWithNullInt64Param() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE 1=@number") - .setLongParam("number", null) - .build(); + BoundStatement s = boundStatementBuilder().setLongParam("number", null).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE 1=@number") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("number", Value.newBuilder().setType(int64Type()).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -180,15 +192,12 @@ public void statementWithNullInt64Param() { @Test public void statementWithBoolParam() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE @bool") - .setBooleanParam("bool", true) - .build(); + BoundStatement s = boundStatementBuilder().setBooleanParam("bool", true).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE @bool") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "bool", Value.newBuilder().setType(boolType()).setBoolValue(true).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) @@ -198,15 +207,12 @@ public void statementWithBoolParam() { @Test public void statementWithNullBoolParam() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE @bool") - .setBooleanParam("bool", null) - .build(); + BoundStatement s = boundStatementBuilder().setBooleanParam("bool", null).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE @bool") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("bool", Value.newBuilder().setType(boolType()).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -215,17 +221,15 @@ public void statementWithNullBoolParam() { @Test public void statementWithTimestampParam() { - Statement s = - Statement.newBuilder( - "SELECT * FROM table WHERE PARSE_TIMESTAMP(\"%Y/%m/%dT%H:%M:%S\", CAST(cf[\"ts\"] AS STRING)) < @timeParam") + BoundStatement s = + boundStatementBuilder() .setTimestampParam("timeParam", Instant.ofEpochSecond(1000, 100)) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT * FROM table WHERE PARSE_TIMESTAMP(\"%Y/%m/%dT%H:%M:%S\", CAST(cf[\"ts\"] AS STRING)) < @timeParam") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "timeParam", Value.newBuilder() @@ -239,17 +243,12 @@ public void statementWithTimestampParam() { @Test public void statementWithNullTimestampParam() { - Statement s = - Statement.newBuilder( - "SELECT * FROM table WHERE PARSE_TIMESTAMP(\"%Y/%m/%dT%H:%M:%S\", CAST(cf[\"ts\"] AS STRING)) < @timeParam") - .setTimestampParam("timeParam", null) - .build(); + BoundStatement s = boundStatementBuilder().setTimestampParam("timeParam", null).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT * FROM table WHERE PARSE_TIMESTAMP(\"%Y/%m/%dT%H:%M:%S\", CAST(cf[\"ts\"] AS STRING)) < @timeParam") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("timeParam", Value.newBuilder().setType(timestampType()).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -258,17 +257,15 @@ public void statementWithNullTimestampParam() { @Test public void statementWithDateParam() { - Statement s = - Statement.newBuilder( - "SELECT * FROM table WHERE PARSE_DATE(\"%Y%m%d\", CAST(cf[\"date\"] AS STRING)) < @dateParam") + BoundStatement s = + boundStatementBuilder() .setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 11)) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT * FROM table WHERE PARSE_DATE(\"%Y%m%d\", CAST(cf[\"date\"] AS STRING)) < @dateParam") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "dateParam", Value.newBuilder() @@ -282,17 +279,12 @@ public void statementWithDateParam() { @Test public void statementWithNullDateParam() { - Statement s = - Statement.newBuilder( - "SELECT * FROM table WHERE PARSE_DATE(\"%Y%m%d\", CAST(cf[\"date\"] AS STRING)) < @dateParam") - .setDateParam("dateParam", null) - .build(); + BoundStatement s = boundStatementBuilder().setDateParam("dateParam", null).build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT * FROM table WHERE PARSE_DATE(\"%Y%m%d\", CAST(cf[\"date\"] AS STRING)) < @dateParam") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams("dateParam", Value.newBuilder().setType(dateType()).build()) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) @@ -301,9 +293,8 @@ public void statementWithNullDateParam() { @Test public void statementWithBytesListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam( "listParam", Arrays.asList(ByteString.copyFromUtf8("foo"), ByteString.copyFromUtf8("bar")), @@ -316,11 +307,10 @@ public void statementWithBytesListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.bytes())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -350,9 +340,8 @@ public void statementWithBytesListParam() { @Test public void statementWithStringListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam( "listParam", Arrays.asList("foo", "bar"), SqlType.arrayOf(SqlType.string())) .setListParam( @@ -363,11 +352,10 @@ public void statementWithStringListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.string())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -397,9 +385,8 @@ public void statementWithStringListParam() { @Test public void statementWithInt64ListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam("listParam", Arrays.asList(1L, 2L), SqlType.arrayOf(SqlType.int64())) .setListParam( "listWithNullElem", Arrays.asList(null, 3L, 4L), SqlType.arrayOf(SqlType.int64())) @@ -407,11 +394,10 @@ public void statementWithInt64ListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.int64())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -439,9 +425,8 @@ public void statementWithInt64ListParam() { @Test public void statementWithFloat32ListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam( "listParam", Arrays.asList(1.1f, 1.2f), SqlType.arrayOf(SqlType.float32())) .setListParam( @@ -452,11 +437,10 @@ public void statementWithFloat32ListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.float32())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -486,9 +470,8 @@ public void statementWithFloat32ListParam() { @Test public void statementWithFloat64ListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam( "listParam", Arrays.asList(1.1d, 1.2d), SqlType.arrayOf(SqlType.float64())) .setListParam( @@ -499,11 +482,10 @@ public void statementWithFloat64ListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.float64())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -532,9 +514,8 @@ public void statementWithFloat64ListParam() { @Test public void statementWithBooleanListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam("listParam", Arrays.asList(true, false), SqlType.arrayOf(SqlType.bool())) .setListParam( "listWithNullElem", @@ -544,11 +525,10 @@ public void statementWithBooleanListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.bool())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -578,9 +558,8 @@ public void statementWithBooleanListParam() { @Test public void statementWithTimestampListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam( "listParam", Arrays.asList(Instant.ofEpochSecond(3000, 100), Instant.ofEpochSecond(4000, 100)), @@ -595,11 +574,10 @@ public void statementWithTimestampListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.timestamp())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -634,9 +612,8 @@ public void statementWithTimestampListParam() { @Test public void statementWithDateListParam() { - Statement s = - Statement.newBuilder( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + BoundStatement s = + boundStatementBuilder() .setListParam( "listParam", Arrays.asList(Date.fromYearMonthDay(2024, 6, 1), Date.fromYearMonthDay(2024, 7, 1)), @@ -650,11 +627,10 @@ public void statementWithDateListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.date())) .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery( - "SELECT cf, @listParam, @listWithNullElem, @emptyList, @nullList FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "listParam", Value.newBuilder() @@ -685,7 +661,7 @@ public void statementWithDateListParam() { @Test public void setListParamRejectsUnsupportedElementTypes() { - Statement.Builder statement = Statement.newBuilder("SELECT @param"); + BoundStatement.Builder statement = boundStatementBuilder(); assertThrows( IllegalArgumentException.class, @@ -704,18 +680,18 @@ public void setListParamRejectsUnsupportedElementTypes() { @Test public void statementBuilderAllowsParamsToBeOverridden() { - Statement s = - Statement.newBuilder("SELECT * FROM table WHERE _key=@key") + BoundStatement s = + boundStatementBuilder() .setStringParam("key", "test1") .setStringParam("key", "test2") .setStringParam("key", "test3") .setStringParam("key", "test4") .build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table WHERE _key=@key") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .putParams( "key", Value.newBuilder().setType(stringType()).setStringValue("test4").build()) .setInstanceName(EXPECTED_INSTANCE_NAME) @@ -725,12 +701,12 @@ public void statementBuilderAllowsParamsToBeOverridden() { @Test public void builderWorksWithNoParams() { - Statement s = Statement.newBuilder("SELECT * FROM table").build(); + BoundStatement s = boundStatementBuilder().build(); - assertThat(s.toProto(REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( ExecuteQueryRequest.newBuilder() - .setQuery("SELECT * FROM table") + .setPreparedQuery(EXPECTED_PREPARED_QUERY) .setInstanceName(EXPECTED_INSTANCE_NAME) .setAppProfileId(EXPECTED_APP_PROFILE) .build()); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java index 258819e24b..5890454b03 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java @@ -856,7 +856,7 @@ public void prepareQuery( responseObserver.onNext( // Need to set metadata for response to parse PrepareQueryResponse.newBuilder() - .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .setMetadata(metadata(columnMetadata("foo", stringType()))) .build()); responseObserver.onCompleted(); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java index 472c74e511..1e185a7038 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java @@ -36,7 +36,7 @@ import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.Row; import com.google.cloud.bigtable.data.v2.models.RowMutation; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Range; @@ -828,7 +828,7 @@ public void executeQuerySettingsAreNotLost() { @Test public void executeQueryHasSaneDefaults() { - ServerStreamingCallSettings.Builder builder = + ServerStreamingCallSettings.Builder builder = EnhancedBigtableStubSettings.newBuilder().executeQuerySettings(); // Retries aren't supported right now @@ -841,7 +841,7 @@ public void executeQueryHasSaneDefaults() { @Test public void executeQueryRetriesAreDisabled() { - ServerStreamingCallSettings.Builder builder = + ServerStreamingCallSettings.Builder builder = EnhancedBigtableStubSettings.newBuilder().executeQuerySettings(); assertThat(builder.getRetrySettings().getMaxAttempts()).isAtMost(1); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index 88ab52fb7d..302ff35039 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -70,6 +70,7 @@ import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; @@ -86,8 +87,8 @@ import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider; import com.google.cloud.bigtable.data.v2.stub.sql.ExecuteQueryCallable; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream; @@ -156,6 +157,14 @@ public class EnhancedBigtableStubTest { private static final String WAIT_TIME_TABLE_ID = "test-wait-timeout"; private static final String WAIT_TIME_QUERY = "test-wait-timeout"; private static final Duration WATCHDOG_CHECK_DURATION = Duration.ofMillis(100); + private static final PrepareResponse PREPARE_RESPONSE = + PrepareResponse.fromProto( + PrepareQueryResponse.newBuilder() + .setPreparedQuery(ByteString.copyFromUtf8(WAIT_TIME_QUERY)) + .setMetadata(metadata(columnMetadata("foo", stringType()))) + .build()); + private static final PreparedStatement WAIT_TIME_PREPARED_STATEMENT = + PreparedStatementImpl.create(PREPARE_RESPONSE); private Server server; private MetadataInterceptor metadataInterceptor; @@ -430,9 +439,7 @@ public void testPrepareQueryRequestResponseConversion() assertThat(protoReq) .isEqualTo(req.toProto(RequestContext.create(PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID))); assertThat(f.get().resultSetMetadata()) - .isEqualTo( - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("foo", stringType())).getMetadata())); + .isEqualTo(ProtoResultSetMetadata.fromProto(metadata(columnMetadata("foo", stringType())))); assertThat(f.get().preparedQuery()).isEqualTo(ByteString.copyFromUtf8("foo")); assertThat(f.get().validUntil()).isEqualTo(Instant.ofEpochSecond(1000, 1000)); } @@ -459,9 +466,7 @@ public void testPrepareQueryRequestParams() throws ExecutionException, Interrupt assertThat(reqMetadata.keys()).contains("bigtable-client-attempt-epoch-usec"); assertThat(f.get().resultSetMetadata()) - .isEqualTo( - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("foo", stringType())).getMetadata())); + .isEqualTo(ProtoResultSetMetadata.fromProto(metadata(columnMetadata("foo", stringType())))); assertThat(f.get().preparedQuery()).isEqualTo(ByteString.copyFromUtf8("foo")); assertThat(f.get().validUntil()).isEqualTo(Instant.ofEpochSecond(1000, 1000)); } @@ -894,13 +899,19 @@ public void testBatchMutationRPCErrorCode() { @Test public void testCreateExecuteQueryCallable() throws InterruptedException { ExecuteQueryCallable streamingCallable = enhancedBigtableStub.createExecuteQueryCallable(); - - SqlServerStream sqlServerStream = streamingCallable.call(Statement.of("SELECT * FROM table")); + PrepareResponse prepareResponse = + PrepareResponse.fromProto( + PrepareQueryResponse.newBuilder() + .setPreparedQuery(ByteString.copyFromUtf8("abc")) + .setMetadata(metadata(columnMetadata("foo", stringType()))) + .build()); + PreparedStatement preparedStatement = PreparedStatementImpl.create(prepareResponse); + SqlServerStream sqlServerStream = streamingCallable.call(preparedStatement.bind().build()); ExecuteQueryRequest expectedRequest = ExecuteQueryRequest.newBuilder() .setInstanceName(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID)) .setAppProfileId(APP_PROFILE_ID) - .setQuery("SELECT * FROM table") + .setPreparedQuery(ByteString.copyFromUtf8("abc")) .build(); assertThat(sqlServerStream.rows().iterator().next()).isNotNull(); assertThat(sqlServerStream.metadataFuture().isDone()).isTrue(); @@ -917,7 +928,10 @@ public void testExecuteQueryWaitTimeoutIsSet() throws IOException { EnhancedBigtableStub stub = EnhancedBigtableStub.create(settings.build()); Iterator iterator = - stub.executeQueryCallable().call(Statement.of(WAIT_TIME_QUERY)).rows().iterator(); + stub.executeQueryCallable() + .call(WAIT_TIME_PREPARED_STATEMENT.bind().build()) + .rows() + .iterator(); WatchdogTimeoutException e = assertThrows(WatchdogTimeoutException.class, iterator::next); assertThat(e).hasMessageThat().contains("Canceled due to timeout waiting for next response"); } @@ -933,7 +947,9 @@ public void testExecuteQueryWaitTimeoutWorksWithMetadataFuture() try (EnhancedBigtableStub stub = EnhancedBigtableStub.create(settings.build())) { ApiFuture future = - stub.executeQueryCallable().call(Statement.of(WAIT_TIME_QUERY)).metadataFuture(); + stub.executeQueryCallable() + .call(WAIT_TIME_PREPARED_STATEMENT.bind().build()) + .metadataFuture(); ExecutionException e = assertThrows(ExecutionException.class, future::get); assertThat(e.getCause()).isInstanceOf(WatchdogTimeoutException.class); @@ -1085,7 +1101,7 @@ public void pingAndWarm( @Override public void executeQuery( ExecuteQueryRequest request, StreamObserver responseObserver) { - if (request.getQuery().contains(WAIT_TIME_QUERY)) { + if (request.getPreparedQuery().startsWith(ByteString.copyFromUtf8(WAIT_TIME_QUERY))) { try { Thread.sleep(WATCHDOG_CHECK_DURATION.toMillis() * 2); } catch (Exception e) { @@ -1093,8 +1109,8 @@ public void executeQuery( } } executeQueryRequests.add(request); - responseObserver.onNext(metadata(columnMetadata("foo", stringType()))); responseObserver.onNext(partialResultSetWithToken(stringValue("test"))); + responseObserver.onCompleted(); } @Override @@ -1111,7 +1127,7 @@ public void prepareQuery( responseObserver.onNext( PrepareQueryResponse.newBuilder() .setPreparedQuery(ByteString.copyFromUtf8("foo")) - .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .setMetadata(metadata(columnMetadata("foo", stringType()))) .setValidUntil(Timestamp.newBuilder().setSeconds(1000).setNanos(1000).build()) .build()); responseObserver.onCompleted(); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java index bee8374d9b..ff65fbf9ce 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java @@ -17,6 +17,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; @@ -41,13 +42,16 @@ import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Mutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.rpc.Status; import io.grpc.Metadata; import io.grpc.Server; @@ -171,7 +175,11 @@ public void readModifyWriteTest() { @Test public void executeQueryTest() { - client.executeQuery(Statement.of("SELECT * FROM table")); + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse(metadata(columnMetadata("foo", stringType()))))); + client.executeQuery(new BoundStatement.Builder(preparedStatement).build()); verifyHeaderSent(true); } @@ -278,7 +286,7 @@ public void prepareQuery( responseObserver.onNext( // Need to set metadata for response to parse PrepareQueryResponse.newBuilder() - .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .setMetadata(metadata(columnMetadata("foo", stringType()))) .build()); responseObserver.onCompleted(); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java index 024d36e018..937b03102a 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java @@ -853,7 +853,7 @@ public void prepareQuery( responseObserver.onNext( // Need to set metadata for response to parse PrepareQueryResponse.newBuilder() - .setMetadata(metadata(columnMetadata("foo", stringType())).getMetadata()) + .setMetadata(metadata(columnMetadata("foo", stringType()))) .build()); responseObserver.onCompleted(); } else { diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java new file mode 100644 index 0000000000..f48eb10623 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java @@ -0,0 +1,89 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.stub.sql; + +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.api.core.SettableApiFuture; +import com.google.bigtable.v2.ExecuteQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.protobuf.ByteString; +import java.util.concurrent.ExecutionException; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ExecuteQueryCallContextTest { + private static final ByteString PREPARED_QUERY = ByteString.copyFromUtf8("test"); + private static final com.google.bigtable.v2.ResultSetMetadata METADATA = + metadata(columnMetadata("foo", stringType()), columnMetadata("bar", bytesType())); + private static final PreparedStatement PREPARED_STATEMENT = + PreparedStatementImpl.create( + PrepareResponse.fromProto(prepareResponse(PREPARED_QUERY, METADATA))); + + @Test + public void testToRequest() { + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create( + PREPARED_STATEMENT.bind().setStringParam("foo", "val").build(), + SettableApiFuture.create()); + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + ExecuteQueryRequest request = callContext.toRequest(requestContext); + + assertThat(request.getPreparedQuery()).isEqualTo(PREPARED_QUERY); + assertThat(request.getAppProfileId()).isEqualTo("profile"); + assertThat(request.getInstanceName()) + .isEqualTo(NameUtil.formatInstanceName("project", "instance")); + assertThat(request.getParamsMap().get("foo").getStringValue()).isEqualTo("val"); + assertThat(request.getParamsMap().get("foo").getType()).isEqualTo(stringType()); + } + + @Test + public void testFirstResponseReceived() throws ExecutionException, InterruptedException { + SettableApiFuture mdFuture = SettableApiFuture.create(); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(PREPARED_STATEMENT.bind().build(), mdFuture); + + callContext.firstResponseReceived(); + assertThat(mdFuture.isDone()).isTrue(); + assertThat(mdFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); + } + + @Test + public void testSetMetadataException() { + SettableApiFuture mdFuture = SettableApiFuture.create(); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(PREPARED_STATEMENT.bind().build(), mdFuture); + + callContext.setMetadataException(new RuntimeException("test")); + assertThat(mdFuture.isDone()).isTrue(); + ExecutionException e = assertThrows(ExecutionException.class, mdFuture::get); + assertThat(e.getCause()).isInstanceOf(RuntimeException.class); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java index 7c6d3df881..56785459e5 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java @@ -17,6 +17,8 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; @@ -30,11 +32,14 @@ import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.ProtoSqlRow; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement.Builder; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; import io.grpc.Context; @@ -57,8 +62,26 @@ @RunWith(JUnit4.class) public class ExecuteQueryCallableTest { + private static final class FakePreparedStatement implements PreparedStatement { + + @Override + public Builder bind() { + return new BoundStatement.Builder(this); + } + + @Override + public PrepareResponse getPrepareResponse() { + return PrepareResponse.fromProto( + prepareResponse(metadata(columnMetadata("foo", stringType())))); + } + + @Override + public void close() throws Exception {} + } + private static final RequestContext REQUEST_CONTEXT = RequestContext.create("fake-project", "fake-instance", "fake-profile"); + private static final PreparedStatement PREPARED_STATEMENT = new FakePreparedStatement(); private Server server; private FakeService fakeService = new FakeService(); @@ -87,13 +110,12 @@ public void tearDown() { public void testCallContextAndServerStreamSetup() { SqlRow row = ProtoSqlRow.create( - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("test", stringType())).getMetadata()), + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("test", stringType()))), Collections.singletonList(stringValue("foo"))); ServerStreamingStashCallable innerCallable = new ServerStreamingStashCallable<>(Collections.singletonList(row)); ExecuteQueryCallable callable = new ExecuteQueryCallable(innerCallable, REQUEST_CONTEXT); - SqlServerStream stream = callable.call(Statement.of("SELECT * FROM table")); + SqlServerStream stream = callable.call(PREPARED_STATEMENT.bind().build()); assertThat(stream.metadataFuture()) .isEqualTo(innerCallable.getActualRequest().resultSetMetadataFuture()); @@ -106,7 +128,7 @@ public void testCallContextAndServerStreamSetup() { public void testExecuteQueryRequestsAreNotRetried() { // TODO: retries for execute query is currently disabled. This test should be // updated once resumption token is in place. - SqlServerStream stream = stub.executeQueryCallable().call(Statement.of("SELECT * FROM table")); + SqlServerStream stream = stub.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); Iterator iterator = stream.rows().iterator(); @@ -128,7 +150,7 @@ public void testExecuteQueryRequestsIgnoreOverriddenMaxAttempts() throws IOExcep try (EnhancedBigtableStub overrideStub = EnhancedBigtableStub.create(overrideSettings.build().getStubSettings())) { SqlServerStream stream = - overrideStub.executeQueryCallable().call(Statement.of("SELECT * FROM table")); + overrideStub.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); Iterator iterator = stream.rows().iterator(); assertThrows(UnavailableException.class, iterator::next).getCause(); @@ -138,7 +160,7 @@ public void testExecuteQueryRequestsIgnoreOverriddenMaxAttempts() throws IOExcep @Test public void testExecuteQueryRequestsSetDefaultDeadline() { - SqlServerStream stream = stub.executeQueryCallable().call(Statement.of("SELECT * FROM table")); + SqlServerStream stream = stub.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); Iterator iterator = stream.rows().iterator(); // We don't care about this but are reusing the fake service that tests retries assertThrows(UnavailableException.class, iterator::next).getCause(); @@ -165,7 +187,7 @@ public void testExecuteQueryRequestsRespectDeadline() throws IOException { try (EnhancedBigtableStub overrideDeadline = EnhancedBigtableStub.create(overrideSettings.build().getStubSettings())) { SqlServerStream streamOverride = - overrideDeadline.executeQueryCallable().call(Statement.of("SELECT * FROM table")); + overrideDeadline.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); Iterator overrideIterator = streamOverride.rows().iterator(); // We don't care about this but are reusing the fake service that tests retries assertThrows(DeadlineExceededException.class, overrideIterator::next).getCause(); @@ -194,7 +216,7 @@ public void executeQuery( throw new RuntimeException(e); } attempts++; - responseObserver.onNext(metadata(columnMetadata("test", stringType()))); + responseObserver.onNext(partialResultSetWithoutToken(stringValue("foo"))); responseObserver.onError(new StatusRuntimeException(Status.UNAVAILABLE)); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java index 1c04a11d33..87909a045f 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java @@ -20,6 +20,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Value; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; @@ -28,7 +29,11 @@ import com.google.api.core.SettableApiFuture; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.stub.sql.MetadataResolvingCallable.MetadataObserver; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; @@ -36,7 +41,6 @@ import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; -import java.util.Arrays; import java.util.Collections; import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; @@ -49,7 +53,7 @@ public class MetadataResolvingCallableTest { private static final ExecuteQueryRequest FAKE_REQUEST = ExecuteQueryRequest.newBuilder().build(); - private static final ExecuteQueryResponse METADATA = + private static final com.google.bigtable.v2.ResultSetMetadata METADATA = metadata(columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())); private static final ExecuteQueryResponse DATA = partialResultSetWithToken(stringValue("fooVal"), int64Value(100)); @@ -61,59 +65,37 @@ public class MetadataResolvingCallableTest { @Before public void setUp() { metadataFuture = SettableApiFuture.create(); + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse( + metadata( + columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))))); + + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); outerObserver = new MockResponseObserver<>(true); - observer = new MetadataObserver(outerObserver, metadataFuture); + observer = new MetadataObserver(outerObserver, callContext); } @Test - public void observer_parsesMetadataSetsFutureAndPassesThroughResponses() + public void observer_setsFutureAndPassesThroughResponses() throws ExecutionException, InterruptedException { ServerStreamingStashCallable innerCallable = - new ServerStreamingStashCallable<>(Arrays.asList(METADATA, DATA)); + new ServerStreamingStashCallable<>(Collections.singletonList(DATA)); innerCallable.call(FAKE_REQUEST, observer); assertThat(metadataFuture.isDone()).isTrue(); - assertThat(metadataFuture.get()) - .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA.getMetadata())); - assertThat(outerObserver.popNextResponse()).isEqualTo(METADATA); + assertThat(metadataFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); assertThat(outerObserver.popNextResponse()).isEqualTo(DATA); assertThat(outerObserver.isDone()).isTrue(); assertThat(outerObserver.getFinalError()).isNull(); } - @Test - public void observer_invalidMetadataFailsFutureAndPassesThroughError() { - ExecuteQueryResponse invalidMetadataResponse = metadata(); - ServerStreamingStashCallable innerCallable = - new ServerStreamingStashCallable<>(Arrays.asList(invalidMetadataResponse, DATA)); - innerCallable.call(FAKE_REQUEST, observer); - - assertThat(metadataFuture.isDone()).isTrue(); - assertThrows(ExecutionException.class, metadataFuture::get); - ExecutionException e = assertThrows(ExecutionException.class, metadataFuture::get); - assertThat(e.getCause()).isInstanceOf(IllegalStateException.class); - assertThat(outerObserver.isDone()).isTrue(); - assertThat(outerObserver.getFinalError()).isInstanceOf(IllegalStateException.class); - } - - @Test - public void observer_invalidFirstResponseFailsFutureAndPassesThroughError() { - ServerStreamingStashCallable innerCallable = - new ServerStreamingStashCallable<>(Collections.singletonList(DATA)); - innerCallable.call(FAKE_REQUEST, observer); - - assertThat(metadataFuture.isDone()).isTrue(); - assertThrows(ExecutionException.class, metadataFuture::get); - ExecutionException e = assertThrows(ExecutionException.class, metadataFuture::get); - assertThat(e.getCause()).isInstanceOf(IllegalStateException.class); - assertThat(outerObserver.isDone()).isTrue(); - assertThat(outerObserver.getFinalError()).isInstanceOf(IllegalStateException.class); - } - // cancel will manifest as an onError call so these are testing both cancellation and // other exceptions @Test - public void observer_passesThroughErrorBeforeReceivingMetadata() { + public void observer_passesThroughErrorBeforeResolvingMetadata() { MockServerStreamingCallable innerCallable = new MockServerStreamingCallable<>(); innerCallable.call(FAKE_REQUEST, observer); @@ -132,7 +114,7 @@ public void observer_passesThroughErrorBeforeReceivingMetadata() { } @Test - public void observer_passesThroughErrorAfterReceivingMetadata() + public void observer_passesThroughErrorAfterSettingMetadata() throws ExecutionException, InterruptedException { MockServerStreamingCallable innerCallable = new MockServerStreamingCallable<>(); @@ -141,13 +123,13 @@ public void observer_passesThroughErrorAfterReceivingMetadata() innerCallable.popLastCall(); MockStreamController innerController = lastCall.getController(); - innerController.getObserver().onResponse(METADATA); + innerController.getObserver().onResponse(ExecuteQueryResponse.getDefaultInstance()); innerController.getObserver().onError(new RuntimeException("exception after metadata")); assertThat(metadataFuture.isDone()).isTrue(); - assertThat(metadataFuture.get()) - .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA.getMetadata())); - assertThat(outerObserver.popNextResponse()).isEqualTo(METADATA); + assertThat(metadataFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); + assertThat(outerObserver.popNextResponse()) + .isEqualTo(ExecuteQueryResponse.getDefaultInstance()); assertThat(outerObserver.isDone()).isTrue(); assertThat(outerObserver.getFinalError()).isInstanceOf(RuntimeException.class); } @@ -165,7 +147,8 @@ public void observer_passThroughOnStart() { } @Test - public void observer_onCompleteBeforeMetadata_throwsException() throws InterruptedException { + public void observer_onCompleteWithNoData_resolvesMetadata() + throws InterruptedException, ExecutionException { MockServerStreamingCallable innerCallable = new MockServerStreamingCallable<>(); innerCallable.call(FAKE_REQUEST, observer); @@ -174,29 +157,34 @@ public void observer_onCompleteBeforeMetadata_throwsException() throws Interrupt MockStreamController innerController = lastCall.getController(); innerController.getObserver().onComplete(); - assertThrows(ExecutionException.class, metadataFuture::get); - ExecutionException e = assertThrows(ExecutionException.class, metadataFuture::get); - assertThat(e.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(metadataFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); assertThat(outerObserver.isDone()).isTrue(); - assertThat(outerObserver.getFinalError()).isInstanceOf(IllegalStateException.class); + assertThat(outerObserver.getFinalError()).isNull(); } @Test public void testCallable() throws ExecutionException, InterruptedException { ServerStreamingStashCallable innerCallable = - new ServerStreamingStashCallable<>(Arrays.asList(METADATA, DATA)); - MetadataResolvingCallable callable = new MetadataResolvingCallable(innerCallable); + new ServerStreamingStashCallable<>(Collections.singletonList(DATA)); + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + MetadataResolvingCallable callable = + new MetadataResolvingCallable(innerCallable, requestContext); MockResponseObserver outerObserver = new MockResponseObserver<>(true); SettableApiFuture metadataFuture = SettableApiFuture.create(); + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse( + metadata( + columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))))); + ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create(FAKE_REQUEST, metadataFuture); + ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); callable.call(callContext, outerObserver); assertThat(metadataFuture.isDone()).isTrue(); - assertThat(metadataFuture.get()) - .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA.getMetadata())); - assertThat(outerObserver.popNextResponse()).isEqualTo(METADATA); + assertThat(metadataFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); assertThat(outerObserver.popNextResponse()).isEqualTo(DATA); assertThat(outerObserver.isDone()).isTrue(); assertThat(outerObserver.getFinalError()).isNull(); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java index c4586a5c13..53edec4438 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java @@ -62,8 +62,7 @@ public static final class IndividualTests { @Test public void stateMachine_hasCompleteBatch_falseWhenEmpty() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); assertThat(stateMachine).hasCompleteBatch(false); } @@ -71,8 +70,7 @@ public void stateMachine_hasCompleteBatch_falseWhenEmpty() { @Test public void stateMachine_hasCompleteBatch_falseWhenAwaitingPartialBatch() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); @@ -82,8 +80,7 @@ public void stateMachine_hasCompleteBatch_falseWhenAwaitingPartialBatch() { @Test public void stateMachine_hasCompleteBatch_trueWhenAwaitingBatchConsume() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); @@ -94,8 +91,7 @@ public void stateMachine_hasCompleteBatch_trueWhenAwaitingBatchConsume() { @Test public void stateMachine_isBatchInProgress_falseWhenEmpty() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); assertThat(stateMachine).isBatchInProgress(false); } @@ -103,8 +99,7 @@ public void stateMachine_isBatchInProgress_falseWhenEmpty() { @Test public void stateMachine_isBatchInProgress_trueWhenAwaitingPartialBatch() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); @@ -114,8 +109,7 @@ public void stateMachine_isBatchInProgress_trueWhenAwaitingPartialBatch() { @Test public void stateMachine_isBatchInProgress_trueWhenAwaitingBatchConsume() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); @@ -126,8 +120,7 @@ public void stateMachine_isBatchInProgress_trueWhenAwaitingBatchConsume() { public void stateMachine_consumeRow_throwsExceptionWhenColumnsArentComplete() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType()), columnMetadata("b", stringType())) - .getMetadata()); + metadata(columnMetadata("a", stringType()), columnMetadata("b", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); // this is a valid partial result set so we don't expect an error until we call populateQueue stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("foo")).getResults()); @@ -138,8 +131,7 @@ public void stateMachine_consumeRow_throwsExceptionWhenColumnsArentComplete() { @Test public void stateMachine_consumeRow_throwsExceptionWhenAwaitingPartialBatch() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); // this doesn't have a token so we shouldn't allow results to be processed stateMachine.addPartialResultSet( @@ -151,8 +143,7 @@ public void stateMachine_consumeRow_throwsExceptionWhenAwaitingPartialBatch() { @Test public void stateMachine_mergesPartialBatches() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); @@ -171,7 +162,7 @@ public void stateMachine_mergesPartialBatches() { public void stateMachine_mergesPartialBatches_withRandomChunks() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("map", mapType(stringType(), bytesType()))).getMetadata()); + metadata(columnMetadata("map", mapType(stringType(), bytesType())))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); Value mapVal = mapValue( @@ -201,11 +192,10 @@ public void stateMachine_reconstructsRowWithMultipleColumns() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( metadata( - columnMetadata("a", stringType()), - columnMetadata("b", bytesType()), - columnMetadata("c", arrayType(stringType())), - columnMetadata("d", mapType(stringType(), bytesType()))) - .getMetadata()); + columnMetadata("a", stringType()), + columnMetadata("b", bytesType()), + columnMetadata("c", arrayType(stringType())), + columnMetadata("d", mapType(stringType(), bytesType())))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); Value stringVal = stringValue("test"); @@ -237,8 +227,7 @@ public void stateMachine_reconstructsRowWithMultipleColumns() { public void stateMachine_throwsExceptionWhenValuesDontMatchSchema() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType()), columnMetadata("b", bytesType())) - .getMetadata()); + metadata(columnMetadata("a", stringType()), columnMetadata("b", bytesType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); // values in wrong order @@ -251,8 +240,7 @@ public void stateMachine_throwsExceptionWhenValuesDontMatchSchema() { @Test public void stateMachine_handlesResumeTokenWithNoValues() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet(partialResultSetWithToken().getResults()); @@ -262,8 +250,7 @@ public void stateMachine_handlesResumeTokenWithNoValues() { @Test public void stateMachine_handlesResumeTokenWithOpenBatch() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet( @@ -277,8 +264,7 @@ public void stateMachine_handlesResumeTokenWithOpenBatch() { @Test public void addPartialResultSet_throwsExceptionWhenAwaitingRowConsume() { ResultSetMetadata metadata = - ProtoResultSetMetadata.fromProto( - metadata(columnMetadata("a", stringType())).getMetadata()); + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("test")).getResults()); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java index 34c49fed2e..4402af5ba9 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java @@ -19,6 +19,7 @@ import com.google.bigtable.v2.ColumnMetadata; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.PartialResultSet; +import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.ProtoRows; import com.google.bigtable.v2.ProtoRowsBatch; import com.google.bigtable.v2.ProtoSchema; @@ -36,6 +37,19 @@ public class SqlProtoFactory { private SqlProtoFactory() {} + public static PrepareQueryResponse prepareResponse( + ByteString preparedQuery, ResultSetMetadata metadata) { + return PrepareQueryResponse.newBuilder() + .setPreparedQuery(preparedQuery) + .setValidUntil(Timestamp.newBuilder().setSeconds(1000).setNanos(1000).build()) + .setMetadata(metadata) + .build(); + } + + public static PrepareQueryResponse prepareResponse(ResultSetMetadata metadata) { + return prepareResponse(ByteString.copyFromUtf8("foo"), metadata); + } + public static ColumnMetadata columnMetadata(String name, Type type) { return ColumnMetadata.newBuilder().setName(name).setType(type).build(); } @@ -190,11 +204,9 @@ public static ExecuteQueryResponse tokenOnlyResultSet(ByteString token) { .build(); } - public static ExecuteQueryResponse metadata(ColumnMetadata... columnMetadata) { + public static ResultSetMetadata metadata(ColumnMetadata... columnMetadata) { ProtoSchema schema = ProtoSchema.newBuilder().addAllColumns(Arrays.asList(columnMetadata)).build(); - ResultSetMetadata metadata = ResultSetMetadata.newBuilder().setProtoSchema(schema).build(); - - return ExecuteQueryResponse.newBuilder().setMetadata(metadata).build(); + return ResultSetMetadata.newBuilder().setProtoSchema(schema).build(); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java index 90e9672998..050f794e98 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java @@ -33,7 +33,6 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlRowMergerSubject.assertThat; import static org.junit.Assert.assertThrows; -import com.google.bigtable.v2.ColumnMetadata; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.Value; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; @@ -42,6 +41,7 @@ import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.util.Arrays; +import java.util.function.Supplier; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -49,67 +49,70 @@ @RunWith(JUnit4.class) public class SqlRowMergerTest { + static Supplier toSupplier( + com.google.bigtable.v2.ResultSetMetadata metadataProto) { + return () -> ProtoResultSetMetadata.fromProto(metadataProto); + } + @Test public void sqlRowMerger_handlesEmptyState() { - SqlRowMerger merger = new SqlRowMerger(); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); assertThat(merger).hasPartialFrame(false); assertThat(merger).hasFullFrame(false); } @Test public void sqlRowMerger_handlesMetadata() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), - columnMetadata("bytes", bytesType()), - columnMetadata("strArr", arrayType(stringType())), - columnMetadata("strByteMap", mapType(stringType(), bytesType())) - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata( + columnMetadata("str", stringType()), + columnMetadata("bytes", bytesType()), + columnMetadata("strArr", arrayType(stringType())), + columnMetadata("strByteMap", mapType(stringType(), bytesType()))); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); assertThat(merger).hasPartialFrame(false); assertThat(merger).hasFullFrame(false); } @Test - public void sqlRowMerger_rejectsMetadataOfUnrecognizedType() { - SqlRowMerger merger = new SqlRowMerger(); - ExecuteQueryResponse unrecognizedMetadata = - ExecuteQueryResponse.newBuilder() - .setMetadata(com.google.bigtable.v2.ResultSetMetadata.newBuilder().build()) - .build(); + public void sqlRowMerger_doesntResolveMetadataUntilFirstPush() { + SqlRowMerger merger = + new SqlRowMerger( + () -> { + throw new RuntimeException("test"); + }); - assertThrows(IllegalStateException.class, () -> merger.push(unrecognizedMetadata)); + assertThat(merger).hasPartialFrame(false); + assertThat(merger).hasFullFrame(false); + assertThrows( + RuntimeException.class, () -> merger.push(ExecuteQueryResponse.getDefaultInstance())); } @Test public void hasPartialFrame_trueWithIncompleteBatch() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()) - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); merger.push(partialResultSetWithoutToken(stringValue("test"))); assertThat(merger).hasPartialFrame(true); } @Test public void hasPartialFrame_trueWithFullRow() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()) - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); merger.push(partialResultSetWithToken(stringValue("test"), bytesValue("test"))); assertThat(merger).hasPartialFrame(true); } @Test public void push_failsOnCompleteBatchWithIncompleteRow() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()) - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); assertThrows( IllegalStateException.class, () -> merger.push(partialResultSetWithToken(stringValue("test")))); @@ -117,11 +120,9 @@ public void push_failsOnCompleteBatchWithIncompleteRow() { @Test public void hasFullFrame_trueWithFullRow() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()) - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); merger.push(partialResultSetWithoutToken(stringValue("test"))); merger.push(partialResultSetWithToken(bytesValue("test"))); assertThat(merger).hasFullFrame(true); @@ -129,26 +130,23 @@ public void hasFullFrame_trueWithFullRow() { @Test public void hasFullFrame_falseWithIncompleteBatch() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()) - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); merger.push(partialResultSetWithoutToken(stringValue("test"))); assertThat(merger).hasFullFrame(false); } @Test public void sqlRowMerger_handlesResponseStream() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), - columnMetadata("bytes", bytesType()), - columnMetadata("strArr", arrayType(stringType())), - columnMetadata("strByteMap", mapType(stringType(), bytesType())) - }; - ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columns).getMetadata()); - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata( + columnMetadata("str", stringType()), + columnMetadata("bytes", bytesType()), + columnMetadata("strArr", arrayType(stringType())), + columnMetadata("strByteMap", mapType(stringType(), bytesType()))); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); // Three logical rows worth of values split across two responses Value[] values = { @@ -181,7 +179,9 @@ public void sqlRowMerger_handlesResponseStream() { @Test public void addValue_failsWithoutMetadataFirst() { - SqlRowMerger merger = new SqlRowMerger(); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); assertThrows( IllegalStateException.class, () -> merger.push(partialResultSetWithToken(stringValue("test")))); @@ -189,12 +189,10 @@ public void addValue_failsWithoutMetadataFirst() { @Test public void sqlRowMerger_handlesTokenWithOpenPartialBatch() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()), - }; - ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columns).getMetadata()); - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); merger.push(partialResultSetWithoutToken(stringValue("test"))); merger.push(partialResultSetWithoutToken(bytesValue("test"))); merger.push(tokenOnlyResultSet(ByteString.copyFromUtf8("token"))); @@ -209,11 +207,9 @@ public void sqlRowMerger_handlesTokenWithOpenPartialBatch() { @Test public void sqlRowMerger_handlesTokensWithNoData() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()), - }; - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); merger.push(tokenOnlyResultSet(ByteString.copyFromUtf8("token1"))); merger.push(tokenOnlyResultSet(ByteString.copyFromUtf8("token2"))); merger.push(tokenOnlyResultSet(ByteString.copyFromUtf8("token3"))); @@ -224,12 +220,10 @@ public void sqlRowMerger_handlesTokensWithNoData() { @Test public void sqlRowMerger_handlesLeadingTokens() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = { - columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()), - }; - ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columns).getMetadata()); - merger.push(metadata(columns)); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); + ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); merger.push(tokenOnlyResultSet(ByteString.copyFromUtf8("token1"))); merger.push(partialResultSetWithoutToken(stringValue("test"))); merger.push(partialResultSetWithToken(bytesValue("test"))); @@ -243,18 +237,21 @@ public void sqlRowMerger_handlesLeadingTokens() { } @Test - public void addValue_failsOnDuplicateMetadata() { - SqlRowMerger merger = new SqlRowMerger(); - ColumnMetadata[] columns = {columnMetadata("str", stringType())}; - merger.push(metadata(columns)); - merger.push(partialResultSetWithToken(stringValue("test"))); + public void addValue_failsOnMetadataResponse() { + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); - assertThrows(IllegalStateException.class, () -> merger.push(metadata(columns))); + ExecuteQueryResponse deprecatedMetadataResponse = + ExecuteQueryResponse.newBuilder().setMetadata(metadataProto).build(); + assertThrows(IllegalStateException.class, () -> merger.push(deprecatedMetadataResponse)); } @Test public void pop_failsWhenQueueIsEmpty() { - SqlRowMerger merger = new SqlRowMerger(); + com.google.bigtable.v2.ResultSetMetadata metadataProto = + metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); + SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); assertThrows(NullPointerException.class, merger::pop); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java index 761ca4090f..ecc41f4f89 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java @@ -23,6 +23,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; @@ -30,11 +31,13 @@ import com.google.api.core.SettableApiFuture; import com.google.api.gax.rpc.ServerStream; -import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; -import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ProtoSqlRow; import com.google.cloud.bigtable.data.v2.internal.SqlRow; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; import com.google.common.collect.Lists; @@ -55,28 +58,31 @@ public class SqlRowMergingCallableTest { @Test public void testMerging() { - ExecuteQueryResponse metadataResponse = - metadata( - columnMetadata("stringCol", stringType()), - columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType()))); ServerStreamingStashCallable inner = new ServerStreamingStashCallable<>( Lists.newArrayList( - metadataResponse, partialResultSetWithoutToken( stringValue("foo"), int64Value(1), arrayValue(stringValue("foo"), stringValue("bar"))), partialResultSetWithToken(stringValue("test"), int64Value(10), arrayValue()))); + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse( + metadata( + columnMetadata("stringCol", stringType()), + columnMetadata("intCol", int64Type()), + columnMetadata("arrayCol", arrayType(stringType())))))); + BoundStatement boundStatement = preparedStatement.bind().build(); + ResultSetMetadata metadata = preparedStatement.getPrepareResponse().resultSetMetadata(); + SettableApiFuture mdFuture = SettableApiFuture.create(); + mdFuture.set(metadata); SqlRowMergingCallable rowMergingCallable = new SqlRowMergingCallable(inner); ServerStream results = - rowMergingCallable.call( - ExecuteQueryCallContext.create( - ExecuteQueryRequest.getDefaultInstance(), SettableApiFuture.create())); + rowMergingCallable.call(ExecuteQueryCallContext.create(boundStatement, mdFuture)); List resultsList = results.stream().collect(Collectors.toList()); - ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataResponse.getMetadata()); assertThat(resultsList) .containsExactly( ProtoSqlRow.create( @@ -91,16 +97,53 @@ public void testMerging() { @Test public void testError() { - // empty metadata is invalid + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse( + metadata( + columnMetadata("stringCol", stringType()), + columnMetadata("intCol", int64Type()), + columnMetadata("arrayCol", arrayType(stringType())))))); + BoundStatement boundStatement = preparedStatement.bind().build(); + + // empty response is invalid ServerStreamingStashCallable inner = - new ServerStreamingStashCallable<>(Lists.newArrayList(metadata())); + new ServerStreamingStashCallable<>( + Lists.newArrayList(ExecuteQueryResponse.getDefaultInstance())); SqlRowMergingCallable rowMergingCallable = new SqlRowMergingCallable(inner); + SettableApiFuture mdFuture = SettableApiFuture.create(); + mdFuture.set(preparedStatement.getPrepareResponse().resultSetMetadata()); ServerStream results = - rowMergingCallable.call( - ExecuteQueryCallContext.create( - ExecuteQueryRequest.getDefaultInstance(), SettableApiFuture.create())); + rowMergingCallable.call(ExecuteQueryCallContext.create(boundStatement, mdFuture)); assertThrows(IllegalStateException.class, () -> results.iterator().next()); } + + @Test + public void testMetdataFutureError() { + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse( + metadata( + columnMetadata("stringCol", stringType()), + columnMetadata("intCol", int64Type()), + columnMetadata("arrayCol", arrayType(stringType())))))); + BoundStatement boundStatement = preparedStatement.bind().build(); + + // empty response is invalid + ServerStreamingStashCallable inner = + new ServerStreamingStashCallable<>( + Lists.newArrayList(ExecuteQueryResponse.getDefaultInstance())); + + SqlRowMergingCallable rowMergingCallable = new SqlRowMergingCallable(inner); + SettableApiFuture mdFuture = SettableApiFuture.create(); + mdFuture.setException(new RuntimeException("test")); + ServerStream results = + rowMergingCallable.call(ExecuteQueryCallContext.create(boundStatement, mdFuture)); + + assertThrows(RuntimeException.class, () -> results.iterator().next()); + } } From 575d0e3b2ba7dc1c42e762d870eaabca65880bc2 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Tue, 4 Feb 2025 17:01:43 -0500 Subject: [PATCH 04/11] Validate BoundStatement params match paramTypes Change-Id: Ifbe25eec9fd7d1a79d67d1629ec1782c067ab008 --- .../bigtable/data/v2/BigtableDataClient.java | 2 +- .../v2/internal/PreparedStatementImpl.java | 13 +- .../data/v2/models/sql/BoundStatement.java | 36 +++- .../data/v2/internal/ResultSetImplTest.java | 15 +- .../v2/models/sql/BoundStatementTest.java | 189 +++++++++++++++--- .../v2/stub/EnhancedBigtableStubTest.java | 5 +- .../bigtable/data/v2/stub/HeadersTest.java | 6 +- .../stub/sql/ExecuteQueryCallContextTest.java | 7 +- .../v2/stub/sql/ExecuteQueryCallableTest.java | 3 +- .../sql/MetadataResolvingCallableTest.java | 7 +- .../stub/sql/SqlRowMergingCallableTest.java | 10 +- 11 files changed, 243 insertions(+), 50 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 215cf0305b..694abd4281 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -2754,7 +2754,7 @@ public ResultSet executeQuery(BoundStatement boundStatement) { public PreparedStatement prepareStatement(String query, Map> paramTypes) { PrepareQueryRequest request = PrepareQueryRequest.create(query, paramTypes); PrepareResponse response = stub.prepareQueryCallable().call(request); - return PreparedStatementImpl.create(response); + return PreparedStatementImpl.create(response, paramTypes); } /** Close the clients and releases all associated resources. */ diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java index 46dbad2c3c..abaece4e87 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java @@ -19,6 +19,8 @@ import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement.Builder; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import java.util.Map; /** * Implementation of PreparedStatement that handles PreparedQuery refresh @@ -29,18 +31,21 @@ @InternalApi("For internal use only") public class PreparedStatementImpl implements PreparedStatement { private PrepareResponse response; + private final Map> paramTypes; - public PreparedStatementImpl(PrepareResponse response) { + public PreparedStatementImpl(PrepareResponse response, Map> paramTypes) { this.response = response; + this.paramTypes = paramTypes; } - public static PreparedStatement create(PrepareResponse response) { - return new PreparedStatementImpl(response); + public static PreparedStatement create( + PrepareResponse response, Map> paramTypes) { + return new PreparedStatementImpl(response, paramTypes); } @Override public BoundStatement.Builder bind() { - return new Builder(this); + return new Builder(this, paramTypes); } // TODO update when plan refresh is implement diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java index ffb8fa9c90..d277126f6c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java @@ -26,6 +26,7 @@ import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.QueryParamUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; import com.google.protobuf.Timestamp; @@ -83,9 +84,9 @@ public PrepareResponse getLatestPrepareResponse() { return preparedStatement.getPrepareResponse(); } - // TODO pass in paramTypes, to support validation public static class Builder { private final PreparedStatement preparedStatement; + private final Map> paramTypes; private final Map params; /** @@ -95,13 +96,21 @@ public static class Builder { * applications. */ @InternalApi("For internal use only") - public Builder(PreparedStatement preparedStatement) { + public Builder(PreparedStatement preparedStatement, Map> paramTypes) { this.preparedStatement = preparedStatement; + this.paramTypes = paramTypes; this.params = new HashMap<>(); } /** Builds a {@code Statement} from the builder */ public BoundStatement build() { + for (Map.Entry> paramType : paramTypes.entrySet()) { + String paramName = paramType.getKey(); + if (!params.containsKey(paramName)) { + throw new IllegalArgumentException( + "Attempting to build BoundStatement without binding parameter: " + paramName); + } + } return new BoundStatement(preparedStatement, ImmutableMap.copyOf(params)); } @@ -110,6 +119,7 @@ public BoundStatement build() { * value} */ public Builder setStringParam(String paramName, @Nullable String value) { + validateMatchesParamTypes(paramName, SqlType.string()); params.put(paramName, stringParamOf(value)); return this; } @@ -119,6 +129,7 @@ public Builder setStringParam(String paramName, @Nullable String value) { * value} */ public Builder setBytesParam(String paramName, @Nullable ByteString value) { + validateMatchesParamTypes(paramName, SqlType.bytes()); params.put(paramName, bytesParamOf(value)); return this; } @@ -128,6 +139,7 @@ public Builder setBytesParam(String paramName, @Nullable ByteString value) { * value} */ public Builder setLongParam(String paramName, @Nullable Long value) { + validateMatchesParamTypes(paramName, SqlType.int64()); params.put(paramName, int64ParamOf(value)); return this; } @@ -137,6 +149,7 @@ public Builder setLongParam(String paramName, @Nullable Long value) { * value} */ public Builder setFloatParam(String paramName, @Nullable Float value) { + validateMatchesParamTypes(paramName, SqlType.float32()); params.put(paramName, float32ParamOf(value)); return this; } @@ -146,6 +159,7 @@ public Builder setFloatParam(String paramName, @Nullable Float value) { * value} */ public Builder setDoubleParam(String paramName, @Nullable Double value) { + validateMatchesParamTypes(paramName, SqlType.float64()); params.put(paramName, float64ParamOf(value)); return this; } @@ -154,6 +168,7 @@ public Builder setDoubleParam(String paramName, @Nullable Double value) { * Sets a query parameter with the name {@code paramName} and the BOOL typed value {@code value} */ public Builder setBooleanParam(String paramName, @Nullable Boolean value) { + validateMatchesParamTypes(paramName, SqlType.bool()); params.put(paramName, booleanParamOf(value)); return this; } @@ -163,6 +178,7 @@ public Builder setBooleanParam(String paramName, @Nullable Boolean value) { * value} */ public Builder setTimestampParam(String paramName, @Nullable Instant value) { + validateMatchesParamTypes(paramName, SqlType.timestamp()); params.put(paramName, timestampParamOf(value)); return this; } @@ -171,6 +187,7 @@ public Builder setTimestampParam(String paramName, @Nullable Instant value) { * Sets a query parameter with the name {@code paramName} and the DATE typed value {@code value} */ public Builder setDateParam(String paramName, @Nullable Date value) { + validateMatchesParamTypes(paramName, SqlType.date()); params.put(paramName, dateParamOf(value)); return this; } @@ -182,10 +199,25 @@ public Builder setDateParam(String paramName, @Nullable Date value) { */ public Builder setListParam( String paramName, @Nullable List value, SqlType.Array arrayType) { + validateMatchesParamTypes(paramName, arrayType); params.put(paramName, arrayParamOf(value, arrayType)); return this; } + private void validateMatchesParamTypes(String paramName, SqlType expectedType) { + Preconditions.checkArgument( + paramTypes.containsKey(paramName), "No parameter named: " + paramName); + SqlType actualType = paramTypes.get(paramName); + Preconditions.checkArgument( + SqlType.typesMatch(expectedType, actualType), + "Invalid type passed for query param '" + + paramName + + "'. Expected: " + + expectedType + + " received: " + + actualType); + } + private static Value stringParamOf(@Nullable String value) { Type type = QueryParamUtil.convertToQueryParamProto(SqlType.string()); Value.Builder builder = nullValueWithType(type); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java index b4f1515923..2aee489f11 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java @@ -47,7 +47,6 @@ import com.google.api.core.SettableApiFuture; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.cloud.Date; -import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; @@ -78,7 +77,7 @@ private static ResultSet resultSetWithFakeStream( ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); future.set(metadata); PrepareResponse response = PrepareResponse.fromProto(prepareResponse(protoMetadata)); - PreparedStatement preparedStatement = PreparedStatementImpl.create(response); + PreparedStatement preparedStatement = PreparedStatementImpl.create(response, new HashMap<>()); ExecuteQueryCallContext fakeCallContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), future); return ResultSetImpl.create(SqlServerStreamImpl.create(future, stream.call(fakeCallContext))); @@ -325,10 +324,10 @@ public void getMetadata_unwrapsExecutionExceptions() { new ServerStreamingStashCallable<>(Collections.emptyList()); PrepareResponse prepareResponse = PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))); - PreparedStatement preparedStatement = PreparedStatementImpl.create(prepareResponse); + PreparedStatement preparedStatement = + PreparedStatementImpl.create(prepareResponse, new HashMap<>()); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create( - new BoundStatement.Builder(preparedStatement).build(), metadataFuture); + ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); ResultSet rs = ResultSetImpl.create( SqlServerStreamImpl.create(metadataFuture, stream.call(fakeCallContext))); @@ -344,10 +343,10 @@ public void getMetadata_returnsNonRuntimeExecutionExceptionsWrapped() { new ServerStreamingStashCallable<>(Collections.emptyList()); PrepareResponse prepareResponse = PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))); - PreparedStatement preparedStatement = PreparedStatementImpl.create(prepareResponse); + PreparedStatement preparedStatement = + PreparedStatementImpl.create(prepareResponse, new HashMap<>()); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create( - new BoundStatement.Builder(preparedStatement).build(), metadataFuture); + ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); ResultSet rs = ResultSetImpl.create( SqlServerStreamImpl.create(metadataFuture, stream.call(fakeCallContext))); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java index a27a10ef05..613c5df365 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java @@ -52,6 +52,7 @@ import java.time.Instant; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -71,7 +72,12 @@ public class BoundStatementTest { columnMetadata("_key", bytesType()), columnMetadata("cf", stringType()) }; - public static BoundStatement.Builder boundStatementBuilder() { + // Use ColumnMetadata as a more concise way of specifying params + public static BoundStatement.Builder boundStatementBuilder(ColumnMetadata... paramColumns) { + HashMap> paramTypes = new HashMap<>(paramColumns.length); + for (ColumnMetadata c : paramColumns) { + paramTypes.put(c.getName(), SqlType.fromProto(c.getType())); + } // This doesn't impact bound statement, but set it so it looks like a real response Instant expiry = Instant.now().plus(Duration.ofMinutes(1)); return PreparedStatementImpl.create( @@ -82,7 +88,8 @@ public static BoundStatement.Builder boundStatementBuilder() { .setValidUntil( Timestamp.ofTimeSecondsAndNanos(expiry.getEpochSecond(), expiry.getNano()) .toProto()) - .build())) + .build()), + paramTypes) .bind(); } @@ -102,7 +109,9 @@ public void statementWithoutParameters() { @Test public void statementWithBytesParam() { BoundStatement s = - boundStatementBuilder().setBytesParam("key", ByteString.copyFromUtf8("test")).build(); + boundStatementBuilder(columnMetadata("key", bytesType())) + .setBytesParam("key", ByteString.copyFromUtf8("test")) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -121,7 +130,10 @@ public void statementWithBytesParam() { @Test public void statementWithNullBytesParam() { - BoundStatement s = boundStatementBuilder().setBytesParam("key", null).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("key", bytesType())) + .setBytesParam("key", null) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -135,7 +147,10 @@ public void statementWithNullBytesParam() { @Test public void statementWithStringParam() { - BoundStatement s = boundStatementBuilder().setStringParam("key", "test").build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("key", stringType())) + .setStringParam("key", "test") + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -150,7 +165,10 @@ public void statementWithStringParam() { @Test public void statementWithNullStringParam() { - BoundStatement s = boundStatementBuilder().setStringParam("key", null).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("key", stringType())) + .setStringParam("key", null) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -164,7 +182,10 @@ public void statementWithNullStringParam() { @Test public void statementWithInt64Param() { - BoundStatement s = boundStatementBuilder().setLongParam("number", 1L).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("number", int64Type())) + .setLongParam("number", 1L) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -178,7 +199,10 @@ public void statementWithInt64Param() { @Test public void statementWithNullInt64Param() { - BoundStatement s = boundStatementBuilder().setLongParam("number", null).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("number", int64Type())) + .setLongParam("number", null) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -192,7 +216,10 @@ public void statementWithNullInt64Param() { @Test public void statementWithBoolParam() { - BoundStatement s = boundStatementBuilder().setBooleanParam("bool", true).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("bool", boolType())) + .setBooleanParam("bool", true) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -207,7 +234,10 @@ public void statementWithBoolParam() { @Test public void statementWithNullBoolParam() { - BoundStatement s = boundStatementBuilder().setBooleanParam("bool", null).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("bool", boolType())) + .setBooleanParam("bool", null) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -222,7 +252,7 @@ public void statementWithNullBoolParam() { @Test public void statementWithTimestampParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder(columnMetadata("timeParam", timestampType())) .setTimestampParam("timeParam", Instant.ofEpochSecond(1000, 100)) .build(); @@ -243,7 +273,10 @@ public void statementWithTimestampParam() { @Test public void statementWithNullTimestampParam() { - BoundStatement s = boundStatementBuilder().setTimestampParam("timeParam", null).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("timeParam", timestampType())) + .setTimestampParam("timeParam", null) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -258,7 +291,7 @@ public void statementWithNullTimestampParam() { @Test public void statementWithDateParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder(columnMetadata("dateParam", dateType())) .setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 11)) .build(); @@ -279,7 +312,10 @@ public void statementWithDateParam() { @Test public void statementWithNullDateParam() { - BoundStatement s = boundStatementBuilder().setDateParam("dateParam", null).build(); + BoundStatement s = + boundStatementBuilder(columnMetadata("dateParam", dateType())) + .setDateParam("dateParam", null) + .build(); assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) .isEqualTo( @@ -294,7 +330,11 @@ public void statementWithNullDateParam() { @Test public void statementWithBytesListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(bytesType())), + columnMetadata("listWithNullElem", arrayType(bytesType())), + columnMetadata("emptyList", arrayType(bytesType())), + columnMetadata("nullList", arrayType(bytesType()))) .setListParam( "listParam", Arrays.asList(ByteString.copyFromUtf8("foo"), ByteString.copyFromUtf8("bar")), @@ -341,7 +381,11 @@ public void statementWithBytesListParam() { @Test public void statementWithStringListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(stringType())), + columnMetadata("listWithNullElem", arrayType(stringType())), + columnMetadata("emptyList", arrayType(stringType())), + columnMetadata("nullList", arrayType(stringType()))) .setListParam( "listParam", Arrays.asList("foo", "bar"), SqlType.arrayOf(SqlType.string())) .setListParam( @@ -386,7 +430,11 @@ public void statementWithStringListParam() { @Test public void statementWithInt64ListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(int64Type())), + columnMetadata("listWithNullElem", arrayType(int64Type())), + columnMetadata("emptyList", arrayType(int64Type())), + columnMetadata("nullList", arrayType(int64Type()))) .setListParam("listParam", Arrays.asList(1L, 2L), SqlType.arrayOf(SqlType.int64())) .setListParam( "listWithNullElem", Arrays.asList(null, 3L, 4L), SqlType.arrayOf(SqlType.int64())) @@ -426,7 +474,11 @@ public void statementWithInt64ListParam() { @Test public void statementWithFloat32ListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(float32Type())), + columnMetadata("listWithNullElem", arrayType(float32Type())), + columnMetadata("emptyList", arrayType(float32Type())), + columnMetadata("nullList", arrayType(float32Type()))) .setListParam( "listParam", Arrays.asList(1.1f, 1.2f), SqlType.arrayOf(SqlType.float32())) .setListParam( @@ -471,7 +523,11 @@ public void statementWithFloat32ListParam() { @Test public void statementWithFloat64ListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(float64Type())), + columnMetadata("listWithNullElem", arrayType(float64Type())), + columnMetadata("emptyList", arrayType(float64Type())), + columnMetadata("nullList", arrayType(float64Type()))) .setListParam( "listParam", Arrays.asList(1.1d, 1.2d), SqlType.arrayOf(SqlType.float64())) .setListParam( @@ -515,7 +571,11 @@ public void statementWithFloat64ListParam() { @Test public void statementWithBooleanListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(boolType())), + columnMetadata("listWithNullElem", arrayType(boolType())), + columnMetadata("emptyList", arrayType(boolType())), + columnMetadata("nullList", arrayType(boolType()))) .setListParam("listParam", Arrays.asList(true, false), SqlType.arrayOf(SqlType.bool())) .setListParam( "listWithNullElem", @@ -559,7 +619,11 @@ public void statementWithBooleanListParam() { @Test public void statementWithTimestampListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(timestampType())), + columnMetadata("listWithNullElem", arrayType(timestampType())), + columnMetadata("emptyList", arrayType(timestampType())), + columnMetadata("nullList", arrayType(timestampType()))) .setListParam( "listParam", Arrays.asList(Instant.ofEpochSecond(3000, 100), Instant.ofEpochSecond(4000, 100)), @@ -613,7 +677,11 @@ public void statementWithTimestampListParam() { @Test public void statementWithDateListParam() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder( + columnMetadata("listParam", arrayType(dateType())), + columnMetadata("listWithNullElem", arrayType(dateType())), + columnMetadata("emptyList", arrayType(dateType())), + columnMetadata("nullList", arrayType(dateType()))) .setListParam( "listParam", Arrays.asList(Date.fromYearMonthDay(2024, 6, 1), Date.fromYearMonthDay(2024, 7, 1)), @@ -681,7 +749,7 @@ public void setListParamRejectsUnsupportedElementTypes() { @Test public void statementBuilderAllowsParamsToBeOverridden() { BoundStatement s = - boundStatementBuilder() + boundStatementBuilder(columnMetadata("key", stringType())) .setStringParam("key", "test1") .setStringParam("key", "test2") .setStringParam("key", "test3") @@ -711,4 +779,79 @@ public void builderWorksWithNoParams() { .setAppProfileId(EXPECTED_APP_PROFILE) .build()); } + + @Test + public void builderValidatesParameterNames() { + BoundStatement.Builder builder = boundStatementBuilder(); + + IllegalArgumentException e = + assertThrows( + IllegalArgumentException.class, () -> builder.setStringParam("non-existent", "test")); + assertThat(e.getMessage()).contains("No parameter named: non-existent"); + } + + @Test + public void builderValidatesTypesMatch() { + BoundStatement.Builder builder = + boundStatementBuilder( + columnMetadata("stringParam", stringType()), + columnMetadata("bytesParam", bytesType()), + columnMetadata("stringListParam", stringType())); + + IllegalArgumentException eString = + assertThrows( + IllegalArgumentException.class, + () -> builder.setBytesParam("stringParam", ByteString.copyFromUtf8("foo"))); + assertThat(eString.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eBytes = + assertThrows( + IllegalArgumentException.class, () -> builder.setStringParam("bytesParam", "foo")); + assertThat(eBytes.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eLong = + assertThrows(IllegalArgumentException.class, () -> builder.setLongParam("bytesParam", 1L)); + assertThat(eLong.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eDouble = + assertThrows( + IllegalArgumentException.class, () -> builder.setDoubleParam("bytesParam", 1.1d)); + assertThat(eLong.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eFloat = + assertThrows( + IllegalArgumentException.class, () -> builder.setFloatParam("bytesParam", 1.1f)); + assertThat(eFloat.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eBool = + assertThrows( + IllegalArgumentException.class, () -> builder.setBooleanParam("bytesParam", true)); + assertThat(eBool.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eTs = + assertThrows( + IllegalArgumentException.class, + () -> builder.setTimestampParam("bytesParam", Instant.now())); + assertThat(eTs.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eDate = + assertThrows( + IllegalArgumentException.class, + () -> builder.setDateParam("bytesParam", Date.fromYearMonthDay(2025, 1, 1))); + assertThat(eDate.getMessage()).contains("Invalid type passed for query param"); + IllegalArgumentException eList = + assertThrows( + IllegalArgumentException.class, + () -> + builder.setListParam( + "stringListParam", + Collections.singletonList(ByteString.copyFromUtf8("foo")), + SqlType.arrayOf(SqlType.bytes()))); + assertThat(eList.getMessage()).contains("Invalid type passed for query param"); + } + + @Test + public void builderValidatesAllParamsAreSet() { + BoundStatement.Builder builder = + boundStatementBuilder( + columnMetadata("stringParam", stringType()), columnMetadata("bytesParam", bytesType())); + builder.setStringParam("stringParam", "s"); + + IllegalArgumentException e = assertThrows(IllegalArgumentException.class, builder::build); + assertThat(e.getMessage()) + .contains("Attempting to build BoundStatement without binding parameter: bytesParam"); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index 302ff35039..930f91b488 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -164,7 +164,7 @@ public class EnhancedBigtableStubTest { .setMetadata(metadata(columnMetadata("foo", stringType()))) .build()); private static final PreparedStatement WAIT_TIME_PREPARED_STATEMENT = - PreparedStatementImpl.create(PREPARE_RESPONSE); + PreparedStatementImpl.create(PREPARE_RESPONSE, new HashMap<>()); private Server server; private MetadataInterceptor metadataInterceptor; @@ -905,7 +905,8 @@ public void testCreateExecuteQueryCallable() throws InterruptedException { .setPreparedQuery(ByteString.copyFromUtf8("abc")) .setMetadata(metadata(columnMetadata("foo", stringType()))) .build()); - PreparedStatement preparedStatement = PreparedStatementImpl.create(prepareResponse); + PreparedStatement preparedStatement = + PreparedStatementImpl.create(prepareResponse, new HashMap<>()); SqlServerStream sqlServerStream = streamingCallable.call(preparedStatement.bind().build()); ExecuteQueryRequest expectedRequest = ExecuteQueryRequest.newBuilder() diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java index ff65fbf9ce..7c5d190fbf 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java @@ -50,7 +50,6 @@ import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; -import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.rpc.Status; import io.grpc.Metadata; @@ -178,8 +177,9 @@ public void executeQueryTest() { PreparedStatement preparedStatement = PreparedStatementImpl.create( PrepareResponse.fromProto( - prepareResponse(metadata(columnMetadata("foo", stringType()))))); - client.executeQuery(new BoundStatement.Builder(preparedStatement).build()); + prepareResponse(metadata(columnMetadata("foo", stringType())))), + new HashMap<>()); + client.executeQuery(preparedStatement.bind().build()); verifyHeaderSent(true); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java index f48eb10623..0a113eb0fc 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java @@ -32,7 +32,10 @@ import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; +import java.util.Map; import java.util.concurrent.ExecutionException; import org.junit.Test; import org.junit.runner.RunWith; @@ -43,9 +46,11 @@ public class ExecuteQueryCallContextTest { private static final ByteString PREPARED_QUERY = ByteString.copyFromUtf8("test"); private static final com.google.bigtable.v2.ResultSetMetadata METADATA = metadata(columnMetadata("foo", stringType()), columnMetadata("bar", bytesType())); + private static final Map> PARAM_TYPES = + ImmutableMap.of("foo", SqlType.string()); private static final PreparedStatement PREPARED_STATEMENT = PreparedStatementImpl.create( - PrepareResponse.fromProto(prepareResponse(PREPARED_QUERY, METADATA))); + PrepareResponse.fromProto(prepareResponse(PREPARED_QUERY, METADATA)), PARAM_TYPES); @Test public void testToRequest() { diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java index 56785459e5..ecf0992063 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java @@ -51,6 +51,7 @@ import java.io.IOException; import java.time.Duration; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.concurrent.TimeUnit; import org.junit.After; @@ -66,7 +67,7 @@ private static final class FakePreparedStatement implements PreparedStatement { @Override public Builder bind() { - return new BoundStatement.Builder(this); + return new BoundStatement.Builder(this, new HashMap<>()); } @Override diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java index 87909a045f..dee00a2ec6 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java @@ -42,6 +42,7 @@ import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; import java.util.Collections; +import java.util.HashMap; import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import org.junit.Before; @@ -70,7 +71,8 @@ public void setUp() { PrepareResponse.fromProto( prepareResponse( metadata( - columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))))); + columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), + new HashMap<>()); ExecuteQueryCallContext callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); @@ -176,7 +178,8 @@ public void testCallable() throws ExecutionException, InterruptedException { PrepareResponse.fromProto( prepareResponse( metadata( - columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))))); + columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), + new HashMap<>()); ExecuteQueryCallContext callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java index ecc41f4f89..65cb680599 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java @@ -42,6 +42,7 @@ import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; import com.google.common.collect.Lists; import java.util.Arrays; +import java.util.HashMap; import java.util.List; import java.util.stream.Collectors; import org.junit.Test; @@ -74,7 +75,8 @@ public void testMerging() { metadata( columnMetadata("stringCol", stringType()), columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType())))))); + columnMetadata("arrayCol", arrayType(stringType()))))), + new HashMap<>()); BoundStatement boundStatement = preparedStatement.bind().build(); ResultSetMetadata metadata = preparedStatement.getPrepareResponse().resultSetMetadata(); SettableApiFuture mdFuture = SettableApiFuture.create(); @@ -104,7 +106,8 @@ public void testError() { metadata( columnMetadata("stringCol", stringType()), columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType())))))); + columnMetadata("arrayCol", arrayType(stringType()))))), + new HashMap<>()); BoundStatement boundStatement = preparedStatement.bind().build(); // empty response is invalid @@ -130,7 +133,8 @@ public void testMetdataFutureError() { metadata( columnMetadata("stringCol", stringType()), columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType())))))); + columnMetadata("arrayCol", arrayType(stringType()))))), + new HashMap<>()); BoundStatement boundStatement = preparedStatement.bind().build(); // empty response is invalid From 29def62cd5fd19a75ac38d409d9f904e112d039c Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Fri, 14 Feb 2025 10:31:20 -0500 Subject: [PATCH 05/11] Update ExecuteQuery protocol This incorporates the checksum & reset logic added since preview. It also updates the metadata resolution to be deffered until the first resume token has been received (see comment) Change-Id: I0eda81f85e86a7747c6ce2750dbeaa1f8d0d09b4 --- .../v2/stub/sql/ExecuteQueryCallContext.java | 4 +- .../stub/sql/MetadataResolvingCallable.java | 32 +- .../sql/ProtoRowsMergingStateMachine.java | 116 +++++-- .../data/v2/stub/sql/SqlRowMerger.java | 19 +- .../v2/internal/SqlRowMergerUtilTest.java | 63 ++-- .../stub/sql/ExecuteQueryCallContextTest.java | 8 +- .../sql/MetadataResolvingCallableTest.java | 32 +- .../ProtoRowsMergingStateMachineSubject.java | 6 +- .../sql/ProtoRowsMergingStateMachineTest.java | 303 ++++++++++++++++-- .../data/v2/stub/sql/SqlProtoFactory.java | 73 ++++- .../data/v2/stub/sql/SqlRowMergerTest.java | 84 ++++- 11 files changed, 585 insertions(+), 155 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java index c777a13fc1..33cecc9c85 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java @@ -56,11 +56,11 @@ ExecuteQueryRequest toRequest(RequestContext requestContext) { } /** - * Metadata can change as the plan is refreshed. Once a response or complete has been received + * Metadata can change as the plan is refreshed. Once a resume token or complete has been received * from the stream we know that the {@link com.google.bigtable.v2.PrepareQueryResponse} can no * longer change, so we can set the metadata. */ - void firstResponseReceived() { + void finalizeMetadata() { metadataFuture.set(latestPrepareResponse.resultSetMetadata()); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java index 0ceacdecf0..8b8f850d5f 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java @@ -60,14 +60,14 @@ static final class MetadataObserver extends SafeResponseObserver outerObserver; // This doesn't need to be synchronized because this is called above the reframer // so onResponse will be called sequentially - private boolean isFirstResponse; + private boolean hasReceivedResumeToken; MetadataObserver( ResponseObserver outerObserver, ExecuteQueryCallContext callContext) { super(outerObserver); this.outerObserver = outerObserver; this.callContext = callContext; - this.isFirstResponse = true; + this.hasReceivedResumeToken = false; } @Override @@ -77,10 +77,25 @@ protected void onStartImpl(StreamController streamController) { @Override protected void onResponseImpl(ExecuteQueryResponse response) { - if (isFirstResponse) { - callContext.firstResponseReceived(); + // Defer finalizing metadata until we receive a resume token, because this is the + // only point we can guarantee it won't change. + // + // An example of why this is necessary, for query "SELECT * FROM table": + // - Make a request, table has one column family 'cf' + // - Return an incomplete batch + // - request fails with transient error + // - Meanwhile the table has had a second column family added 'cf2' + // - Retry the request, get an error indicating the `prepared_query` has expired + // - Refresh the prepared_query and retry the request, the new prepared_query + // contains both 'cf' & 'cf2' + // - It sends a new incomplete batch and resets the old outdated batch + // - It send the next chunk with a checksum and resume_token, closing the batch. + // In this case the row merger and the ResultSet should be using the updated schema from + // the refreshed prepare request. + if (!hasReceivedResumeToken && !response.getResults().getResumeToken().isEmpty()) { + callContext.finalizeMetadata(); + hasReceivedResumeToken = true; } - isFirstResponse = false; outerObserver.onResponse(response); } @@ -92,12 +107,11 @@ protected void onErrorImpl(Throwable throwable) { outerObserver.onError(throwable); } - // TODO this becomes a valid state @Override protected void onCompleteImpl() { - if (isFirstResponse) { - // If the stream completes successfully we know we used the current metadata - callContext.firstResponseReceived(); + if (!callContext.resultSetMetadataFuture().isDone()) { + // If stream succeeds with no responses, we can finalize the metadata + callContext.finalizeMetadata(); } outerObserver.onComplete(); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachine.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachine.java index deefda4cad..30899e5d52 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachine.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachine.java @@ -26,11 +26,18 @@ import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.common.hash.HashCode; +import com.google.common.hash.HashFunction; +import com.google.common.hash.Hashing; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; +import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.Queue; +import java.util.function.Supplier; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Used to transform a stream of {@link com.google.bigtable.v2.ProtoRowsBatch} bytes chunks into @@ -41,7 +48,7 @@ * *

      *
    • Add results with {@link #addPartialResultSet(PartialResultSet)} until {@link - * #hasCompleteBatch()} is true + * #hasCompleteBatches()} is true *
    • Call {@link #populateQueue(Queue)} to materialize results from the complete batch. *
    • Repeat until all {@link PartialResultSet}s have been processed *
    • Ensure that there is no incomplete data using {@link #isBatchInProgress()} @@ -52,75 +59,119 @@ @InternalApi final class ProtoRowsMergingStateMachine { enum State { - /** Waiting for the first chunk of bytes for a new batch */ - AWAITING_NEW_BATCH, - /** Waiting for the next chunk of bytes, to combine with the bytes currently being buffered. */ - AWAITING_PARTIAL_BATCH, - /** Buffering a complete batch of rows, waiting for populateQueue to be called for the batch */ + /** Waiting for data to be added to the state machine */ + AWAITING_NEW_DATA, + /** Buffering a complete set of rows, waiting for populateQueue to be called */ AWAITING_BATCH_CONSUME, } - private final ResultSetMetadata metadata; + private static final HashFunction CRC32C = Hashing.crc32c(); + + private final Supplier metadataSupplier; + private @Nullable ResultSetMetadata metadata; private State state; private ByteString batchBuffer; - private ProtoRows completeBatch; + private List> parsedBatches; + private boolean hasReceivedFirstResumeToken; - ProtoRowsMergingStateMachine(ResultSetMetadata metadata) { - this.metadata = metadata; - state = State.AWAITING_NEW_BATCH; + ProtoRowsMergingStateMachine(Supplier metadataSupplier) { + this.metadataSupplier = metadataSupplier; + state = State.AWAITING_NEW_DATA; batchBuffer = ByteString.empty(); + parsedBatches = new ArrayList<>(); + hasReceivedFirstResumeToken = false; } /** * Adds the bytes from the given PartialResultSet to the current buffer. If a resume token is * present, attempts to parse the bytes to the underlying protobuf row format + * + *

      See the comments on {@link PartialResultSet} protobuf message definition for explanation of + * the protocol implemented below. + * + * 

      Translated to use local variable names the expected logic is as follows:
+   * if results.reset {
+   *   reset batchBuffer
+   *   reset parsedBatches
+   * }
+   * if results.proto_rows_batch is set {
+   *   append result.proto_rows_batch.batch_data to batchBuffer
+   * }
+   * if results.batch_checksum is set {
+   *   validate the checksum matches the crc32c hash of batchBuffer
+   *   parse batchBuffer as a ProtoRows message, clearing batchBuffer
+   *   add the parsed data to parsedBatches
+   * }
+   * if results.resume_token is set {
+   *   yield the results in parsedBatches to the row merger.
+   *   this is controlled by the AWAITING_BATCH_CONSUME state.
+   * }
+   *
      */ void addPartialResultSet(PartialResultSet results) { Preconditions.checkState( state != State.AWAITING_BATCH_CONSUME, "Attempting to add partial result set to state machine in state AWAITING_BATCH_CONSUME"); + // If the API indicates we should reset we need to clear buffered data + if (results.getReset()) { + batchBuffer = ByteString.EMPTY; + parsedBatches.clear(); + } // ByteString has an efficient concat which generally involves no copying batchBuffer = batchBuffer.concat(results.getProtoRowsBatch().getBatchData()); - state = State.AWAITING_PARTIAL_BATCH; - if (results.getResumeToken().isEmpty()) { - return; - } - // A resume token means the batch is complete and safe to yield - // We can receive resume tokens with no new data. In this case we yield an empty batch. - if (batchBuffer.isEmpty()) { - completeBatch = ProtoRows.getDefaultInstance(); - } else { + if (results.hasBatchChecksum()) { + HashCode hash = CRC32C.hashBytes(batchBuffer.toByteArray()); + Preconditions.checkState( + hash.hashCode() == results.getBatchChecksum(), "Unexpected checksum mismatch"); try { - completeBatch = ProtoRows.parseFrom(batchBuffer); + ProtoRows completeBatch = ProtoRows.parseFrom(batchBuffer); + batchBuffer = ByteString.EMPTY; + parsedBatches.add(completeBatch.getValuesList()); } catch (InvalidProtocolBufferException e) { throw new InternalError("Unexpected exception parsing response protobuf", e); } } - // Empty buffers can benefit from resetting because ByteString.concat builds a rope - batchBuffer = ByteString.empty(); - state = State.AWAITING_BATCH_CONSUME; + boolean hasResumeToken = !results.getResumeToken().isEmpty(); + if (hasResumeToken) { + if (!hasReceivedFirstResumeToken) { + // Don't resolve the metadata until we receive the first resume token. + // This is safe because we only use the metadata in populateQueue, which can't be called + // until we receive a resume token. For details on why this is necessary, see + // MetadataResolvingCallable + metadata = metadataSupplier.get(); + hasReceivedFirstResumeToken = true; + } + Preconditions.checkState( + batchBuffer.isEmpty(), "Received resumeToken with buffered data and no checksum"); + state = State.AWAITING_BATCH_CONSUME; + } } - /** Returns true if there is a complete batch buffered, false otherwise */ - boolean hasCompleteBatch() { + /** Returns true if there are complete batches, ready to yield. False otherwise */ + boolean hasCompleteBatches() { return state == State.AWAITING_BATCH_CONSUME; } /** Returns true if there is a partial or complete batch buffered, false otherwise */ boolean isBatchInProgress() { - return hasCompleteBatch() || state == State.AWAITING_PARTIAL_BATCH; + boolean hasBufferedData = !batchBuffer.isEmpty() || !parsedBatches.isEmpty(); + return hasCompleteBatches() || hasBufferedData; } /** - * Populates the given queue with the complete batch of rows + * Populates the given queue with the currently buffered rows of rows * - * @throws IllegalStateException if there is not a complete batch + * @throws IllegalStateException if there is no yieldable data */ void populateQueue(Queue queue) { Preconditions.checkState( state == State.AWAITING_BATCH_CONSUME, "Attempting to populate Queue from state machine without completed batch"); - Iterator valuesIterator = completeBatch.getValuesList().iterator(); + Preconditions.checkState( + batchBuffer.isEmpty(), "Unexpected buffered partial batch while consuming rows."); + Preconditions.checkNotNull(metadata, "Unexpected empty metadata when parsing response"); + + Iterator valuesIterator = Iterables.concat(parsedBatches).iterator(); while (valuesIterator.hasNext()) { ImmutableList.Builder rowDataBuilder = ImmutableList.builder(); for (ColumnMetadata c : metadata.getColumns()) { @@ -132,9 +183,8 @@ void populateQueue(Queue queue) { } queue.add(ProtoSqlRow.create(metadata, rowDataBuilder.build())); } - // reset the batch to be empty - completeBatch = ProtoRows.getDefaultInstance(); - state = State.AWAITING_NEW_BATCH; + this.parsedBatches = new ArrayList<>(); + state = State.AWAITING_NEW_DATA; } @InternalApi("VisibleForTestingOnly") diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java index c44ead34f3..a4f2c618e9 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger.java @@ -34,9 +34,7 @@ public final class SqlRowMerger implements Reframer { private final Queue queue; - private ProtoRowsMergingStateMachine stateMachine; - private final Supplier metadataSupplier; - private Boolean isFirstResponse; + private final ProtoRowsMergingStateMachine stateMachine; /** * @param metadataSupplier a supplier of {@link ResultSetMetadata}. This is expected to return @@ -44,9 +42,8 @@ public final class SqlRowMerger implements ReframerThis exists to facilitate plan refresh that can happen after creation of the row merger. */ public SqlRowMerger(Supplier metadataSupplier) { - this.metadataSupplier = metadataSupplier; queue = new ArrayDeque<>(); - isFirstResponse = true; + stateMachine = new ProtoRowsMergingStateMachine(metadataSupplier); } /** @@ -56,13 +53,6 @@ public SqlRowMerger(Supplier metadataSupplier) { */ @Override public void push(ExecuteQueryResponse response) { - if (isFirstResponse) { - // Wait until we've received the first response to get the metadata, as a - // PreparedQuery may need to be refreshed based on initial errors. Once we've - // received a response, it will never change, even upon request resumption. - stateMachine = new ProtoRowsMergingStateMachine(metadataSupplier.get()); - isFirstResponse = false; - } Preconditions.checkState( response.hasResults(), "Expected results response, but received: %s", @@ -73,7 +63,7 @@ public void push(ExecuteQueryResponse response) { private void processProtoRows(PartialResultSet results) { stateMachine.addPartialResultSet(results); - if (stateMachine.hasCompleteBatch()) { + if (stateMachine.hasCompleteBatches()) { stateMachine.populateQueue(queue); } } @@ -95,9 +85,6 @@ public boolean hasFullFrame() { */ @Override public boolean hasPartialFrame() { - if (isFirstResponse) { - return false; - } return hasFullFrame() || stateMachine.isBatchInProgress(); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java index 9986476abb..3de9821147 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtilTest.java @@ -26,6 +26,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSets; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; @@ -74,7 +75,7 @@ public void parseExecuteQueryResponses_handlesSingleValue_serializedProtoRows() @Test public void - parseExecuteQueryResponses_handlesMultipleValuesAccrossMultipleRows_serializedProtoRows() { + parseExecuteQueryResponses_handlesMultipleValuesAcrossMultipleRows_serializedProtoRows() { ColumnMetadata[] columns = { columnMetadata("str", stringType()), columnMetadata("bytes", bytesType()), @@ -84,22 +85,20 @@ public void parseExecuteQueryResponses_handlesSingleValue_serializedProtoRows() com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columns); ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); ImmutableList responses = - ImmutableList.of( - partialResultSetWithoutToken( - stringValue("str1"), - bytesValue("bytes1"), - arrayValue(stringValue("arr1")), - mapValue(mapElement(stringValue("key1"), bytesValue("val1"))), - stringValue("str2")), - partialResultSetWithoutToken( - bytesValue("bytes2"), - arrayValue(stringValue("arr2")), - mapValue(mapElement(stringValue("key2"), bytesValue("val2")))), - partialResultSetWithToken( - stringValue("str3"), - bytesValue("bytes3"), - arrayValue(stringValue("arr3")), - mapValue(mapElement(stringValue("key3"), bytesValue("val3"))))); + partialResultSets( + 3, + stringValue("str1"), + bytesValue("bytes1"), + arrayValue(stringValue("arr1")), + mapValue(mapElement(stringValue("key1"), bytesValue("val1"))), + stringValue("str2"), + bytesValue("bytes2"), + arrayValue(stringValue("arr2")), + mapValue(mapElement(stringValue("key2"), bytesValue("val2"))), + stringValue("str3"), + bytesValue("bytes3"), + arrayValue(stringValue("arr3")), + mapValue(mapElement(stringValue("key3"), bytesValue("val3")))); try (SqlRowMergerUtil util = new SqlRowMergerUtil(metadataProto)) { List rows = util.parseExecuteQueryResponses(responses); assertThat(rows) @@ -165,22 +164,20 @@ public void parseExecuteQueryResponses_worksWithIncrementalSetsOfResponses_seria com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columns); ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadataProto); ImmutableList responses = - ImmutableList.of( - partialResultSetWithoutToken( - stringValue("str1"), - bytesValue("bytes1"), - arrayValue(stringValue("arr1")), - mapValue(mapElement(stringValue("key1"), bytesValue("val1"))), - stringValue("str2")), - partialResultSetWithoutToken( - bytesValue("bytes2"), - arrayValue(stringValue("arr2")), - mapValue(mapElement(stringValue("key2"), bytesValue("val2")))), - partialResultSetWithToken( - stringValue("str3"), - bytesValue("bytes3"), - arrayValue(stringValue("arr3")), - mapValue(mapElement(stringValue("key3"), bytesValue("val3"))))); + partialResultSets( + 3, + stringValue("str1"), + bytesValue("bytes1"), + arrayValue(stringValue("arr1")), + mapValue(mapElement(stringValue("key1"), bytesValue("val1"))), + stringValue("str2"), + bytesValue("bytes2"), + arrayValue(stringValue("arr2")), + mapValue(mapElement(stringValue("key2"), bytesValue("val2"))), + stringValue("str3"), + bytesValue("bytes3"), + arrayValue(stringValue("arr3")), + mapValue(mapElement(stringValue("key3"), bytesValue("val3")))); try (SqlRowMergerUtil util = new SqlRowMergerUtil(metadataProto)) { List rows = new ArrayList<>(); rows.addAll(util.parseExecuteQueryResponses(responses.subList(0, 1))); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java index 0a113eb0fc..f688d96669 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java @@ -73,9 +73,10 @@ public void testToRequest() { public void testFirstResponseReceived() throws ExecutionException, InterruptedException { SettableApiFuture mdFuture = SettableApiFuture.create(); ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create(PREPARED_STATEMENT.bind().build(), mdFuture); + ExecuteQueryCallContext.create( + PREPARED_STATEMENT.bind().setStringParam("foo", "val").build(), mdFuture); - callContext.firstResponseReceived(); + callContext.finalizeMetadata(); assertThat(mdFuture.isDone()).isTrue(); assertThat(mdFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); } @@ -84,7 +85,8 @@ public void testFirstResponseReceived() throws ExecutionException, InterruptedEx public void testSetMetadataException() { SettableApiFuture mdFuture = SettableApiFuture.create(); ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create(PREPARED_STATEMENT.bind().build(), mdFuture); + ExecuteQueryCallContext.create( + PREPARED_STATEMENT.bind().setStringParam("foo", "val").build(), mdFuture); callContext.setMetadataException(new RuntimeException("test")); assertThat(mdFuture.isDone()).isTrue(); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java index dee00a2ec6..854268102e 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java @@ -20,11 +20,15 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Value; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.tokenOnlyResultSet; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import com.google.api.core.SettableApiFuture; import com.google.bigtable.v2.ExecuteQueryRequest; @@ -41,6 +45,7 @@ import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; +import com.google.protobuf.ByteString; import java.util.Collections; import java.util.HashMap; import java.util.concurrent.CancellationException; @@ -59,6 +64,7 @@ public class MetadataResolvingCallableTest { private static final ExecuteQueryResponse DATA = partialResultSetWithToken(stringValue("fooVal"), int64Value(100)); + ExecuteQueryCallContext callContext; MockResponseObserver outerObserver; SettableApiFuture metadataFuture; MetadataResolvingCallable.MetadataObserver observer; @@ -74,15 +80,33 @@ public void setUp() { columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), new HashMap<>()); - ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); + callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); outerObserver = new MockResponseObserver<>(true); observer = new MetadataObserver(outerObserver, callContext); } + @Test + public void observer_doesNotSetFutureUntilTokenReceived() + throws ExecutionException, InterruptedException { + MockServerStreamingCallable innerCallable = + new MockServerStreamingCallable<>(); + innerCallable.call(FAKE_REQUEST, observer); + MockServerStreamingCall lastCall = + innerCallable.popLastCall(); + MockStreamController innerController = lastCall.getController(); + + innerController.getObserver().onResponse(partialResultSetWithoutToken(stringValue("foo"))); + assertFalse(callContext.resultSetMetadataFuture().isDone()); + innerController.getObserver().onResponse(partialResultSetWithToken(stringValue("bar"))); + assertTrue(callContext.resultSetMetadataFuture().isDone()); + assertThat(callContext.resultSetMetadataFuture().get()) + .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); + } + @Test public void observer_setsFutureAndPassesThroughResponses() throws ExecutionException, InterruptedException { + // This has a token so it should finalize the metadata ServerStreamingStashCallable innerCallable = new ServerStreamingStashCallable<>(Collections.singletonList(DATA)); innerCallable.call(FAKE_REQUEST, observer); @@ -125,13 +149,13 @@ public void observer_passesThroughErrorAfterSettingMetadata() innerCallable.popLastCall(); MockStreamController innerController = lastCall.getController(); - innerController.getObserver().onResponse(ExecuteQueryResponse.getDefaultInstance()); + innerController.getObserver().onResponse(tokenOnlyResultSet(ByteString.copyFromUtf8("token"))); innerController.getObserver().onError(new RuntimeException("exception after metadata")); assertThat(metadataFuture.isDone()).isTrue(); assertThat(metadataFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); assertThat(outerObserver.popNextResponse()) - .isEqualTo(ExecuteQueryResponse.getDefaultInstance()); + .isEqualTo(tokenOnlyResultSet(ByteString.copyFromUtf8("token"))); assertThat(outerObserver.isDone()).isTrue(); assertThat(outerObserver.getFinalError()).isInstanceOf(RuntimeException.class); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineSubject.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineSubject.java index 9ec406d71e..e9f6bf09e6 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineSubject.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineSubject.java @@ -46,11 +46,11 @@ public static ProtoRowsMergingStateMachineSubject assertThat( return assertAbout(stateMachine()).that(actual); } - public void hasCompleteBatch(boolean expectation) { + public void hasCompleteBatches(boolean expectation) { if (expectation) { - check("hasCompleteBatch()").that(actual.hasCompleteBatch()).isTrue(); + check("hasCompleteBatch()").that(actual.hasCompleteBatches()).isTrue(); } else { - check("hasCompleteBatch()").that(actual.hasCompleteBatch()).isFalse(); + check("hasCompleteBatch()").that(actual.hasCompleteBatches()).isFalse(); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java index 53edec4438..8a848a8dfd 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ProtoRowsMergingStateMachineTest.java @@ -20,21 +20,30 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.arrayValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.checksum; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Value; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapElement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSets; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.structType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.structValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.tokenOnlyResultSet; import static com.google.common.truth.Truth.assertWithMessage; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.rpc.ApiExceptions; +import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.PartialResultSet; import com.google.bigtable.v2.ProtoRows; import com.google.bigtable.v2.ProtoRowsBatch; @@ -47,6 +56,7 @@ import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.util.ArrayDeque; +import java.util.List; import org.junit.Test; import org.junit.experimental.runners.Enclosed; import org.junit.runner.RunWith; @@ -60,39 +70,39 @@ public final class ProtoRowsMergingStateMachineTest { public static final class IndividualTests { @Test - public void stateMachine_hasCompleteBatch_falseWhenEmpty() { + public void stateMachine_hasCompleteBatches_falseWhenEmpty() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); - assertThat(stateMachine).hasCompleteBatch(false); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + assertThat(stateMachine).hasCompleteBatches(false); } @Test - public void stateMachine_hasCompleteBatch_falseWhenAwaitingPartialBatch() { + public void stateMachine_hasCompleteBatches_falseWhenAwaitingPartialBatch() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); - assertThat(stateMachine).hasCompleteBatch(false); + assertThat(stateMachine).hasCompleteBatches(false); } @Test - public void stateMachine_hasCompleteBatch_trueWhenAwaitingBatchConsume() { + public void stateMachine_hasCompleteBatches_trueWhenAwaitingBatchConsume() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("bar")).getResults()); - assertThat(stateMachine).hasCompleteBatch(true); + assertThat(stateMachine).hasCompleteBatches(true); } @Test public void stateMachine_isBatchInProgress_falseWhenEmpty() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); assertThat(stateMachine).isBatchInProgress(false); } @@ -100,7 +110,7 @@ public void stateMachine_isBatchInProgress_falseWhenEmpty() { public void stateMachine_isBatchInProgress_trueWhenAwaitingPartialBatch() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); assertThat(stateMachine).isBatchInProgress(true); @@ -110,7 +120,7 @@ public void stateMachine_isBatchInProgress_trueWhenAwaitingPartialBatch() { public void stateMachine_isBatchInProgress_trueWhenAwaitingBatchConsume() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); assertThat(stateMachine).isBatchInProgress(true); @@ -121,7 +131,7 @@ public void stateMachine_consumeRow_throwsExceptionWhenColumnsArentComplete() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( metadata(columnMetadata("a", stringType()), columnMetadata("b", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); // this is a valid partial result set so we don't expect an error until we call populateQueue stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("foo")).getResults()); assertThrows( @@ -132,7 +142,7 @@ public void stateMachine_consumeRow_throwsExceptionWhenColumnsArentComplete() { public void stateMachine_consumeRow_throwsExceptionWhenAwaitingPartialBatch() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); // this doesn't have a token so we shouldn't allow results to be processed stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("foo")).getResults()); @@ -144,12 +154,12 @@ public void stateMachine_consumeRow_throwsExceptionWhenAwaitingPartialBatch() { public void stateMachine_mergesPartialBatches() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); - stateMachine.addPartialResultSet( - partialResultSetWithoutToken(stringValue("foo")).getResults()); - stateMachine.addPartialResultSet( - partialResultSetWithoutToken(stringValue("bar")).getResults()); - stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("baz")).getResults()); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + List partialBatches = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + for (ExecuteQueryResponse res : partialBatches) { + stateMachine.addPartialResultSet(res.getResults()); + } assertThat(stateMachine) .populateQueueYields( @@ -163,7 +173,7 @@ public void stateMachine_mergesPartialBatches_withRandomChunks() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( metadata(columnMetadata("map", mapType(stringType(), bytesType())))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); Value mapVal = mapValue( mapElement( @@ -181,6 +191,7 @@ public void stateMachine_mergesPartialBatches_withRandomChunks() { PartialResultSet.newBuilder() .setResumeToken(ByteString.copyFromUtf8("token")) .setProtoRowsBatch(ProtoRowsBatch.newBuilder().setBatchData(chunk2).build()) + .setBatchChecksum(checksum(rows.toByteString())) .build()); assertThat(stateMachine) @@ -196,7 +207,7 @@ public void stateMachine_reconstructsRowWithMultipleColumns() { columnMetadata("b", bytesType()), columnMetadata("c", arrayType(stringType())), columnMetadata("d", mapType(stringType(), bytesType())))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); Value stringVal = stringValue("test"); stateMachine.addPartialResultSet(partialResultSetWithoutToken(stringVal).getResults()); @@ -210,14 +221,14 @@ public void stateMachine_reconstructsRowWithMultipleColumns() { mapElement(stringValue("b"), bytesValue("bVal"))); stateMachine.addPartialResultSet(partialResultSetWithToken(mapVal).getResults()); - assertThat(stateMachine).hasCompleteBatch(true); + assertThat(stateMachine).hasCompleteBatches(true); assertThat(stateMachine) .populateQueueYields( ProtoSqlRow.create( metadata, ImmutableList.of(stringVal, bytesVal, arrayVal, mapVal))); // Once we consume a completed row the state machine should be reset - assertThat(stateMachine).hasCompleteBatch(false); + assertThat(stateMachine).hasCompleteBatches(false); assertThrows( IllegalStateException.class, () -> stateMachine.populateQueue(new ArrayDeque<>())); assertThat(stateMachine).isBatchInProgress(false); @@ -228,7 +239,7 @@ public void stateMachine_throwsExceptionWhenValuesDontMatchSchema() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto( metadata(columnMetadata("a", stringType()), columnMetadata("b", bytesType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); // values in wrong order stateMachine.addPartialResultSet( @@ -241,7 +252,7 @@ public void stateMachine_throwsExceptionWhenValuesDontMatchSchema() { public void stateMachine_handlesResumeTokenWithNoValues() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet(partialResultSetWithToken().getResults()); assertThat(stateMachine).populateQueueYields(new ProtoSqlRow[] {}); @@ -251,7 +262,7 @@ public void stateMachine_handlesResumeTokenWithNoValues() { public void stateMachine_handlesResumeTokenWithOpenBatch() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet( partialResultSetWithoutToken(stringValue("test")).getResults()); @@ -265,7 +276,7 @@ public void stateMachine_handlesResumeTokenWithOpenBatch() { public void addPartialResultSet_throwsExceptionWhenAwaitingRowConsume() { ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); - ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(metadata); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("test")).getResults()); assertThrows( @@ -274,6 +285,242 @@ public void addPartialResultSet_throwsExceptionWhenAwaitingRowConsume() { stateMachine.addPartialResultSet( partialResultSetWithToken(stringValue("test2")).getResults())); } + + @Test + public void stateMachine_throwsExceptionOnChecksumMismatch() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + List responses = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + + // Override the checksum of the final response + PartialResultSet lastResultsWithBadChecksum = + responses.get(2).getResults().toBuilder().setBatchChecksum(1234).build(); + + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + stateMachine.addPartialResultSet(responses.get(0).getResults()); + stateMachine.addPartialResultSet(responses.get(1).getResults()); + + assertThrows( + IllegalStateException.class, + () -> stateMachine.addPartialResultSet(lastResultsWithBadChecksum)); + } + + @Test + public void stateMachine_handlesResetOnPartialBatch() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + // Initial response here has reset bit set + List responses = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + stateMachine.addPartialResultSet(responses.get(0).getResults()); + stateMachine.addPartialResultSet(responses.get(1).getResults()); + + // The two results above should be discarded by reset + for (ExecuteQueryResponse response : responses) { + stateMachine.addPartialResultSet(response.getResults()); + } + + assertThat(stateMachine) + .populateQueueYields( + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("foo"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("bar"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("baz")))); + } + + @Test + public void stateMachine_handlesResetWithUncommittedBatches() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + // Create 2 batches split into multiple chunks. Neither containing a resume token + List firstBatch = + partialResultSets( + 2, + true, + ByteString.EMPTY, + stringValue("foo"), + stringValue("bar"), + stringValue("baz")); + List secondBatch = + partialResultSets( + 3, false, ByteString.EMPTY, stringValue("a"), stringValue("b"), stringValue("c")); + + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + for (ExecuteQueryResponse res : firstBatch) { + stateMachine.addPartialResultSet(res.getResults()); + } + for (ExecuteQueryResponse res : secondBatch) { + stateMachine.addPartialResultSet(res.getResults()); + } + // Nothing should be yielded yet + assertThrows( + IllegalStateException.class, () -> stateMachine.populateQueue(new ArrayDeque<>())); + + List resetBatch = + partialResultSets( + 2, + true, + ByteString.EMPTY, + stringValue("foo2"), + stringValue("bar2"), + stringValue("baz2")); + List batchAfterReset = + partialResultSets( + 3, + false, + ByteString.copyFromUtf8("token"), + stringValue("a2"), + stringValue("b2"), + stringValue("c2")); + for (ExecuteQueryResponse res : resetBatch) { + stateMachine.addPartialResultSet(res.getResults()); + } + for (ExecuteQueryResponse res : batchAfterReset) { + stateMachine.addPartialResultSet(res.getResults()); + } + assertThat(stateMachine) + .populateQueueYields( + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("foo2"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("bar2"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("baz2"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("a2"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("b2"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("c2")))); + } + + @Test + public void stateMachine_handlesMultipleCompleteBatchesBeforeToken() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + // Create 2 batches split into multiple chunks. Neither containing a resume token + List firstBatch = + partialResultSets( + 2, + true, + ByteString.EMPTY, + stringValue("foo"), + stringValue("bar"), + stringValue("baz")); + List secondBatch = + partialResultSets( + 3, false, ByteString.EMPTY, stringValue("a"), stringValue("b"), stringValue("c")); + + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + for (ExecuteQueryResponse res : firstBatch) { + stateMachine.addPartialResultSet(res.getResults()); + } + for (ExecuteQueryResponse res : secondBatch) { + stateMachine.addPartialResultSet(res.getResults()); + } + // Nothing should be yielded yet + assertThrows( + IllegalStateException.class, () -> stateMachine.populateQueue(new ArrayDeque<>())); + ExecuteQueryResponse resultWithToken = partialResultSetWithToken(stringValue("final")); + stateMachine.addPartialResultSet(resultWithToken.getResults()); + assertThat(stateMachine) + .populateQueueYields( + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("foo"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("bar"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("baz"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("a"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("b"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("c"))), + ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("final")))); + } + + @Test + public void stateMachine_throwsExceptionWithChecksumButNoData() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + + PartialResultSet invalid = PartialResultSet.newBuilder().setBatchChecksum(1234).build(); + assertThrows(IllegalStateException.class, () -> stateMachine.addPartialResultSet(invalid)); + } + + @Test + public void stateMachine_resolvesMetadataOnlyAfterFirstToken() { + final boolean[] metadataHasBeenAccessed = {false}; + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + ProtoRowsMergingStateMachine stateMachine = + new ProtoRowsMergingStateMachine( + () -> { + // hacky way to check if supplier has been resolved + // This is in an array so the variable can be final + metadataHasBeenAccessed[0] = true; + return metadata; + }); + + stateMachine.addPartialResultSet(partialResultSetWithoutToken(stringValue("s")).getResults()); + assertFalse(metadataHasBeenAccessed[0]); + stateMachine.addPartialResultSet(partialResultSetWithToken(stringValue("b")).getResults()); + assertTrue(metadataHasBeenAccessed[0]); + } + + @Test + public void stateMachine_handlesSchemaChangeAfterResetOfInitialBatch() { + SettableApiFuture mdFuture = SettableApiFuture.create(); + ProtoRowsMergingStateMachine stateMachine = + new ProtoRowsMergingStateMachine( + () -> ApiExceptions.callAndTranslateApiException(mdFuture)); + stateMachine.addPartialResultSet( + partialResultSetWithoutToken(stringValue("discard")).getResults()); + + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto( + metadata(columnMetadata("a", bytesType()), columnMetadata("b", int64Type()))); + mdFuture.set(metadata); + List retryResponses = + partialResultSets(2, bytesValue("bytes"), int64Value(123)); + for (ExecuteQueryResponse res : retryResponses) { + stateMachine.addPartialResultSet(res.getResults()); + } + assertThat(stateMachine) + .populateQueueYields( + ProtoSqlRow.create(metadata, ImmutableList.of(bytesValue("bytes"), int64Value(123)))); + } + + @Test + public void stateMachine_throwsExceptionWithTokenAndIncompleteBatch() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + + List responses = + partialResultSets(2, stringValue("foo"), stringValue("bar")); + stateMachine.addPartialResultSet(responses.get(0).getResults()); + // We haven't added the second response above, this should error + assertThrows( + IllegalStateException.class, + () -> + stateMachine.addPartialResultSet( + tokenOnlyResultSet(ByteString.copyFromUtf8("token")).getResults())); + } + + @Test + public void isBatchInProgress_trueWithUncommitedCompleteBatches() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + + stateMachine.addPartialResultSet( + partialResultSetWithoutToken(stringValue("foo")).getResults()); + assertThat(stateMachine).isBatchInProgress(true); + } + + @Test + public void hasCompleteBatches_falseWithUncommitedCompleteBatches() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + ProtoRowsMergingStateMachine stateMachine = new ProtoRowsMergingStateMachine(() -> metadata); + + stateMachine.addPartialResultSet( + partialResultSetWithoutToken(stringValue("foo")).getResults()); + assertThat(stateMachine).hasCompleteBatches(false); + } } @RunWith(Parameterized.class) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java index 4402af5ba9..c3643789a3 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java @@ -27,6 +27,9 @@ import com.google.bigtable.v2.Type; import com.google.bigtable.v2.Type.Struct.Field; import com.google.bigtable.v2.Value; +import com.google.common.collect.ImmutableList; +import com.google.common.hash.HashFunction; +import com.google.common.hash.Hashing; import com.google.protobuf.ByteString; import com.google.protobuf.Timestamp; import com.google.type.Date; @@ -35,6 +38,8 @@ /** Utilities for creating sql proto objects in tests */ public class SqlProtoFactory { + private static final HashFunction CRC32C = Hashing.crc32c(); + private SqlProtoFactory() {} public static PrepareQueryResponse prepareResponse( @@ -177,25 +182,14 @@ public static Value mapElement(Value... fields) { return structValue(fields); } - private static ProtoRowsBatch protoRowsBatch(Value... values) { - ProtoRows protoRows = ProtoRows.newBuilder().addAllValues(Arrays.asList(values)).build(); - return ProtoRowsBatch.newBuilder().setBatchData(protoRows.toByteString()).build(); - } - + /** Creates a single response representing a complete batch, with no token */ public static ExecuteQueryResponse partialResultSetWithoutToken(Value... values) { - return ExecuteQueryResponse.newBuilder() - .setResults(PartialResultSet.newBuilder().setProtoRowsBatch(protoRowsBatch(values)).build()) - .build(); + return partialResultSets(1, false, ByteString.EMPTY, values).get(0); } + /** Creates a single response representing a complete batch, with a resume token of 'test' */ public static ExecuteQueryResponse partialResultSetWithToken(Value... values) { - return ExecuteQueryResponse.newBuilder() - .setResults( - PartialResultSet.newBuilder() - .setProtoRowsBatch(protoRowsBatch(values)) - .setResumeToken(ByteString.copyFromUtf8("test")) - .build()) - .build(); + return partialResultSets(1, false, ByteString.copyFromUtf8("test"), values).get(0); } public static ExecuteQueryResponse tokenOnlyResultSet(ByteString token) { @@ -204,9 +198,58 @@ public static ExecuteQueryResponse tokenOnlyResultSet(ByteString token) { .build(); } + /** + * splits values across specified number of batches. Sets reset on first response, and resume + * token on final response + */ + public static ImmutableList partialResultSets( + int batches, Value... values) { + return partialResultSets(batches, true, ByteString.copyFromUtf8("test"), values); + } + + /** + * @param batches number of {@link ProtoRowsBatch}s to split values across + * @param reset whether to set the reset bit on the first response + * @param resumeToken resumption token for the final response. Unset if empty + * @param values List of values to split across batches + * @return List of responses with length equal to number of batches + */ + public static ImmutableList partialResultSets( + int batches, boolean reset, ByteString resumeToken, Value... values) { + ProtoRows protoRows = ProtoRows.newBuilder().addAllValues(Arrays.asList(values)).build(); + ByteString batchData = protoRows.toByteString(); + int batch_checksum = checksum(batchData); + ImmutableList.Builder responses = ImmutableList.builder(); + int batchSize = batchData.size() / batches; + for (int i = 0; i < batches; i++) { + boolean finalBatch = i == batches - 1; + int batchStart = i * batchSize; + int batchEnd = finalBatch ? batchData.size() : batchStart + batchSize; + ProtoRowsBatch.Builder batchBuilder = ProtoRowsBatch.newBuilder(); + batchBuilder.setBatchData(batchData.substring(batchStart, batchEnd)); + PartialResultSet.Builder resultSetBuilder = PartialResultSet.newBuilder(); + if (reset && i == 0) { + resultSetBuilder.setReset(true); + } + if (finalBatch) { + resultSetBuilder.setBatchChecksum(batch_checksum); + if (!resumeToken.isEmpty()) { + resultSetBuilder.setResumeToken(resumeToken); + } + } + resultSetBuilder.setProtoRowsBatch(batchBuilder.build()); + responses.add(ExecuteQueryResponse.newBuilder().setResults(resultSetBuilder.build()).build()); + } + return responses.build(); + } + public static ResultSetMetadata metadata(ColumnMetadata... columnMetadata) { ProtoSchema schema = ProtoSchema.newBuilder().addAllColumns(Arrays.asList(columnMetadata)).build(); return ResultSetMetadata.newBuilder().setProtoSchema(schema).build(); } + + public static int checksum(ByteString bytes) { + return CRC32C.hashBytes(bytes.toByteArray()).asInt(); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java index 050f794e98..d61d9d5f20 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergerTest.java @@ -27,6 +27,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSets; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.tokenOnlyResultSet; @@ -34,6 +35,7 @@ import static org.junit.Assert.assertThrows; import com.google.bigtable.v2.ExecuteQueryResponse; +import com.google.bigtable.v2.PartialResultSet; import com.google.bigtable.v2.Value; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.ProtoSqlRow; @@ -41,6 +43,7 @@ import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.util.Arrays; +import java.util.List; import java.util.function.Supplier; import org.junit.Test; import org.junit.runner.RunWith; @@ -91,7 +94,21 @@ public void sqlRowMerger_doesntResolveMetadataUntilFirstPush() { } @Test - public void hasPartialFrame_trueWithIncompleteBatch() { + public void hasPartialFrame_trueWithPartialBatch() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + SqlRowMerger merger = new SqlRowMerger(() -> metadata); + // Initial response here has reset bit set + List responses = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + + merger.push(responses.get(0)); + merger.push(responses.get(1)); + assertThat(merger).hasPartialFrame(true); + } + + @Test + public void hasPartialFrame_trueWithUncommittedBatch() { com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); @@ -130,6 +147,20 @@ public void hasFullFrame_trueWithFullRow() { @Test public void hasFullFrame_falseWithIncompleteBatch() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + SqlRowMerger merger = new SqlRowMerger(() -> metadata); + // Initial response here has reset bit set + List responses = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + + merger.push(responses.get(0)); + merger.push(responses.get(1)); + assertThat(merger).hasFullFrame(false); + } + + @Test + public void hasFullFrame_falseWithUncommittedBatches() { com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); @@ -178,17 +209,52 @@ public void sqlRowMerger_handlesResponseStream() { } @Test - public void addValue_failsWithoutMetadataFirst() { - com.google.bigtable.v2.ResultSetMetadata metadataProto = - metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); - SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); - assertThrows( - IllegalStateException.class, - () -> merger.push(partialResultSetWithToken(stringValue("test")))); + public void sqlRowMerger_handlesReset() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + SqlRowMerger merger = new SqlRowMerger(() -> metadata); + // Initial response here has reset bit set + List responses = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + + merger.push(responses.get(0)); + merger.push(responses.get(1)); + assertThat(merger).hasPartialFrame(true); + assertThat(merger).hasFullFrame(false); + + for (ExecuteQueryResponse res : responses) { + merger.push(res); + } + assertThat(merger).hasFullFrame(true); + assertThat(merger.pop()) + .isEqualTo(ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("foo")))); + assertThat(merger.pop()) + .isEqualTo(ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("bar")))); + assertThat(merger.pop()) + .isEqualTo(ProtoSqlRow.create(metadata, ImmutableList.of(stringValue("baz")))); + assertThat(merger).hasFullFrame(false); + } + + @Test + public void sqlRowMerger_throwsExceptionOnChecksumMismatch() { + ResultSetMetadata metadata = + ProtoResultSetMetadata.fromProto(metadata(columnMetadata("a", stringType()))); + SqlRowMerger merger = new SqlRowMerger(() -> metadata); + List responses = + partialResultSets(3, stringValue("foo"), stringValue("bar"), stringValue("baz")); + + // Override the checksum of the final response + PartialResultSet lastResultsWithBadChecksum = + responses.get(2).getResults().toBuilder().setBatchChecksum(1234).build(); + ExecuteQueryResponse badChecksum = + ExecuteQueryResponse.newBuilder().setResults(lastResultsWithBadChecksum).build(); + merger.push(responses.get(0)); + merger.push(responses.get(1)); + assertThrows(IllegalStateException.class, () -> merger.push(badChecksum)); } @Test - public void sqlRowMerger_handlesTokenWithOpenPartialBatch() { + public void sqlRowMerger_handlesTokenWithUncommittedBatches() { com.google.bigtable.v2.ResultSetMetadata metadataProto = metadata(columnMetadata("str", stringType()), columnMetadata("bytes", bytesType())); SqlRowMerger merger = new SqlRowMerger(toSupplier(metadataProto)); From 2fe674ea56fc3583d2b5e3e22812022f5fd32f58 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Mon, 17 Feb 2025 16:37:57 -0500 Subject: [PATCH 06/11] Implement streaming retries for executeQuery Renames MetadataResolvingCallable to PlanRefreshingCallable as this will be the place we handle plan refresh & metadata resolution Change-Id: I6d5329ced6d9a4228a3a49fbc8b5f497d8678cce Change-Id: I92dcde9e67459e0b8119b4c9d216a47fb6819758 --- .../data/v2/models/sql/BoundStatement.java | 24 +- .../data/v2/stub/EnhancedBigtableStub.java | 72 ++- .../v2/stub/EnhancedBigtableStubSettings.java | 26 +- .../v2/stub/sql/ExecuteQueryCallContext.java | 11 +- .../sql/ExecuteQueryResumptionStrategy.java | 61 +++ .../sql/MetadataErrorHandlingCallable.java | 88 ++++ ...lable.java => PlanRefreshingCallable.java} | 14 +- .../v2/models/sql/BoundStatementTest.java | 63 ++- .../EnhancedBigtableStubSettingsTest.java | 19 +- .../v2/stub/sql/ExecuteQueryCallableTest.java | 58 +-- .../ExecuteQueryResumptionStrategyTest.java | 72 +++ .../v2/stub/sql/ExecuteQueryRetryTest.java | 441 ++++++++++++++++++ .../MetadataErrorHandlingCallableTest.java | 87 ++++ ...t.java => PlanRefreshingCallableTest.java} | 55 +-- .../data/v2/stub/sql/SqlProtoFactory.java | 5 + 15 files changed, 897 insertions(+), 199 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategy.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallable.java rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/{MetadataResolvingCallable.java => PlanRefreshingCallable.java} (91%) create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java rename google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/{MetadataResolvingCallableTest.java => PlanRefreshingCallableTest.java} (75%) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java index d277126f6c..4692923d0b 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java @@ -373,14 +373,20 @@ private static Value.Builder nullValueWithType(Type type) { * not meant to be used by applications. */ @InternalApi("For internal use only") - public ExecuteQueryRequest toProto(ByteString preparedQuery, RequestContext requestContext) { - return ExecuteQueryRequest.newBuilder() - .setInstanceName( - NameUtil.formatInstanceName( - requestContext.getProjectId(), requestContext.getInstanceId())) - .setAppProfileId(requestContext.getAppProfileId()) - .setPreparedQuery(preparedQuery) - .putAllParams(params) - .build(); + public ExecuteQueryRequest toProto( + ByteString preparedQuery, RequestContext requestContext, @Nullable ByteString resumeToken) { + ExecuteQueryRequest.Builder requestBuilder = + ExecuteQueryRequest.newBuilder() + .setInstanceName( + NameUtil.formatInstanceName( + requestContext.getProjectId(), requestContext.getInstanceId())) + .setAppProfileId(requestContext.getAppProfileId()) + .setPreparedQuery(preparedQuery) + .putAllParams(params); + + if (resumeToken != null) { + requestBuilder.setResumeToken(resumeToken); + } + return requestBuilder.build(); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 966263c1ac..2dfbdc181c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -45,7 +45,6 @@ import com.google.api.gax.rpc.RequestParamsExtractor; import com.google.api.gax.rpc.ServerStreamingCallSettings; import com.google.api.gax.rpc.ServerStreamingCallable; -import com.google.api.gax.rpc.StatusCode.Code; import com.google.api.gax.rpc.UnaryCallSettings; import com.google.api.gax.rpc.UnaryCallable; import com.google.api.gax.tracing.ApiTracerFactory; @@ -121,7 +120,9 @@ import com.google.cloud.bigtable.data.v2.stub.readrows.RowMergingCallable; import com.google.cloud.bigtable.data.v2.stub.sql.ExecuteQueryCallContext; import com.google.cloud.bigtable.data.v2.stub.sql.ExecuteQueryCallable; -import com.google.cloud.bigtable.data.v2.stub.sql.MetadataResolvingCallable; +import com.google.cloud.bigtable.data.v2.stub.sql.ExecuteQueryResumptionStrategy; +import com.google.cloud.bigtable.data.v2.stub.sql.MetadataErrorHandlingCallable; +import com.google.cloud.bigtable.data.v2.stub.sql.PlanRefreshingCallable; import com.google.cloud.bigtable.data.v2.stub.sql.SqlRowMergingCallable; import com.google.cloud.bigtable.gaxx.retrying.ApiResultRetryAlgorithm; import com.google.cloud.bigtable.gaxx.retrying.RetryInfoRetryAlgorithm; @@ -144,10 +145,8 @@ import io.opentelemetry.api.common.Attributes; import java.io.IOException; import java.time.Duration; -import java.util.Collections; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.function.Function; import javax.annotation.Nonnull; @@ -1151,9 +1150,6 @@ private UnaryCallable createReadModifyWriteRowCallable( */ @InternalApi("For internal use only") public ExecuteQueryCallable createExecuteQueryCallable() { - // TODO support resumption - // TODO update codes once resumption is implemented - Set retryableCodes = Collections.emptySet(); ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( GrpcCallSettings.newBuilder() @@ -1168,54 +1164,54 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { } }) .build(), - retryableCodes); + settings.executeQuerySettings().getRetryableCodes()); ServerStreamingCallable withStatsHeaders = new StatsHeadersServerStreamingCallable<>(base); - ServerStreamingCallSettings watchdogSettings = - ServerStreamingCallSettings.newBuilder() + ServerStreamingCallable withMetadataObserver = + new PlanRefreshingCallable(withStatsHeaders, requestContext); + + ServerStreamingCallSettings retrySettings = + ServerStreamingCallSettings.newBuilder() + .setResumptionStrategy(new ExecuteQueryResumptionStrategy()) + .setRetryableCodes(settings.executeQuerySettings().getRetryableCodes()) + .setRetrySettings(settings.executeQuerySettings().getRetrySettings()) .setIdleTimeout(settings.executeQuerySettings().getIdleTimeout()) .setWaitTimeout(settings.executeQuerySettings().getWaitTimeout()) .build(); - // Watchdog needs to stay above the metadata observer so that watchdog errors - // are passed through to the metadata future. - ServerStreamingCallable watched = - Callables.watched(withStatsHeaders, watchdogSettings, clientContext); - - ServerStreamingCallable withMetadataObserver = - new MetadataResolvingCallable(watched, requestContext); + // Retries need to happen before row merging, because the resumeToken is part + // of the ExecuteQueryResponse. This is okay because the first response in every + // attempt stream will have reset set to true, so any unyielded data from the previous + // attempt will be reset properly + ServerStreamingCallable retries = + withRetries(withMetadataObserver, retrySettings); ServerStreamingCallable merging = - new SqlRowMergingCallable(withMetadataObserver); - - ServerStreamingCallable withBigtableTracer = - new BigtableTracerStreamingCallable<>(merging); + new SqlRowMergingCallable(retries); - ServerStreamingCallSettings retrySettings = + ServerStreamingCallSettings watchdogSettings = ServerStreamingCallSettings.newBuilder() - // TODO add resumption strategy and pass through retry settings unchanged - // we pass through retry settings to use the deadlines now but don't - // support retries - .setRetrySettings( - settings - .executeQuerySettings() - .getRetrySettings() - .toBuilder() - // override maxAttempts as a safeguard against changes from user - .setMaxAttempts(1) - .build()) + .setIdleTimeout(settings.executeQuerySettings().getIdleTimeout()) + .setWaitTimeout(settings.executeQuerySettings().getWaitTimeout()) .build(); - // Adding RetryingCallable to the callable chain so that client side metrics can be - // measured correctly and deadlines are set. Retries are currently disabled. - ServerStreamingCallable retries = - withRetries(withBigtableTracer, retrySettings); + // Watchdog needs to stay above the metadata error handling so that watchdog errors + // are passed through to the metadata future. + ServerStreamingCallable watched = + Callables.watched(merging, watchdogSettings, clientContext); + + ServerStreamingCallable passingThroughErrorsToMetadata = + new MetadataErrorHandlingCallable(watched); + + ServerStreamingCallable withBigtableTracer = + new BigtableTracerStreamingCallable<>(passingThroughErrorsToMetadata); SpanName span = getSpanName("ExecuteQuery"); ServerStreamingCallable traced = - new TracedServerStreamingCallable<>(retries, clientContext.getTracerFactory(), span); + new TracedServerStreamingCallable<>( + withBigtableTracer, clientContext.getTracerFactory(), span); return new ExecuteQueryCallable( traced.withDefaultCallContext( diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index 218fb682ad..d756a52370 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -61,7 +61,6 @@ import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.Base64; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; @@ -201,21 +200,27 @@ public class EnhancedBigtableStubSettings extends StubSettings EXECUTE_QUERY_RETRY_CODES = Collections.emptySet(); + // Allow retrying ABORTED statuses. These will be returned by the server when the client is + // too slow to read the responses. + private static final Set EXECUTE_QUERY_RETRY_CODES = + ImmutableSet.builder().addAll(IDEMPOTENT_RETRY_CODES).add(Code.ABORTED).build(); - // We still setup retry settings in order to set default deadlines + // We use the same configuration as READ_ROWS private static final RetrySettings EXECUTE_QUERY_RETRY_SETTINGS = RetrySettings.newBuilder() - .setMaxAttempts(1) - // Set a conservative deadline to start for preview. We'll increase this in the future - .setInitialRpcTimeout(Duration.ofSeconds(30)) - .setMaxRpcTimeout(Duration.ofSeconds(30)) + .setInitialRetryDelay(Duration.ofMillis(10)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelay(Duration.ofMinutes(1)) + .setMaxAttempts(10) + .setJittered(true) + .setInitialRpcTimeout(Duration.ofMinutes(30)) + .setRpcTimeoutMultiplier(1.0) + .setMaxRpcTimeout(Duration.ofMinutes(30)) + .setTotalTimeout(Duration.ofHours(12)) .build(); // Similar to IDEMPOTENT but with a lower initial rpc timeout since we expect - // these calls to be quick in most circumestances + // these calls to be quick in most circumstances private static final RetrySettings PREPARE_QUERY_RETRY_SETTINGS = RetrySettings.newBuilder() .setInitialRetryDelay(Duration.ofMillis(10)) @@ -884,7 +889,6 @@ private Builder() { executeQuerySettings = ServerStreamingCallSettings.newBuilder(); executeQuerySettings .setRetryableCodes(EXECUTE_QUERY_RETRY_CODES) - // This is used to set deadlines. We do not support retries yet. .setRetrySettings(EXECUTE_QUERY_RETRY_SETTINGS) .setIdleTimeout(Duration.ofMinutes(5)) .setWaitTimeout(Duration.ofMinutes(5)); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java index 33cecc9c85..337dc2d497 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java @@ -22,6 +22,8 @@ import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.protobuf.ByteString; +import javax.annotation.Nullable; /** * Used to provide a future to the ExecuteQuery callable chain in order to return metadata to users @@ -37,7 +39,7 @@ public class ExecuteQueryCallContext { private final BoundStatement boundStatement; private final SettableApiFuture metadataFuture; private final PrepareResponse latestPrepareResponse; - // TODO this will be used to track latest resume token here + private @Nullable ByteString resumeToken; private ExecuteQueryCallContext( BoundStatement boundStatement, SettableApiFuture metadataFuture) { @@ -52,7 +54,8 @@ public static ExecuteQueryCallContext create( } ExecuteQueryRequest toRequest(RequestContext requestContext) { - return boundStatement.toProto(latestPrepareResponse.preparedQuery(), requestContext); + return boundStatement.toProto( + latestPrepareResponse.preparedQuery(), requestContext, resumeToken); } /** @@ -75,4 +78,8 @@ void setMetadataException(Throwable t) { SettableApiFuture resultSetMetadataFuture() { return this.metadataFuture; } + + void setLatestResumeToken(ByteString resumeToken) { + this.resumeToken = resumeToken; + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategy.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategy.java new file mode 100644 index 0000000000..e6e2562c33 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategy.java @@ -0,0 +1,61 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.stub.sql; + +import com.google.api.core.InternalApi; +import com.google.api.gax.retrying.StreamResumptionStrategy; +import com.google.bigtable.v2.ExecuteQueryResponse; +import com.google.protobuf.ByteString; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +@InternalApi +public class ExecuteQueryResumptionStrategy + implements StreamResumptionStrategy { + + private ByteString latestResumeToken = null; + + @Nonnull + @Override + public StreamResumptionStrategy createNew() { + return new ExecuteQueryResumptionStrategy(); + } + + @Nonnull + @Override + public ExecuteQueryResponse processResponse(ExecuteQueryResponse response) { + if (!response.getResults().getResumeToken().isEmpty()) { + latestResumeToken = response.getResults().getResumeToken(); + } + return response; + } + + @Nullable + @Override + public ExecuteQueryCallContext getResumeRequest(ExecuteQueryCallContext originalRequest) { + if (latestResumeToken != null) { + // ExecuteQueryCallContext can handle null token, but we don't bother setting it for + // clarity + originalRequest.setLatestResumeToken(latestResumeToken); + } + return originalRequest; + } + + @Override + public boolean canResume() { + return true; + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallable.java new file mode 100644 index 0000000000..e36bfa57fc --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallable.java @@ -0,0 +1,88 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.stub.sql; + +import com.google.api.core.InternalApi; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.api.gax.rpc.StreamController; +import com.google.cloud.bigtable.data.v2.internal.SqlRow; +import com.google.cloud.bigtable.data.v2.stub.SafeResponseObserver; + +/** + * Callable that handles passing execeptions through to the metadata future. This needs to be used + * after all retries, so that non-retriable errors don't surface as errors to users accessing the + * metadata. + * + *

      In non-error cases the metadata future is resolved by the {@link PlanRefreshingCallable} + * because the metadata needs to resolve before the SqlRowMerger starts yielding rows + * + *

      This is considered an internal implementation detail and should not be used by applications. + */ +@InternalApi("For internal use only") +public class MetadataErrorHandlingCallable + extends ServerStreamingCallable { + private final ServerStreamingCallable inner; + + public MetadataErrorHandlingCallable( + ServerStreamingCallable inner) { + this.inner = inner; + } + + @Override + public void call( + ExecuteQueryCallContext request, + ResponseObserver responseObserver, + ApiCallContext context) { + MetadataErrorHandlingObserver observer = + new MetadataErrorHandlingObserver(responseObserver, request); + inner.call(request, observer, context); + } + + static final class MetadataErrorHandlingObserver extends SafeResponseObserver { + private final ExecuteQueryCallContext callContext; + private final ResponseObserver outerObserver; + + MetadataErrorHandlingObserver( + ResponseObserver outerObserver, ExecuteQueryCallContext callContext) { + super(outerObserver); + this.outerObserver = outerObserver; + this.callContext = callContext; + } + + @Override + protected void onStartImpl(StreamController streamController) { + outerObserver.onStart(streamController); + } + + @Override + protected void onResponseImpl(SqlRow response) { + outerObserver.onResponse(response); + } + + @Override + protected void onErrorImpl(Throwable throwable) { + callContext.setMetadataException(throwable); + outerObserver.onError(throwable); + } + + @Override + protected void onCompleteImpl() { + outerObserver.onComplete(); + } + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java similarity index 91% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java index 8b8f850d5f..ba15e1368e 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java @@ -33,12 +33,12 @@ *

      This is considered an internal implementation detail and should not be used by applications. */ @InternalApi("For internal use only") -public class MetadataResolvingCallable +public class PlanRefreshingCallable extends ServerStreamingCallable { private final ServerStreamingCallable inner; private final RequestContext requestContext; - public MetadataResolvingCallable( + public PlanRefreshingCallable( ServerStreamingCallable inner, RequestContext requestContext) { this.inner = inner; @@ -51,6 +51,8 @@ public void call( ResponseObserver responseObserver, ApiCallContext apiCallContext) { MetadataObserver observer = new MetadataObserver(responseObserver, callContext); + // TODO toRequest will return a future. We need to timeout waiting for future + // based on the totalTimeout inner.call(callContext.toRequest(requestContext), observer, apiCallContext); } @@ -101,9 +103,11 @@ protected void onResponseImpl(ExecuteQueryResponse response) { @Override protected void onErrorImpl(Throwable throwable) { - // When we support retries this will have to move after the retrying callable in a separate - // observer. - callContext.setMetadataException(throwable); + // TODO translate plan refresh errors and trigger plan refresh + + // Note that we do not set exceptions on the metadata future here. This + // needs to be done after the retries, so that retryable errors aren't set on + // the future outerObserver.onError(throwable); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java index 613c5df365..f3a37f99d5 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java @@ -53,6 +53,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import javax.annotation.Nullable; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -71,6 +72,7 @@ public class BoundStatementTest { private static final ColumnMetadata[] DEFAULT_COLUMNS = { columnMetadata("_key", bytesType()), columnMetadata("cf", stringType()) }; + private static final @Nullable ByteString NO_RESUME_TOKEN = null; // Use ColumnMetadata as a more concise way of specifying params public static BoundStatement.Builder boundStatementBuilder(ColumnMetadata... paramColumns) { @@ -97,7 +99,7 @@ public static BoundStatement.Builder boundStatementBuilder(ColumnMetadata... par public void statementWithoutParameters() { BoundStatement s = boundStatementBuilder().build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -106,6 +108,21 @@ public void statementWithoutParameters() { .build()); } + @Test + public void statementWithResumeToken() { + BoundStatement s = boundStatementBuilder().build(); + + assertThat( + s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, ByteString.copyFromUtf8("token"))) + .isEqualTo( + ExecuteQueryRequest.newBuilder() + .setPreparedQuery(EXPECTED_PREPARED_QUERY) + .setInstanceName(EXPECTED_INSTANCE_NAME) + .setAppProfileId(EXPECTED_APP_PROFILE) + .setResumeToken(ByteString.copyFromUtf8("token")) + .build()); + } + @Test public void statementWithBytesParam() { BoundStatement s = @@ -113,7 +130,7 @@ public void statementWithBytesParam() { .setBytesParam("key", ByteString.copyFromUtf8("test")) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -135,7 +152,7 @@ public void statementWithNullBytesParam() { .setBytesParam("key", null) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -152,7 +169,7 @@ public void statementWithStringParam() { .setStringParam("key", "test") .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -170,7 +187,7 @@ public void statementWithNullStringParam() { .setStringParam("key", null) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -187,7 +204,7 @@ public void statementWithInt64Param() { .setLongParam("number", 1L) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -204,7 +221,7 @@ public void statementWithNullInt64Param() { .setLongParam("number", null) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -221,7 +238,7 @@ public void statementWithBoolParam() { .setBooleanParam("bool", true) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -239,7 +256,7 @@ public void statementWithNullBoolParam() { .setBooleanParam("bool", null) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -256,7 +273,7 @@ public void statementWithTimestampParam() { .setTimestampParam("timeParam", Instant.ofEpochSecond(1000, 100)) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -278,7 +295,7 @@ public void statementWithNullTimestampParam() { .setTimestampParam("timeParam", null) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -295,7 +312,7 @@ public void statementWithDateParam() { .setDateParam("dateParam", Date.fromYearMonthDay(2024, 6, 11)) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -317,7 +334,7 @@ public void statementWithNullDateParam() { .setDateParam("dateParam", null) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -347,7 +364,7 @@ public void statementWithBytesListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.bytes())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -396,7 +413,7 @@ public void statementWithStringListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.string())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -442,7 +459,7 @@ public void statementWithInt64ListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.int64())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -489,7 +506,7 @@ public void statementWithFloat32ListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.float32())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -538,7 +555,7 @@ public void statementWithFloat64ListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.float64())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -585,7 +602,7 @@ public void statementWithBooleanListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.bool())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -638,7 +655,7 @@ public void statementWithTimestampListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.timestamp())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -695,7 +712,7 @@ public void statementWithDateListParam() { .setListParam("nullList", null, SqlType.arrayOf(SqlType.date())) .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -756,7 +773,7 @@ public void statementBuilderAllowsParamsToBeOverridden() { .setStringParam("key", "test4") .build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) @@ -771,7 +788,7 @@ public void statementBuilderAllowsParamsToBeOverridden() { public void builderWorksWithNoParams() { BoundStatement s = boundStatementBuilder().build(); - assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT)) + assertThat(s.toProto(EXPECTED_PREPARED_QUERY, REQUEST_CONTEXT, NO_RESUME_TOKEN)) .isEqualTo( ExecuteQueryRequest.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java index 1e185a7038..34c2fbd5a8 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java @@ -45,7 +45,6 @@ import java.lang.reflect.Modifier; import java.net.URI; import java.util.Arrays; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Set; @@ -833,19 +832,11 @@ public void executeQueryHasSaneDefaults() { // Retries aren't supported right now // call verifyRetrySettingAreSane when we do - assertThat(builder.getRetryableCodes()).containsExactlyElementsIn(Collections.emptySet()); - assertThat(builder.getRetrySettings().getInitialRpcTimeout()).isEqualTo(Duration.ofSeconds(30)); - assertThat(builder.getRetrySettings().getMaxRpcTimeout()).isEqualTo(Duration.ofSeconds(30)); - assertThat(builder.getRetrySettings().getMaxAttempts()).isEqualTo(1); - } - - @Test - public void executeQueryRetriesAreDisabled() { - ServerStreamingCallSettings.Builder builder = - EnhancedBigtableStubSettings.newBuilder().executeQuerySettings(); - - assertThat(builder.getRetrySettings().getMaxAttempts()).isAtMost(1); - assertThat(builder.getRetrySettings().getInitialRpcTimeout()).isAtMost(Duration.ofSeconds(30)); + assertThat(builder.getRetryableCodes()) + .containsAtLeast(Code.ABORTED, Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE); + assertThat(builder.getRetrySettings().getInitialRpcTimeout()).isEqualTo(Duration.ofMinutes(30)); + assertThat(builder.getRetrySettings().getMaxRpcTimeout()).isEqualTo(Duration.ofMinutes(30)); + assertThat(builder.getRetrySettings().getMaxAttempts()).isEqualTo(10); } @Test diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java index ecf0992063..62286368a7 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java @@ -17,6 +17,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; @@ -26,7 +27,6 @@ import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.DeadlineExceededException; -import com.google.api.gax.rpc.UnavailableException; import com.google.bigtable.v2.BigtableGrpc; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; @@ -45,8 +45,6 @@ import io.grpc.Context; import io.grpc.Deadline; import io.grpc.Server; -import io.grpc.Status; -import io.grpc.StatusRuntimeException; import io.grpc.stub.StreamObserver; import java.io.IOException; import java.time.Duration; @@ -125,49 +123,18 @@ public void testCallContextAndServerStreamSetup() { assertThat(responseIterator.hasNext()).isFalse(); } - @Test - public void testExecuteQueryRequestsAreNotRetried() { - // TODO: retries for execute query is currently disabled. This test should be - // updated once resumption token is in place. - SqlServerStream stream = stub.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); - - Iterator iterator = stream.rows().iterator(); - - assertThrows(UnavailableException.class, iterator::next).getCause(); - assertThat(fakeService.attempts).isEqualTo(1); - } - - @Test - public void testExecuteQueryRequestsIgnoreOverriddenMaxAttempts() throws IOException { - BigtableDataSettings.Builder overrideSettings = - BigtableDataSettings.newBuilderForEmulator(server.getPort()) - .setProjectId("fake-project") - .setInstanceId("fake-instance"); - overrideSettings - .stubSettings() - .executeQuerySettings() - .setRetrySettings(RetrySettings.newBuilder().setMaxAttempts(10).build()); - - try (EnhancedBigtableStub overrideStub = - EnhancedBigtableStub.create(overrideSettings.build().getStubSettings())) { - SqlServerStream stream = - overrideStub.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); - Iterator iterator = stream.rows().iterator(); - - assertThrows(UnavailableException.class, iterator::next).getCause(); - assertThat(fakeService.attempts).isEqualTo(1); - } - } - @Test public void testExecuteQueryRequestsSetDefaultDeadline() { SqlServerStream stream = stub.executeQueryCallable().call(PREPARED_STATEMENT.bind().build()); - Iterator iterator = stream.rows().iterator(); - // We don't care about this but are reusing the fake service that tests retries - assertThrows(UnavailableException.class, iterator::next).getCause(); - // We have 30s default, we give it a wide range to avoid flakiness, this is mostly just checking - // that some default is set - assertThat(fakeService.deadlineMillisRemaining).isLessThan(30001L); + // We don't care about this, just assert we get a response + boolean rowReceived = false; + for (SqlRow sqlRow : stream.rows()) { + rowReceived = true; + } + assertThat(rowReceived).isTrue(); + // We have 30m default, we give it a wide range to avoid flakiness, this is mostly just + // checking that some default is set + assertThat(fakeService.deadlineMillisRemaining).isLessThan(1800000L); } @Test @@ -197,7 +164,6 @@ public void testExecuteQueryRequestsRespectDeadline() throws IOException { private static class FakeService extends BigtableGrpc.BigtableImplBase { - private int attempts = 0; private long deadlineMillisRemaining; @Override @@ -216,9 +182,9 @@ public void executeQuery( } catch (InterruptedException e) { throw new RuntimeException(e); } - attempts++; responseObserver.onNext(partialResultSetWithoutToken(stringValue("foo"))); - responseObserver.onError(new StatusRuntimeException(Status.UNAVAILABLE)); + responseObserver.onNext(partialResultSetWithToken(stringValue("bar"))); + responseObserver.onCompleted(); } } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java new file mode 100644 index 0000000000..7bd860115b --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java @@ -0,0 +1,72 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.stub.sql; + +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; +import static com.google.common.truth.extensions.proto.ProtoTruth.assertThat; + +import com.google.api.core.SettableApiFuture; +import com.google.bigtable.v2.ExecuteQueryRequest; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.protobuf.ByteString; +import java.util.HashMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ExecuteQueryResumptionStrategyTest { + + @Test + public void tracksResumeToken() { + ExecuteQueryResumptionStrategy resumptionStrategy = new ExecuteQueryResumptionStrategy(); + + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("s", stringType())))), + new HashMap<>()); + SettableApiFuture mdFuture = SettableApiFuture.create(); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(preparedStatement.bind().build(), mdFuture); + + resumptionStrategy.processResponse( + partialResultSetWithToken(ByteString.copyFromUtf8("token"), stringValue("s"))); + // This should not change the token + resumptionStrategy.processResponse(partialResultSetWithoutToken(stringValue("bar"))); + + ExecuteQueryCallContext updatedCallContext = resumptionStrategy.getResumeRequest(callContext); + assertThat( + updatedCallContext.toRequest(RequestContext.create("project", "instance", "profile"))) + .isEqualTo( + ExecuteQueryRequest.newBuilder() + .setInstanceName(NameUtil.formatInstanceName("project", "instance")) + .setAppProfileId("profile") + .setPreparedQuery(ByteString.copyFromUtf8("foo")) + .setResumeToken(ByteString.copyFromUtf8("token")) + .build()); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java new file mode 100644 index 0000000000..ae295c1720 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java @@ -0,0 +1,441 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.stub.sql; + +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSets; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.tokenOnlyResultSet; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.api.gax.rpc.StatusCode; +import com.google.bigtable.v2.BigtableGrpc; +import com.google.bigtable.v2.ExecuteQueryRequest; +import com.google.bigtable.v2.ExecuteQueryResponse; +import com.google.bigtable.v2.ResultSetMetadata; +import com.google.bigtable.v2.Value; +import com.google.cloud.bigtable.data.v2.BigtableDataClient; +import com.google.cloud.bigtable.data.v2.BigtableDataSettings; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; +import com.google.cloud.bigtable.gaxx.reframing.IncompleteStreamException; +import com.google.common.collect.ImmutableMap; +import com.google.common.truth.Truth; +import com.google.protobuf.ByteString; +import io.grpc.Status; +import io.grpc.Status.Code; +import io.grpc.stub.StreamObserver; +import io.grpc.testing.GrpcServerRule; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.LinkedBlockingDeque; +import javax.annotation.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ExecuteQueryRetryTest { + private static final String PROJECT_ID = "fake-project"; + private static final String INSTANCE_ID = "fake-instance"; + private static final String APP_PROFILE_ID = "fake-app-profile"; + private static final ByteString PREPARED_QUERY = ByteString.copyFromUtf8("foo"); + private static final ResultSetMetadata DEFAULT_METADATA = + metadata(columnMetadata("strCol", stringType())); + + @Rule public GrpcServerRule serverRule = new GrpcServerRule(); + private TestBigtableService service; + private BigtableDataClient client; + private PreparedStatement preparedStatement; + + @Before + public void setUp() throws IOException { + service = new TestBigtableService(); + serverRule.getServiceRegistry().addService(service); + + BigtableDataSettings.Builder settings = + BigtableDataSettings.newBuilder() + .setProjectId(PROJECT_ID) + .setInstanceId(INSTANCE_ID) + .setAppProfileId(APP_PROFILE_ID) + .setCredentialsProvider(NoCredentialsProvider.create()); + + settings + .stubSettings() + .setTransportChannelProvider( + FixedTransportChannelProvider.create( + GrpcTransportChannel.create(serverRule.getChannel()))) + // Refreshing channel doesn't work with FixedTransportChannelProvider + .setRefreshingChannel(false) + .build(); + + client = BigtableDataClient.create(settings.build()); + preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto(prepareResponse(PREPARED_QUERY, DEFAULT_METADATA)), + new HashMap<>()); + } + + @After + public void tearDown() { + if (client != null) { + client.close(); + } + } + + @Test + public void testAllSuccesses() { + service.addExpectation( + RpcExpectation.create() + .respondWith( + partialResultSetWithoutToken(stringValue("foo")), + partialResultSetWithoutToken(stringValue("bar")), + partialResultSetWithToken(stringValue("baz")))); + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThat(rs.getMetadata().getColumns()).hasSize(1); + assertThat(rs.getMetadata().getColumns().get(0).name()).isEqualTo("strCol"); + assertThat(rs.getMetadata().getColumns().get(0).type()).isEqualTo(SqlType.string()); + + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("bar"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("baz"); + assertThat(rs.next()).isFalse(); + rs.close(); + } + + @Test + public void testRetryOnInitialError() { + // - First attempt immediately fails + // - Second attempt returns 'foo', w a token, and succeeds + // Expect result to be 'foo' + service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create().respondWith(partialResultSetWithToken(stringValue("foo")))); + + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isFalse(); + rs.close(); + assertThat(service.requestCount).isEqualTo(2); + } + + @Test + public void testResumptionToken() { + // - First attempt gets a response with a token, and then fails with unavailable + // - Second Expects the request to contain the previous token, and returns a new response w + // token and then fails with unavailable + // - Third expects the request to contain the second token, returns a new response w token + // and then succeeds + // We expect the results to contain all of the returned data (no reset batches) + service.addExpectation( + RpcExpectation.create() + .respondWith( + partialResultSetWithToken(ByteString.copyFromUtf8("token1"), stringValue("foo"))) + .respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create() + .withResumeToken(ByteString.copyFromUtf8("token1")) + .respondWith( + partialResultSetWithToken(ByteString.copyFromUtf8("token2"), stringValue("bar"))) + .respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create() + .withResumeToken(ByteString.copyFromUtf8("token2")) + .respondWith( + partialResultSetWithToken(ByteString.copyFromUtf8("final"), stringValue("baz")))); + + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("bar"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("baz"); + assertThat(rs.next()).isFalse(); + rs.close(); + assertThat(service.requestCount).isEqualTo(3); + } + + @Test + public void testResetOnResumption() { + // - First attempt returns 'foo' with 'token1', then 'discard' w no token, then fails + // - Second attempt should resume w 'token1', returns an incomplete batch of two chunks. First + // chunk contains the reset bit and a some data, second contains some data, we fail w/o + // returning the final chunk w a token. + // - Third attempt should resume w 'token1', we return 'baz' w reset & a token, succeed + // Expect the results to be 'foo' and 'baz' + service.addExpectation( + RpcExpectation.create() + .respondWith( + partialResultSetWithToken(ByteString.copyFromUtf8("token1"), stringValue("foo")), + // This is after the token so should be dropped + partialResultSetWithoutToken(stringValue("discard"))) + .respondWithStatus(Code.UNAVAILABLE)); + List chunkedResponses = + partialResultSets( + 3, + true, + ByteString.copyFromUtf8("token2"), + stringValue("longerStringDiscard"), + stringValue("discard")); + service.addExpectation( + RpcExpectation.create() + .withResumeToken(ByteString.copyFromUtf8("token1")) + // Skip the last response, so we don't send a new token + .respondWith(chunkedResponses.get(0), chunkedResponses.get(1)) + .respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create() + .withResumeToken(ByteString.copyFromUtf8("token1")) + .respondWith( + partialResultSets(1, true, ByteString.copyFromUtf8("final"), stringValue("baz")) + .get(0))); + + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("baz"); + assertThat(rs.next()).isFalse(); + rs.close(); + assertThat(service.requestCount).isEqualTo(3); + } + + @Test + public void testErrorAfterFinalData() { + // - First attempt returns 'foo', 'bar', 'baz' w 'finalToken' but fails w unavailable + // - Second attempt uses 'finalToken' and succeeds + // Expect results to be 'foo', 'bar', 'baz' + service.addExpectation( + RpcExpectation.create() + .respondWith( + partialResultSetWithoutToken(stringValue("foo")), + partialResultSetWithoutToken(stringValue("bar")), + partialResultSetWithToken( + ByteString.copyFromUtf8("finalToken"), stringValue("baz"))) + .respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create().withResumeToken(ByteString.copyFromUtf8("finalToken"))); + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThat(rs.getMetadata().getColumns()).hasSize(1); + assertThat(rs.getMetadata().getColumns().get(0).name()).isEqualTo("strCol"); + assertThat(rs.getMetadata().getColumns().get(0).type()).isEqualTo(SqlType.string()); + + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("bar"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("baz"); + assertThat(rs.next()).isFalse(); + rs.close(); + } + + // TODO test changing metadata when plan refresh is implemented + + @Test + public void permanentErrorPropagatesToMetadata() { + service.addExpectation(RpcExpectation.create().respondWithStatus(Code.INVALID_ARGUMENT)); + + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + ApiException e = assertThrows(ApiException.class, rs::getMetadata); + assertThat(e.getStatusCode().getCode()).isEqualTo(StatusCode.Code.INVALID_ARGUMENT); + } + + @Test + public void exhaustedRetriesPropagatesToMetadata() throws IOException { + int attempts = + EnhancedBigtableStubSettings.newBuilder() + .executeQuerySettings() + .getRetrySettings() + .getMaxAttempts(); + assertThat(attempts).isGreaterThan(1); + for (int i = 0; i < attempts; i++) { + service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + } + + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + ApiException e = assertThrows(ApiException.class, rs::getMetadata); + assertThat(e.getStatusCode().getCode()).isEqualTo(StatusCode.Code.UNAVAILABLE); + } + + @Test + public void retryableErrorWithSuccessfulRetryDoesNotPropagateToMetadata() { + service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create().respondWith(tokenOnlyResultSet(ByteString.copyFromUtf8("t")))); + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThat(rs.getMetadata().getColumns()).hasSize(1); + } + + @Test + public void preservesParamsOnRetry() { + Map> paramTypes = ImmutableMap.of("strParam", SqlType.string()); + PreparedStatement preparedStatementWithParams = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse(metadata(columnMetadata("strCol", stringType())))), + paramTypes); + Map params = + ImmutableMap.of("strParam", stringValue("foo").toBuilder().setType(stringType()).build()); + service.addExpectation( + RpcExpectation.create() + .withParams(params) + .respondWith( + partialResultSetWithToken(ByteString.copyFromUtf8("token1"), stringValue("foo"))) + .respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + RpcExpectation.create() + .withParams(params) + .withResumeToken(ByteString.copyFromUtf8("token1")) + .respondWith( + partialResultSetWithToken(ByteString.copyFromUtf8("token2"), stringValue("bar")))); + + ResultSet rs = + client.executeQuery( + preparedStatementWithParams.bind().setStringParam("strParam", "foo").build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("bar"); + assertThat(rs.next()).isFalse(); + } + + @Test + public void failsOnCompleteWithOpenPartialBatch() { + // Return 'foo' with no token, followed by ok + // This should throw an error, as the backend has violated its contract + service.addExpectation( + RpcExpectation.create() + .respondWith(partialResultSetWithoutToken(stringValue("foo"))) + .respondWithStatus(Code.OK)); + ResultSet rs = client.executeQuery(preparedStatement.bind().build()); + assertThrows(IncompleteStreamException.class, rs::next); + } + + private static class TestBigtableService extends BigtableGrpc.BigtableImplBase { + Queue expectations = new LinkedBlockingDeque<>(); + int requestCount = 0; + + void addExpectation(RpcExpectation expectation) { + expectations.add(expectation); + } + + @Override + public void executeQuery( + ExecuteQueryRequest request, StreamObserver responseObserver) { + RpcExpectation expectedRpc = expectations.poll(); + requestCount++; + int requestIndex = requestCount - 1; + + Truth.assertWithMessage("Unexpected request#" + requestIndex + ":" + request.toString()) + .that(expectedRpc) + .isNotNull(); + Truth.assertWithMessage("Unexpected request#" + requestIndex) + .that(request) + .isEqualTo(expectedRpc.getExpectedRequest()); + + for (ExecuteQueryResponse response : expectedRpc.responses) { + responseObserver.onNext(response); + } + if (expectedRpc.statusCode.toStatus().isOk()) { + responseObserver.onCompleted(); + } else if (expectedRpc.exception != null) { + responseObserver.onError(expectedRpc.exception); + } else { + responseObserver.onError(expectedRpc.statusCode.toStatus().asRuntimeException()); + } + } + } + + private static class RpcExpectation { + ExecuteQueryRequest.Builder request; + Status.Code statusCode; + @Nullable ApiException exception; + List responses; + + private RpcExpectation() { + this.request = ExecuteQueryRequest.newBuilder(); + this.request.setPreparedQuery(PREPARED_QUERY); + this.request.setInstanceName(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID)); + this.request.setAppProfileId(APP_PROFILE_ID); + this.statusCode = Code.OK; + this.responses = new ArrayList<>(); + } + + static RpcExpectation create() { + return new RpcExpectation(); + } + + RpcExpectation withResumeToken(ByteString resumeToken) { + this.request.setResumeToken(resumeToken); + return this; + } + + RpcExpectation withParams(Map params) { + this.request.putAllParams(params); + return this; + } + + RpcExpectation respondWithStatus(Status.Code code) { + this.statusCode = code; + return this; + } + + RpcExpectation respondWithException(Status.Code code, ApiException exception) { + this.statusCode = code; + this.exception = exception; + return this; + } + + RpcExpectation respondWith(ExecuteQueryResponse... responses) { + this.responses = Arrays.asList(responses); + return this; + } + + ExecuteQueryRequest getExpectedRequest() { + return this.request.build(); + } + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java new file mode 100644 index 0000000000..77ec69da9d --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java @@ -0,0 +1,87 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.stub.sql; + +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.api.core.SettableApiFuture; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.SqlRow; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.cloud.bigtable.data.v2.stub.sql.MetadataErrorHandlingCallable.MetadataErrorHandlingObserver; +import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockResponseObserver; +import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall; +import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; +import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; +import java.util.HashMap; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class MetadataErrorHandlingCallableTest { + private ExecuteQueryCallContext callContext; + private MockResponseObserver outerObserver; + private SettableApiFuture metadataFuture; + private MetadataErrorHandlingCallable.MetadataErrorHandlingObserver observer; + + @Before + public void setUp() { + metadataFuture = SettableApiFuture.create(); + PreparedStatement preparedStatement = + PreparedStatementImpl.create( + PrepareResponse.fromProto( + prepareResponse( + metadata( + columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), + new HashMap<>()); + + callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); + outerObserver = new MockResponseObserver<>(true); + observer = new MetadataErrorHandlingObserver(outerObserver, callContext); + } + + // cancel will manifest as an onError call so these are testing both cancellation and + // other exceptions + @Test + public void observer_passesThroughErrorAndSetsMetadataException() { + MockServerStreamingCallable innerCallable = + new MockServerStreamingCallable<>(); + innerCallable.call(callContext, observer); + MockServerStreamingCall lastCall = innerCallable.popLastCall(); + MockStreamController innerController = lastCall.getController(); + + innerController.getObserver().onError(new CancellationException("Cancelled")); + + assertThat(metadataFuture.isDone()).isTrue(); + assertThrows(ExecutionException.class, metadataFuture::get); + ExecutionException e = assertThrows(ExecutionException.class, metadataFuture::get); + assertThat(e.getCause()).isInstanceOf(CancellationException.class); + assertThat(outerObserver.isDone()).isTrue(); + assertThat(outerObserver.getFinalError()).isInstanceOf(CancellationException.class); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java similarity index 75% rename from google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java rename to google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java index 854268102e..2ab050c573 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java @@ -24,10 +24,8 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; -import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.tokenOnlyResultSet; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import com.google.api.core.SettableApiFuture; @@ -39,16 +37,14 @@ import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; -import com.google.cloud.bigtable.data.v2.stub.sql.MetadataResolvingCallable.MetadataObserver; +import com.google.cloud.bigtable.data.v2.stub.sql.PlanRefreshingCallable.MetadataObserver; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockResponseObserver; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; -import com.google.protobuf.ByteString; import java.util.Collections; import java.util.HashMap; -import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import org.junit.Before; import org.junit.Test; @@ -56,7 +52,7 @@ import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public class MetadataResolvingCallableTest { +public class PlanRefreshingCallableTest { private static final ExecuteQueryRequest FAKE_REQUEST = ExecuteQueryRequest.newBuilder().build(); private static final com.google.bigtable.v2.ResultSetMetadata METADATA = @@ -67,7 +63,7 @@ public class MetadataResolvingCallableTest { ExecuteQueryCallContext callContext; MockResponseObserver outerObserver; SettableApiFuture metadataFuture; - MetadataResolvingCallable.MetadataObserver observer; + PlanRefreshingCallable.MetadataObserver observer; @Before public void setUp() { @@ -118,48 +114,6 @@ public void observer_setsFutureAndPassesThroughResponses() assertThat(outerObserver.getFinalError()).isNull(); } - // cancel will manifest as an onError call so these are testing both cancellation and - // other exceptions - @Test - public void observer_passesThroughErrorBeforeResolvingMetadata() { - MockServerStreamingCallable innerCallable = - new MockServerStreamingCallable<>(); - innerCallable.call(FAKE_REQUEST, observer); - MockServerStreamingCall lastCall = - innerCallable.popLastCall(); - MockStreamController innerController = lastCall.getController(); - - innerController.getObserver().onError(new CancellationException("Cancelled")); - - assertThat(metadataFuture.isDone()).isTrue(); - assertThrows(ExecutionException.class, metadataFuture::get); - ExecutionException e = assertThrows(ExecutionException.class, metadataFuture::get); - assertThat(e.getCause()).isInstanceOf(CancellationException.class); - assertThat(outerObserver.isDone()).isTrue(); - assertThat(outerObserver.getFinalError()).isInstanceOf(CancellationException.class); - } - - @Test - public void observer_passesThroughErrorAfterSettingMetadata() - throws ExecutionException, InterruptedException { - MockServerStreamingCallable innerCallable = - new MockServerStreamingCallable<>(); - innerCallable.call(FAKE_REQUEST, observer); - MockServerStreamingCall lastCall = - innerCallable.popLastCall(); - MockStreamController innerController = lastCall.getController(); - - innerController.getObserver().onResponse(tokenOnlyResultSet(ByteString.copyFromUtf8("token"))); - innerController.getObserver().onError(new RuntimeException("exception after metadata")); - - assertThat(metadataFuture.isDone()).isTrue(); - assertThat(metadataFuture.get()).isEqualTo(ProtoResultSetMetadata.fromProto(METADATA)); - assertThat(outerObserver.popNextResponse()) - .isEqualTo(tokenOnlyResultSet(ByteString.copyFromUtf8("token"))); - assertThat(outerObserver.isDone()).isTrue(); - assertThat(outerObserver.getFinalError()).isInstanceOf(RuntimeException.class); - } - @Test public void observer_passThroughOnStart() { MockServerStreamingCallable innerCallable = @@ -193,8 +147,7 @@ public void testCallable() throws ExecutionException, InterruptedException { ServerStreamingStashCallable innerCallable = new ServerStreamingStashCallable<>(Collections.singletonList(DATA)); RequestContext requestContext = RequestContext.create("project", "instance", "profile"); - MetadataResolvingCallable callable = - new MetadataResolvingCallable(innerCallable, requestContext); + PlanRefreshingCallable callable = new PlanRefreshingCallable(innerCallable, requestContext); MockResponseObserver outerObserver = new MockResponseObserver<>(true); SettableApiFuture metadataFuture = SettableApiFuture.create(); PreparedStatement preparedStatement = diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java index c3643789a3..23fc8bbef3 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java @@ -192,6 +192,11 @@ public static ExecuteQueryResponse partialResultSetWithToken(Value... values) { return partialResultSets(1, false, ByteString.copyFromUtf8("test"), values).get(0); } + /** Creates a single response representing a complete batch, with a resume token of token */ + public static ExecuteQueryResponse partialResultSetWithToken(ByteString token, Value... values) { + return partialResultSets(1, false, token, values).get(0); + } + public static ExecuteQueryResponse tokenOnlyResultSet(ByteString token) { return ExecuteQueryResponse.newBuilder() .setResults(PartialResultSet.newBuilder().setResumeToken(token)) From 913f748740a551072610810c56ebd57d51d23e22 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Wed, 19 Feb 2025 15:01:47 -0500 Subject: [PATCH 07/11] Implement refresh for PreparedStatement Change-Id: I7a83e84d689e438c32c4e9a320b6fff52d86e0de --- .../bigtable/data/v2/BigtableDataClient.java | 33 +- .../v2/internal/PreparedStatementImpl.java | 252 ++++++- .../data/v2/models/sql/BoundStatement.java | 38 +- .../data/v2/models/sql/PreparedStatement.java | 13 +- ...paredStatementRefreshTimeoutException.java | 30 + .../data/v2/stub/EnhancedBigtableStub.java | 14 +- .../v2/stub/sql/ExecuteQueryCallContext.java | 90 ++- .../v2/stub/sql/ExecuteQueryCallable.java | 15 +- .../v2/stub/sql/PlanRefreshingCallable.java | 136 +++- .../data/v2/BigtableDataClientTests.java | 12 + .../internal/PreparedStatementImplTest.java | 417 ++++++++++++ .../data/v2/internal/ResultSetImplTest.java | 21 +- .../v2/models/sql/BoundStatementTest.java | 4 +- .../v2/stub/EnhancedBigtableStubTest.java | 22 +- .../bigtable/data/v2/stub/HeadersTest.java | 9 +- .../stub/sql/ExecuteQueryCallContextTest.java | 102 ++- .../v2/stub/sql/ExecuteQueryCallableTest.java | 27 +- .../ExecuteQueryResumptionStrategyTest.java | 20 +- .../v2/stub/sql/ExecuteQueryRetryTest.java | 617 ++++++++++++++---- .../MetadataErrorHandlingCallableTest.java | 15 +- .../stub/sql/PlanRefreshingCallableTest.java | 146 ++++- .../data/v2/stub/sql/SqlProtoFactory.java | 365 ++++++++++- .../data/v2/stub/sql/SqlProtoFactoryTest.java | 15 + .../stub/sql/SqlRowMergingCallableTest.java | 68 +- .../gaxx/testing/MockStreamingApi.java | 14 +- test-proxy/pom.xml | 4 +- 26 files changed, 2151 insertions(+), 348 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatementRefreshTimeoutException.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 694abd4281..7f16645060 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -2711,32 +2711,35 @@ public void readChangeStreamAsync( * Executes a SQL Query and returns a ResultSet to iterate over the results. The returned * ResultSet instance is not threadsafe, it can only be used from single thread. * + *

      The {@link BoundStatement} must be built from a {@link PreparedStatement} created using + * the same instance and app profile. + * *

      Sample code: * * 

      {@code
    * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
    *   String query = "SELECT CAST(cf['stringCol'] AS STRING) FROM [TABLE]";
    *   Map> paramTypes = new HashMap<>();
-   *   try (PreparedStatement preparedStatement = bigtableDataClient.prepareStatement(query, paramTypes)) {
-   *       // Ideally one PreparedStatement should be reused across requests
-   *       BoundStatement boundStatement = preparedStatement.bind()
-   *          // set any query params before calling build
-   *          .build();
-   *       try (ResultSet resultSet = bigtableDataClient.executeQuery(boundStatement)) {
-   *           while (resultSet.next()) {
-   *               String s = resultSet.getString("stringCol");
-   *                // do something with data
-   *           }
-   *        } catch (RuntimeException e) {
-   *            e.printStackTrace();
-   *        }
+   *   PreparedStatement preparedStatement = bigtableDataClient.prepareStatement(query, paramTypes));
+   *   // Ideally one PreparedStatement should be reused across requests
+   *   BoundStatement boundStatement = preparedStatement.bind()
+   *      // set any query params before calling build
+   *      .build();
+   *   try (ResultSet resultSet = bigtableDataClient.executeQuery(boundStatement)) {
+   *       while (resultSet.next()) {
+   *           String s = resultSet.getString("stringCol");
+   *            // do something with data
+   *       }
+   *    } catch (RuntimeException e) {
+   *        e.printStackTrace();
    *   }
    * }
      * - * @see {@link PreparedStatement} & {@link BoundStatement} For query options. + * @see {@link PreparedStatement} & {@link BoundStatement} for query options. */ @BetaApi public ResultSet executeQuery(BoundStatement boundStatement) { + boundStatement.assertUsingSameStub(stub); SqlServerStream stream = stub.createExecuteQueryCallable().call(boundStatement); return ResultSetImpl.create(stream); } @@ -2754,7 +2757,7 @@ public ResultSet executeQuery(BoundStatement boundStatement) { public PreparedStatement prepareStatement(String query, Map> paramTypes) { PrepareQueryRequest request = PrepareQueryRequest.create(query, paramTypes); PrepareResponse response = stub.prepareQueryCallable().call(request); - return PreparedStatementImpl.create(response, paramTypes); + return PreparedStatementImpl.create(response, paramTypes, request, stub); } /** Close the clients and releases all associated resources. */ diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java index abaece4e87..e6690d5a8c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImpl.java @@ -15,32 +15,63 @@ */ package com.google.cloud.bigtable.data.v2.internal; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.core.InternalApi; +import com.google.auto.value.AutoValue; import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement.Builder; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.Futures; +import java.time.Duration; +import java.time.Instant; import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicReference; /** - * Implementation of PreparedStatement that handles PreparedQuery refresh + * Implementation of PreparedStatement that handles PreparedQuery refresh. + * + *

      This allows for both hard refresh and background refresh of the current PreparedQueryData. + * When the server returns an error indicating that a plan is expired, hardRefresh should be used. + * Otherwise this will handle updating the PreparedQuery in the background, whenever it is accessed + * within one second of expiry. * *

      This is considered an internal implementation detail and should not be used by applications. */ -// TODO implement plan refresh @InternalApi("For internal use only") public class PreparedStatementImpl implements PreparedStatement { - private PrepareResponse response; + // Time before plan expiry to trigger background refresh + private static final Duration EXPIRY_REFRESH_WINDOW = Duration.ofSeconds(1L); + private final AtomicReference currentState; private final Map> paramTypes; + private final PrepareQueryRequest prepareRequest; + private final EnhancedBigtableStub stub; - public PreparedStatementImpl(PrepareResponse response, Map> paramTypes) { - this.response = response; + @VisibleForTesting + protected PreparedStatementImpl( + PrepareResponse response, + Map> paramTypes, + PrepareQueryRequest request, + EnhancedBigtableStub stub) { + this.currentState = new AtomicReference<>(PrepareQueryState.createInitialState(response)); this.paramTypes = paramTypes; + this.prepareRequest = request; + this.stub = stub; } public static PreparedStatement create( - PrepareResponse response, Map> paramTypes) { - return new PreparedStatementImpl(response, paramTypes); + PrepareResponse response, + Map> paramTypes, + PrepareQueryRequest request, + EnhancedBigtableStub stub) { + return new PreparedStatementImpl(response, paramTypes, request, stub); } @Override @@ -48,14 +79,207 @@ public BoundStatement.Builder bind() { return new Builder(this, paramTypes); } - // TODO update when plan refresh is implement - @Override - public PrepareResponse getPrepareResponse() { - return response; + /** + * Asserts that the given stub matches the stub used for plan refresh. This is necessary to ensure + * that the request comes from the same client and uses the same configuration. We enforce this + * make sure plan refresh will continue to work as expected throughout the lifecycle of + * executeQuery requests. + */ + public void assertUsingSameStub(EnhancedBigtableStub stub) { + Preconditions.checkArgument( + this.stub == stub, + "executeQuery must be called from the same client instance that created the PreparedStatement being used."); } - @Override - public void close() throws Exception { - // TODO cancel any background refresh + /** + * When the client receives an error indicating the current plan has expired, it should call + * immediate refresh with the version of the expired plan. UID is used to handle concurrent + * refresh without making duplicate calls. + * + * @param expiredPreparedQueryVersion version of the PreparedQuery used to make the request that + * triggered immediate refresh + * @return refreshed PreparedQuery to use for retry. + */ + public synchronized PreparedQueryData markExpiredAndStartRefresh( + PreparedQueryVersion expiredPreparedQueryVersion) { + PrepareQueryState localState = this.currentState.get(); + // Check if the expired plan is the current plan. If it's not, then the plan has already + // been refreshed by another thread. + if (!(localState.current().version() == expiredPreparedQueryVersion)) { + return localState.current(); + } + startBackgroundRefresh(expiredPreparedQueryVersion); + // Immediately promote the refresh we just started + return promoteBackgroundRefreshingPlan(expiredPreparedQueryVersion); + } + + private synchronized PreparedQueryData promoteBackgroundRefreshingPlan( + PreparedQueryVersion expiredPreparedQueryVersion) { + PrepareQueryState localState = this.currentState.get(); + // If the expired plan has already been removed, return the current plan + if (!(localState.current().version() == expiredPreparedQueryVersion)) { + return localState.current(); + } + // There is a chance that the background plan could be expired if the PreparedStatement + // isn't used for a long time. It will be refreshed on the next retry if necessary. + PrepareQueryState nextState = localState.promoteBackgroundPlan(); + this.currentState.set(nextState); + return nextState.current(); + } + + /** + * If planNearExpiry is still the latest plan, and there is no ongoing background refresh, start a + * background refresh. Otherwise, refresh has already been triggered for this plan, so do nothing. + */ + private synchronized void startBackgroundRefresh(PreparedQueryVersion planVersionNearExpiry) { + PrepareQueryState localState = this.currentState.get(); + // We've already updated the plan we are triggering refresh based on + if (!(localState.current().version() == planVersionNearExpiry)) { + return; + } + // Another thread already started the refresh + if (localState.maybeBackgroundRefresh().isPresent()) { + return; + } + ApiFuture nextPlanFuture = getFreshPlan(); + PrepareQueryState withRefresh = localState.withBackgroundPlan(nextPlanFuture); + this.currentState.set(withRefresh); + } + + ApiFuture getFreshPlan() { + return this.stub.prepareQueryCallable().futureCall(this.prepareRequest); + } + + /** + * Check the expiry of the current plan, if it's future is resolved. If we are within 1s of + * expiry, call startBackgroundRefresh with the version of the latest PrepareQuery. + */ + void backgroundRefreshIfNeeded() { + PrepareQueryState localState = this.currentState.get(); + if (localState.maybeBackgroundRefresh().isPresent()) { + // We already have an ongoing refresh + return; + } + PreparedQueryData currentPlan = localState.current(); + // Can't access ttl until the current prepare future has resolved + if (!currentPlan.prepareFuture().isDone()) { + return; + } + try { + // Trigger a background refresh if within 1 second of TTL + Instant currentPlanExpireTime = Futures.getDone(currentPlan.prepareFuture()).validUntil(); + Instant backgroundRefreshTime = currentPlanExpireTime.minus(EXPIRY_REFRESH_WINDOW); + if (Instant.now().isAfter(backgroundRefreshTime)) { + // Initiate a background refresh. startBackgroundRefresh handles deduplication. + startBackgroundRefresh(currentPlan.version()); + } + } catch (ExecutionException | CancellationException e) { + // Do nothing if we can't get the future result, a refresh will be done when it's actually + // needed, or during the next call to this method + } + } + + /** + * Returns the most recently refreshed PreparedQueryData. It may still be refreshing if the + * previous plan has expired. + */ + public PreparedQueryData getLatestPrepareResponse() { + PrepareQueryState localState = currentState.get(); + if (localState.maybeBackgroundRefresh().isPresent() + && localState.maybeBackgroundRefresh().get().prepareFuture().isDone()) { + // TODO: consider checking if background plan has already expired and triggering + // a new refresh if so. Right now we are ok with attempting a request w an expired + // plan + + // Current background refresh has completed, so we should make it the current plan. + // promoteBackgroundRefreshingPlan handles duplicate calls. + return promoteBackgroundRefreshingPlan(localState.current().version()); + } else { + backgroundRefreshIfNeeded(); + return localState.current(); + } + } + + /** + * Used to compare different versions of a PreparedQuery by comparing reference equality. + * + *

      This is considered an internal implementation detail and not meant to be used by + * applications. + */ + @InternalApi("For internal use only") + public static class PreparedQueryVersion {} + + /** + * Manages the data around the latest prepared query + * + *

      This is considered an internal implementation detail and not meant to be used by + * applications. + */ + @InternalApi("For internal use only") + @AutoValue + public abstract static class PreparedQueryData { + /** + * Unique identifier for each version of a PreparedQuery. Changes each time the plan is + * refreshed + */ + public abstract PreparedQueryVersion version(); + + /** + * A future holding the prepareResponse. It will never fail, so the caller is responsible for + * timing out requests based on the retry settings of the execute query request + */ + public abstract ApiFuture prepareFuture(); + + public static PreparedQueryData create(ApiFuture prepareFuture) { + return new AutoValue_PreparedStatementImpl_PreparedQueryData( + new PreparedQueryVersion(), prepareFuture); + } + } + + /** + * Encapsulates the state needed to for PreparedStatementImpl. This is both the latest + * PrepareQuery response and, when present, any ongoing background refresh. + * + *

      This is stored together because it is accessed concurrently. This makes it easy to reason + * about and mutate the state atomically. + */ + @AutoValue + abstract static class PrepareQueryState { + /** The data representing the latest PrepareQuery response */ + abstract PreparedQueryData current(); + + /** An Optional, that if present represents an ongoing background refresh attempt */ + abstract Optional maybeBackgroundRefresh(); + + /** Creates a fresh state, using initialPlan as current, with no backgroundRefresh */ + static PrepareQueryState createInitialState(PrepareResponse initialPlan) { + PreparedQueryData initialData = + PreparedQueryData.create(ApiFutures.immediateFuture(initialPlan)); + return new AutoValue_PreparedStatementImpl_PrepareQueryState(initialData, Optional.empty()); + } + + /** + * Returns a new state with the same current PreparedQueryData, using the given PrepareResponse + * future to add a backgroundRefresh + */ + PrepareQueryState withBackgroundPlan(ApiFuture backgroundPlan) { + return new AutoValue_PreparedStatementImpl_PrepareQueryState( + current(), Optional.of(PreparedQueryData.create(backgroundPlan))); + } + + /** + * Returns a new state with the background plan promoted to current, and without a new + * background refresh. This should be used to update the state once a backgroundRefresh has + * completed. + */ + PrepareQueryState promoteBackgroundPlan() { + if (maybeBackgroundRefresh().isPresent()) { + return new AutoValue_PreparedStatementImpl_PrepareQueryState( + maybeBackgroundRefresh().get(), Optional.empty()); + } + // We don't expect this to happen, but if so returning the current plan allows retry on + // subsequent attempts + return this; + } } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java index 4692923d0b..3165edb23c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java @@ -23,9 +23,12 @@ import com.google.bigtable.v2.Value; import com.google.cloud.Date; import com.google.cloud.bigtable.data.v2.internal.NameUtil; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryData; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryVersion; import com.google.cloud.bigtable.data.v2.internal.QueryParamUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; @@ -65,27 +68,26 @@ @BetaApi public class BoundStatement { - private final PreparedStatement preparedStatement; + private final PreparedStatementImpl preparedStatement; private final Map params; - private BoundStatement(PreparedStatement preparedStatement, Map params) { + private BoundStatement(PreparedStatementImpl preparedStatement, Map params) { this.preparedStatement = preparedStatement; this.params = params; } - // TODO return a future when plan refresh is implemented /** - * Get's the most recent version of the PrepareResponse associated with this query. + * Gets the most recent version of the PrepareResponse associated with this query. * *

      This is considered an internal implementation detail and should not be used by applications. */ @InternalApi("For internal use only") - public PrepareResponse getLatestPrepareResponse() { - return preparedStatement.getPrepareResponse(); + public PreparedQueryData getLatestPrepareResponse() { + return preparedStatement.getLatestPrepareResponse(); } public static class Builder { - private final PreparedStatement preparedStatement; + private final PreparedStatementImpl preparedStatement; private final Map> paramTypes; private final Map params; @@ -96,7 +98,7 @@ public static class Builder { * applications. */ @InternalApi("For internal use only") - public Builder(PreparedStatement preparedStatement, Map> paramTypes) { + public Builder(PreparedStatementImpl preparedStatement, Map> paramTypes) { this.preparedStatement = preparedStatement; this.paramTypes = paramTypes; this.params = new HashMap<>(); @@ -389,4 +391,22 @@ public ExecuteQueryRequest toProto( } return requestBuilder.build(); } + + @InternalApi("For internal use only") + public PreparedQueryData markExpiredAndStartRefresh( + PreparedQueryVersion expiredPreparedQueryVersion) { + return this.preparedStatement.markExpiredAndStartRefresh(expiredPreparedQueryVersion); + } + + /** + * Asserts that the given stub matches the stub used for plan refresh. This is necessary to ensure + * that the request comes from the same client and uses the same configuration. + * + *

      This is considered an internal implementation detail and not meant to be used by + * applications + */ + @InternalApi + public void assertUsingSameStub(EnhancedBigtableStub stub) { + this.preparedStatement.assertUsingSameStub(stub); + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java index cbcaba829c..46c0ec59b7 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java @@ -16,27 +16,20 @@ package com.google.cloud.bigtable.data.v2.models.sql; import com.google.api.core.BetaApi; -import com.google.api.core.InternalApi; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; /** * The results of query preparation that can be used to create {@link BoundStatement}s to execute * queries. * - *

      Whenever possible this should be shared across different instances the same query, in order to - * amortize query preparation costs. + *

      Whenever possible this should be shared across different instances of the same query, in order + * to amortize query preparation costs. */ @BetaApi -public interface PreparedStatement extends AutoCloseable { +public interface PreparedStatement { /** * @return {@link BoundStatement.Builder} to bind query params to and pass to {@link * com.google.cloud.bigtable.data.v2.BigtableDataClient#executeQuery(BoundStatement)} */ BoundStatement.Builder bind(); - - // TODO once refresh is implemented this will return a future so we can - // wait while plan is refreshing - @InternalApi("For internal use only") - PrepareResponse getPrepareResponse(); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatementRefreshTimeoutException.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatementRefreshTimeoutException.java new file mode 100644 index 0000000000..413997aff8 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatementRefreshTimeoutException.java @@ -0,0 +1,30 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.models.sql; + +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.ApiException; +import io.grpc.Status.Code; + +/** + * Error thrown when an executeQuery attempt hits the attempt deadline waiting for {@link + * PreparedStatement} to refresh it's underlying plan. + */ +public class PreparedStatementRefreshTimeoutException extends ApiException { + public PreparedStatementRefreshTimeoutException(String message) { + super(message, null, GrpcStatusCode.of(Code.DEADLINE_EXCEEDED), true); + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 2dfbdc181c..db75050a4f 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -1141,8 +1141,11 @@ private UnaryCallable createReadModifyWriteRowCallable( *

    • Convert a {@link BoundStatement} into a {@link ExecuteQueryCallContext}, which passes the * {@link BoundStatement} & a future for the {@link * com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata} up the call chain. - *
    • Upon receiving the response stream, it will set the metadata future and translate the + *
    • Refresh expired {@link PrepareResponse} when the server returns a specific error} + *
    • Add retry/resume on failures + *
    • Upon receiving the first resume_token, it will set the metadata future and translate the * {@link com.google.bigtable.v2.PartialResultSet}s into {@link SqlRow}s + *
    • Pass through non-retryable errors to the metadata future *
    • Add tracing & metrics. *
    • Wrap the metadata future & row stream into a {@link * com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream} @@ -1169,7 +1172,7 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { ServerStreamingCallable withStatsHeaders = new StatsHeadersServerStreamingCallable<>(base); - ServerStreamingCallable withMetadataObserver = + ServerStreamingCallable withPlanRefresh = new PlanRefreshingCallable(withStatsHeaders, requestContext); ServerStreamingCallSettings retrySettings = @@ -1186,7 +1189,7 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { // attempt stream will have reset set to true, so any unyielded data from the previous // attempt will be reset properly ServerStreamingCallable retries = - withRetries(withMetadataObserver, retrySettings); + withRetries(withPlanRefresh, retrySettings); ServerStreamingCallable merging = new SqlRowMergingCallable(retries); @@ -1217,8 +1220,7 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { traced.withDefaultCallContext( clientContext .getDefaultCallContext() - .withRetrySettings(settings.executeQuerySettings().getRetrySettings())), - requestContext); + .withRetrySettings(settings.executeQuerySettings().getRetrySettings()))); } private UnaryCallable createPrepareQueryCallable() { @@ -1486,10 +1488,10 @@ public ExecuteQueryCallable executeQueryCallable() { return executeQueryCallable; } + @InternalApi public UnaryCallable prepareQueryCallable() { return prepareQueryCallable; } - // private SpanName getSpanName(String methodName) { diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java index 337dc2d497..c4deda0387 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext.java @@ -17,17 +17,31 @@ import com.google.api.core.InternalApi; import com.google.api.core.SettableApiFuture; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ApiExceptionFactory; +import com.google.api.gax.rpc.ApiExceptions; +import com.google.api.gax.rpc.StatusCode; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryData; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatementRefreshTimeoutException; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; +import com.google.common.base.Preconditions; import com.google.protobuf.ByteString; +import io.grpc.Deadline; +import io.grpc.Status.Code; +import java.time.Instant; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import javax.annotation.Nullable; /** - * Used to provide a future to the ExecuteQuery callable chain in order to return metadata to users - * outside of the stream of rows. + * Used to handle the state associated with an ExecuteQuery call. This includes plan refresh, resume + * tokens, and metadata resolution. * *

      This should only be constructed by {@link ExecuteQueryCallable} not directly by users. * @@ -38,14 +52,16 @@ public class ExecuteQueryCallContext { private final BoundStatement boundStatement; private final SettableApiFuture metadataFuture; - private final PrepareResponse latestPrepareResponse; + private PreparedQueryData latestPrepareResponse; private @Nullable ByteString resumeToken; + private final Instant startTimeOfCall; private ExecuteQueryCallContext( BoundStatement boundStatement, SettableApiFuture metadataFuture) { this.boundStatement = boundStatement; this.metadataFuture = metadataFuture; this.latestPrepareResponse = boundStatement.getLatestPrepareResponse(); + this.startTimeOfCall = Instant.now(); } public static ExecuteQueryCallContext create( @@ -53,9 +69,44 @@ public static ExecuteQueryCallContext create( return new ExecuteQueryCallContext(boundStatement, metadataFuture); } - ExecuteQueryRequest toRequest(RequestContext requestContext) { - return boundStatement.toProto( - latestPrepareResponse.preparedQuery(), requestContext, resumeToken); + /** + * Builds a request using the latest PrepareQuery data, blocking if necessary for prepare refresh + * to complete. If waiting on refresh, throws a {@link PreparedStatementRefreshTimeoutException} + * exception based on the passed deadline. + * + *

      translates all other exceptions to be retryable so that ExecuteQuery can refresh the plan + * and try again if it has not exhausted its retries + * + *

      If currentAttemptDeadline is null it times out after Long.MAX_VALUE nanoseconds + */ + ExecuteQueryRequest buildRequestWithDeadline( + RequestContext requestContext, @Nullable Deadline currentAttemptDeadline) + throws PreparedStatementRefreshTimeoutException { + // Use max Long as default timeout for simplicity if no deadline is set + long planRefreshWaitTimeoutNanos = Long.MAX_VALUE; + if (currentAttemptDeadline != null) { + planRefreshWaitTimeoutNanos = currentAttemptDeadline.timeRemaining(TimeUnit.NANOSECONDS); + } + try { + PrepareResponse response = + latestPrepareResponse + .prepareFuture() + .get(planRefreshWaitTimeoutNanos, TimeUnit.NANOSECONDS); + return boundStatement.toProto(response.preparedQuery(), requestContext, resumeToken); + } catch (TimeoutException e) { + throw new PreparedStatementRefreshTimeoutException( + "Exceeded deadline waiting for PreparedQuery to refresh"); + } catch (ExecutionException e) { + StatusCode retryStatusCode = GrpcStatusCode.of(Code.FAILED_PRECONDITION); + Throwable cause = e.getCause(); + if (cause instanceof ApiException) { + retryStatusCode = ((ApiException) cause).getStatusCode(); + } + throw ApiExceptionFactory.createException("Plan refresh error", cause, retryStatusCode, true); + } catch (InterruptedException e) { + throw ApiExceptionFactory.createException( + "Plan refresh error", e, GrpcStatusCode.of(Code.FAILED_PRECONDITION), true); + } } /** @@ -64,7 +115,19 @@ ExecuteQueryRequest toRequest(RequestContext requestContext) { * longer change, so we can set the metadata. */ void finalizeMetadata() { - metadataFuture.set(latestPrepareResponse.resultSetMetadata()); + // We don't ever expect an exception here, since we've already received responses at the point + // this is called + try { + Preconditions.checkState( + latestPrepareResponse.prepareFuture().isDone(), + "Unexpected attempt to finalize metadata with unresolved prepare response. This should never as this is called after we receive ExecuteQuery responses, which requires the future to be resolved"); + PrepareResponse response = + ApiExceptions.callAndTranslateApiException(latestPrepareResponse.prepareFuture()); + metadataFuture.set(response.resultSetMetadata()); + } catch (Throwable t) { + metadataFuture.setException(t); + throw t; + } } /** @@ -82,4 +145,17 @@ SettableApiFuture resultSetMetadataFuture() { void setLatestResumeToken(ByteString resumeToken) { this.resumeToken = resumeToken; } + + boolean hasResumeToken() { + return this.resumeToken != null; + } + + void triggerImmediateRefreshOfPreparedQuery() { + latestPrepareResponse = + this.boundStatement.markExpiredAndStartRefresh(latestPrepareResponse.version()); + } + + Instant startTimeOfCall() { + return this.startTimeOfCall; + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java index 534964a97c..687bcdce30 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable.java @@ -22,7 +22,6 @@ import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.bigtable.v2.ExecuteQueryRequest; -import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; @@ -39,15 +38,19 @@ public class ExecuteQueryCallable extends ServerStreamingCallable { private final ServerStreamingCallable inner; - private final RequestContext requestContext; - public ExecuteQueryCallable( - ServerStreamingCallable inner, - RequestContext requestContext) { + public ExecuteQueryCallable(ServerStreamingCallable inner) { this.inner = inner; - this.requestContext = requestContext; } + /** + * This should be used to create execute query calls. This replaces the typical API which allows + * passing of an {@link ApiCallContext}. + * + *

      This class is considered an internal implementation detail and not meant to be used by + * applications. Users should only use executeQuery through the {@link + * com.google.cloud.bigtable.data.v2.BigtableDataClient} + */ public SqlServerStream call(BoundStatement boundStatement) { SettableApiFuture metadataFuture = SettableApiFuture.create(); ServerStream rowStream = diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java index ba15e1368e..521b09da43 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java @@ -16,15 +16,30 @@ package com.google.cloud.bigtable.data.v2.stub.sql; import com.google.api.core.InternalApi; +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.api.gax.rpc.StatusCode.Code; import com.google.api.gax.rpc.StreamController; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.cloud.bigtable.data.v2.internal.RequestContext; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatementRefreshTimeoutException; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.stub.SafeResponseObserver; +import com.google.rpc.PreconditionFailure; +import com.google.rpc.PreconditionFailure.Violation; +import io.grpc.Deadline; +import io.grpc.Status; +import java.time.Duration; +import java.time.Instant; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import javax.annotation.Nullable; /** * Callable that allows passing of {@link ResultSetMetadata} back to users throught the {@link @@ -47,16 +62,98 @@ public PlanRefreshingCallable( @Override public void call( - ExecuteQueryCallContext callContext, + ExecuteQueryCallContext executeQueryCallContext, ResponseObserver responseObserver, - ApiCallContext apiCallContext) { - MetadataObserver observer = new MetadataObserver(responseObserver, callContext); - // TODO toRequest will return a future. We need to timeout waiting for future - // based on the totalTimeout - inner.call(callContext.toRequest(requestContext), observer, apiCallContext); + @Nullable ApiCallContext apiCallContext) { + PlanRefreshingObserver observer = + new PlanRefreshingObserver(responseObserver, executeQueryCallContext); + ExecuteQueryRequest request; + @Nullable GrpcCallContext grpcCallContext = (GrpcCallContext) apiCallContext; + // Convert timeout to an absolute deadline, so we can use it for both the plan refresh and + // the ExecuteQuery rpc + Deadline deadline = getDeadline(grpcCallContext, executeQueryCallContext.startTimeOfCall()); + try { + // TODO: this blocks. That is ok because ResultSet is synchronous. If we ever + // need to make this async that needs to change + request = executeQueryCallContext.buildRequestWithDeadline(requestContext, deadline); + } catch (PreparedStatementRefreshTimeoutException e) { + // If we timed out waiting for refresh, return the retryable error, but don't trigger a + // new refresh since one is ongoing + responseObserver.onError(e); + return; + } catch (Throwable throwable) { + // If we already have a resumeToken we can't refresh the plan, so we throw an error. + // This is not expected to happen, as the plan must be resolved in order for us to + // receive a token + if (executeQueryCallContext.hasResumeToken()) { + responseObserver.onError( + new IllegalStateException( + "Unexpected plan refresh attempt after first token", throwable)); + } + // We trigger refresh so the next attempt will use a fresh plan + executeQueryCallContext.triggerImmediateRefreshOfPreparedQuery(); + responseObserver.onError(throwable); + return; + } + ApiCallContext contextWithAbsoluteDeadline = + Optional.ofNullable(grpcCallContext) + .map(c -> c.withCallOptions(grpcCallContext.getCallOptions().withDeadline(deadline))) + .orElse(null); + inner.call(request, observer, contextWithAbsoluteDeadline); } - static final class MetadataObserver extends SafeResponseObserver { + // Checks for an attempt timeout first, then a total timeout. If found, converts the timeout + // to an absolute deadline. Adjusts totalTimeout based on the time since startTimeOfOverallRequest + private static @Nullable Deadline getDeadline( + GrpcCallContext grpcCallContext, Instant startTimeOfOverallRequest) { + Optional attemptDeadline = + Optional.ofNullable(grpcCallContext) + .flatMap(c -> Optional.ofNullable(c.getTimeoutDuration())) + .map(d -> Deadline.after(d.toNanos(), TimeUnit.NANOSECONDS)); + if (attemptDeadline.isPresent()) { + return attemptDeadline.get(); + } + return Optional.ofNullable(grpcCallContext) + .flatMap(c -> Optional.ofNullable(c.getRetrySettings())) + .map(RetrySettings::getTotalTimeoutDuration) + // TotalTimeout of zero means there is no timeout + .filter(duration -> !duration.isZero()) + .map( + d -> { + Duration elapsedTime = Duration.between(startTimeOfOverallRequest, Instant.now()); + Duration remaining = d.minus(elapsedTime); + // zero is treated as no deadline, so if full deadline is elapsed pass 1 nano + long adjusted = Math.max(remaining.getNano(), 1); + return Deadline.after(adjusted, TimeUnit.NANOSECONDS); + }) + .orElse(null); + } + + @InternalApi + static boolean isPlanRefreshError(Throwable t) { + if (!(t instanceof ApiException)) { + return false; + } + ApiException e = (ApiException) t; + if (!e.getStatusCode().getCode().equals(Code.FAILED_PRECONDITION)) { + return false; + } + if (e.getErrorDetails() == null) { + return false; + } + PreconditionFailure preconditionFailure = e.getErrorDetails().getPreconditionFailure(); + if (preconditionFailure == null) { + return false; + } + for (Violation violation : preconditionFailure.getViolationsList()) { + if (violation.getType().contains("PREPARED_QUERY_EXPIRED")) { + return true; + } + } + return false; + } + + static final class PlanRefreshingObserver extends SafeResponseObserver { private final ExecuteQueryCallContext callContext; private final ResponseObserver outerObserver; @@ -64,7 +161,7 @@ static final class MetadataObserver extends SafeResponseObserver outerObserver, ExecuteQueryCallContext callContext) { super(outerObserver); this.outerObserver = outerObserver; @@ -103,12 +200,23 @@ protected void onResponseImpl(ExecuteQueryResponse response) { @Override protected void onErrorImpl(Throwable throwable) { - // TODO translate plan refresh errors and trigger plan refresh - - // Note that we do not set exceptions on the metadata future here. This - // needs to be done after the retries, so that retryable errors aren't set on - // the future - outerObserver.onError(throwable); + boolean refreshPlan = isPlanRefreshError(throwable); + // If we've received a resume token we shouldn't receive this error. Safeguard against + // accidentally changing the schema mid-response though + if (refreshPlan && !hasReceivedResumeToken) { + callContext.triggerImmediateRefreshOfPreparedQuery(); + outerObserver.onError( + new ApiException(throwable, GrpcStatusCode.of(Status.Code.FAILED_PRECONDITION), true)); + } else if (refreshPlan) { + outerObserver.onError( + new IllegalStateException( + "Unexpected plan refresh attempt after first token", throwable)); + } else { + // Note that we do not set exceptions on the metadata future here. This + // needs to be done after the retries, so that retryable errors aren't set on + // the future + outerObserver.onError(throwable); + } } @Override diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java index f6d50e1c39..eaf5a40abb 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java @@ -21,6 +21,7 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.gax.batching.Batcher; +import com.google.api.gax.rpc.ClientContext; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.gax.rpc.UnaryCallable; @@ -78,6 +79,7 @@ public class BigtableDataClientTests { @Rule public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.WARN); @Mock private EnhancedBigtableStub mockStub; + @Mock private ClientContext mockContext; @Mock(answer = Answers.RETURNS_DEEP_STUBS) private ServerStreamingCallable mockReadRowsCallable; @@ -1075,4 +1077,14 @@ public void prepareQueryTest() { bigtableDataClient.prepareStatement(query, paramTypes); Mockito.verify(mockPrepareQueryCallable).call(PrepareQueryRequest.create(query, paramTypes)); } + + @Test + public void executeQueryMustUseSameClientAsPrepare() { + Mockito.when(mockStub.prepareQueryCallable()).thenReturn(mockPrepareQueryCallable); + + String query = "SELECT * FROM table"; + Map> paramTypes = new HashMap<>(); + bigtableDataClient.prepareStatement(query, paramTypes); + Mockito.verify(mockPrepareQueryCallable).call(PrepareQueryRequest.create(query, paramTypes)); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java new file mode 100644 index 0000000000..28775af904 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java @@ -0,0 +1,417 @@ +/* + * 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. + */ +package com.google.cloud.bigtable.data.v2.internal; + +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.api.core.ApiFutures; +import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.rpc.ApiExceptions; +import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.bigtable.v2.ResultSetMetadata; +import com.google.cloud.bigtable.data.v2.BigtableDataClient; +import com.google.cloud.bigtable.data.v2.BigtableDataSettings; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PrepareQueryState; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryData; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryVersion; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.PrepareRpcExpectation; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.TestBigtableSqlService; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import io.grpc.Status.Code; +import io.grpc.testing.GrpcServerRule; +import java.io.IOException; +import java.lang.ref.WeakReference; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class PreparedStatementImplTest { + + private static final ResultSetMetadata METADATA_PROTO = + metadata(columnMetadata("_key", bytesType()), columnMetadata("p", stringType())); + + @Rule public GrpcServerRule serverRule = new GrpcServerRule(); + private TestBigtableSqlService service; + private BigtableDataClient client; + private Map> paramTypes; + private int prepareAttempts; + + @Before + public void setUp() throws IOException { + service = new TestBigtableSqlService(); + serverRule.getServiceRegistry().addService(service); + BigtableDataSettings.Builder settings = + BigtableDataSettings.newBuilder() + .setProjectId(TestBigtableSqlService.DEFAULT_PROJECT_ID) + .setInstanceId(TestBigtableSqlService.DEFAULT_INSTANCE_ID) + .setAppProfileId(TestBigtableSqlService.DEFAULT_APP_PROFILE_ID) + .setCredentialsProvider(NoCredentialsProvider.create()); + settings + .stubSettings() + .setTransportChannelProvider( + FixedTransportChannelProvider.create( + GrpcTransportChannel.create(serverRule.getChannel()))) + // Refreshing channel doesn't work with FixedTransportChannelProvider + .setRefreshingChannel(false) + .build(); + // Remove log noise from client side metrics + settings.setMetricsProvider(NoopMetricsProvider.INSTANCE); + prepareAttempts = + settings.stubSettings().prepareQuerySettings().retrySettings().getMaxAttempts(); + client = BigtableDataClient.create(settings.build()); + paramTypes = ImmutableMap.of("param", SqlType.string()); + } + + private PreparedStatementImpl getDefaultPrepareStatement() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("plan"), + METADATA_PROTO, + // Plan expires right away + Instant.now()))); + return (PreparedStatementImpl) + client.prepareStatement("SELECT _key, @param AS p FROM table", paramTypes); + } + + @After + public void tearDown() { + if (client != null) { + client.close(); + } + } + + @Test + public void testBackgroundRefresh() throws InterruptedException, ExecutionException { + PreparedStatementImpl preparedStatement = getDefaultPrepareStatement(); + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse(ByteString.copyFromUtf8("plan2"), METADATA_PROTO, Instant.now()))); + // Refresh won't be triggered until this call + PreparedQueryData initialPlan = preparedStatement.getLatestPrepareResponse(); + PrepareResponse initialResponse = initialPlan.prepareFuture().get(); + // wait for the second call + do { + Thread.sleep(10); + } while (service.prepareCount < 2); + PreparedQueryData updatedPlan = preparedStatement.getLatestPrepareResponse(); + PrepareResponse updatedResponse = updatedPlan.prepareFuture().get(); + assertThat(updatedPlan.version()).isNotEqualTo(initialPlan.version()); + assertThat(initialResponse.preparedQuery()).isEqualTo(ByteString.copyFromUtf8("plan")); + assertThat(initialResponse.resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA_PROTO)); + assertThat(updatedResponse.preparedQuery()).isEqualTo(ByteString.copyFromUtf8("plan2")); + assertThat(updatedResponse.resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA_PROTO)); + // We don't expect any additional calls + assertThat(service.prepareCount).isEqualTo(2); + } + + @Test + public void noRefreshBeforeExpiryWindow() throws ExecutionException, InterruptedException { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @other AS o FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("other_plan"), + METADATA_PROTO, + // Don't expire + Instant.now().plus(Duration.ofMinutes(10))))); + PreparedStatementImpl unexpired = + (PreparedStatementImpl) + client.prepareStatement("SELECT _key, @other AS o FROM table", paramTypes); + // Don't expect any refresh + PreparedQueryData initialPlan = unexpired.getLatestPrepareResponse(); + PrepareResponse initialResponse = initialPlan.prepareFuture().get(); + + assertThat(initialResponse.preparedQuery()).isEqualTo(ByteString.copyFromUtf8("other_plan")); + assertThat(service.prepareCount).isEqualTo(1); + } + + @Test + public void testMarkExpiredAndStartRefresh() throws ExecutionException, InterruptedException { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("plan"), + METADATA_PROTO, + // Plan expires right away + Instant.now().plusSeconds(2L)))); + PreparedStatementImpl preparedStatement = + (PreparedStatementImpl) + client.prepareStatement("SELECT _key, @param AS p FROM table", paramTypes); + PreparedQueryData initialPlan = preparedStatement.getLatestPrepareResponse(); + PrepareResponse initialPrepareResponse = initialPlan.prepareFuture().get(); + + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("hardRefreshPlan"), + METADATA_PROTO, + Instant.now().plus(Duration.ofMinutes(10))))); + + PreparedQueryData updatedPlan = + preparedStatement.markExpiredAndStartRefresh(initialPlan.version()); + PrepareResponse updatedPrepareResponse = updatedPlan.prepareFuture().get(); + + assertThat(updatedPlan.version()).isNotEqualTo(initialPlan.version()); + assertThat(initialPrepareResponse.preparedQuery()).isEqualTo(ByteString.copyFromUtf8("plan")); + assertThat(initialPrepareResponse.resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA_PROTO)); + assertThat(updatedPrepareResponse.preparedQuery()) + .isEqualTo(ByteString.copyFromUtf8("hardRefreshPlan")); + assertThat(updatedPrepareResponse.resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA_PROTO)); + // We don't expect any additional calls + assertThat(service.prepareCount).isEqualTo(2); + } + + @Test + public void testConcurrentBackgroundRefreshCalls() + throws InterruptedException, ExecutionException { + PreparedStatementImpl preparedStatement = getDefaultPrepareStatement(); + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("plan2"), + METADATA_PROTO, + Instant.now().plus(Duration.ofMinutes(10))))); + ExecutorService executor = Executors.newFixedThreadPool(50); + List> callableList = new ArrayList<>(); + for (int i = 0; i < 50; i++) { + callableList.add(preparedStatement::getLatestPrepareResponse); + } + List> results = executor.invokeAll(callableList); + executor.shutdown(); + boolean done = executor.awaitTermination(1, TimeUnit.MINUTES); + assertThat(done).isTrue(); + assertThat(service.prepareCount).isEqualTo(2); + for (Future prepareFuture : results) { + PreparedQueryData response = prepareFuture.get(); + assertThat(response.prepareFuture().get().preparedQuery()) + .isIn( + // Some will get the first plan, some might get the result of refresh + ImmutableList.of(ByteString.copyFromUtf8("plan"), ByteString.copyFromUtf8("plan2"))); + } + } + + @Test + public void testConcurrentMarkExpiredAndStartRefreshCalls() + throws InterruptedException, ExecutionException { + PreparedStatementImpl preparedStatement = getDefaultPrepareStatement(); + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("plan2"), + METADATA_PROTO, + Instant.now().plus(Duration.ofMinutes(10))))); + PreparedQueryData initialData = preparedStatement.getLatestPrepareResponse(); + ExecutorService executor = Executors.newFixedThreadPool(50); + List> callableList = new ArrayList<>(); + for (int i = 0; i < 50; i++) { + callableList.add(() -> preparedStatement.markExpiredAndStartRefresh(initialData.version())); + } + List> results = executor.invokeAll(callableList); + executor.shutdown(); + boolean done = executor.awaitTermination(1, TimeUnit.MINUTES); + assertThat(done).isTrue(); + for (Future refreshFuture : results) { + PreparedQueryData response = refreshFuture.get(); + assertThat(response.version()).isNotEqualTo(initialData.version()); + assertThat(response.prepareFuture().get().resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(METADATA_PROTO)); + } + assertThat(service.prepareCount).isEqualTo(2); + } + + @Test + public void testPrepareFailuresAreRetried() throws ExecutionException, InterruptedException { + PreparedStatementImpl preparedStatement = getDefaultPrepareStatement(); + int failures = 0; + // Exhaust all the retries w unavailables + for (int i = 0; i <= prepareAttempts; i++) { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWithStatus(Code.UNAVAILABLE)); + failures++; + } + // Now exhaust all the retries again w deadline exceeded + for (int i = 0; i <= prepareAttempts; i++) { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .respondWithStatus(Code.DEADLINE_EXCEEDED)); + failures++; + } + // then succeed + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .withDelay(Duration.ofMillis(20)) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("plan2"), + METADATA_PROTO, + Instant.now().plus(Duration.ofMinutes(10))))); + PreparedQueryData initialData = preparedStatement.getLatestPrepareResponse(); + PreparedQueryData nextData = + preparedStatement.markExpiredAndStartRefresh(initialData.version()); + + assertThat(nextData.prepareFuture().get().preparedQuery()) + .isEqualTo(ByteString.copyFromUtf8("plan2")); + // initial request + failures + final success + assertThat(service.prepareCount).isEqualTo(1 + failures + 1); + } + + @Test + public void garbageCollectionWorksWhenRetryIsOngoing() throws InterruptedException { + PreparedStatementImpl preparedStatement = getDefaultPrepareStatement(); + + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT _key, @param AS p FROM table") + .withParamTypes(paramTypes) + .withDelay(Duration.ofSeconds(1)) + // Return a permanent error so the stub doesn't retry + .respondWithStatus(Code.INTERNAL)); + WeakReference weakRef = new WeakReference<>(preparedStatement); + PreparedQueryVersion initialPlanId = preparedStatement.getLatestPrepareResponse().version(); + PreparedQueryData next = preparedStatement.markExpiredAndStartRefresh(initialPlanId); + preparedStatement = null; + for (int i = 0; i < 5; i++) { + // This isn't guaranteed to run GC, so call it a few times. Testing has shown that this + // is enough to prevent any flakes in 1000 runs + System.gc(); + Thread.sleep(10); + } + assertThat(service.prepareCount).isEqualTo(2); + assertThat(weakRef.get()).isNull(); + // The plan refresh stops retrying after the PreparedStatement is garbage collected. + // Because this means it isn't needed anymore, we don't want to keep refreshing. + assertThrows( + RuntimeException.class, + () -> ApiExceptions.callAndTranslateApiException(next.prepareFuture())); + } + + @Test + public void testPrepareQueryStateInitialState() throws ExecutionException, InterruptedException { + ResultSetMetadata md = metadata(columnMetadata("strCol", stringType())); + PrepareQueryState state = + PrepareQueryState.createInitialState( + PrepareResponse.fromProto(prepareResponse(ByteString.copyFromUtf8("plan"), md))); + assertThat(state.current().prepareFuture().isDone()).isTrue(); + assertThat(state.current().prepareFuture().get().resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(md)); + assertThat(state.maybeBackgroundRefresh()).isEmpty(); + } + + @Test + public void testPrepareQueryStateWithBackgroundPlan() + throws ExecutionException, InterruptedException { + ResultSetMetadata md = metadata(columnMetadata("strCol", stringType())); + PrepareQueryState state = + PrepareQueryState.createInitialState( + PrepareResponse.fromProto(prepareResponse(ByteString.copyFromUtf8("plan"), md))); + + PrepareQueryState withBackgroundPlan = + state.withBackgroundPlan( + ApiFutures.immediateFuture(PrepareResponse.fromProto(prepareResponse(md)))); + assertThat(withBackgroundPlan.current().prepareFuture().isDone()).isTrue(); + assertThat(withBackgroundPlan.current().prepareFuture().get().resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(md)); + assertThat(withBackgroundPlan.current().version()).isEqualTo(state.current().version()); + assertThat(withBackgroundPlan.maybeBackgroundRefresh()).isPresent(); + assertThat(withBackgroundPlan.maybeBackgroundRefresh().get().version()) + .isNotEqualTo(withBackgroundPlan.current().version()); + assertThat( + withBackgroundPlan + .maybeBackgroundRefresh() + .get() + .prepareFuture() + .get() + .resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(md)); + } + + @Test + public void testPrepareQueryStatePromoteBackgroundPlan() + throws ExecutionException, InterruptedException { + ResultSetMetadata md = metadata(columnMetadata("strCol", stringType())); + PrepareQueryState state = + PrepareQueryState.createInitialState( + PrepareResponse.fromProto(prepareResponse(ByteString.copyFromUtf8("plan"), md))); + PrepareQueryState withBackgroundPlan = + state.withBackgroundPlan( + ApiFutures.immediateFuture(PrepareResponse.fromProto(prepareResponse(md)))); + PrepareQueryState finalState = withBackgroundPlan.promoteBackgroundPlan(); + + assertThat(finalState.current().version()) + .isEqualTo(withBackgroundPlan.maybeBackgroundRefresh().get().version()); + assertThat(finalState.current().prepareFuture().get().resultSetMetadata()) + .isEqualTo(ProtoResultSetMetadata.fromProto(md)); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java index 2aee489f11..4c3e9443d0 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/ResultSetImplTest.java @@ -21,6 +21,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.boolValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.callContext; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.dateType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.dateValue; @@ -33,7 +34,6 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.mapValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; -import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.structField; @@ -52,6 +52,7 @@ import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import com.google.cloud.bigtable.data.v2.stub.sql.ExecuteQueryCallContext; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStreamImpl; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; @@ -76,10 +77,8 @@ private static ResultSet resultSetWithFakeStream( SettableApiFuture future = SettableApiFuture.create(); ResultSetMetadata metadata = ProtoResultSetMetadata.fromProto(protoMetadata); future.set(metadata); - PrepareResponse response = PrepareResponse.fromProto(prepareResponse(protoMetadata)); - PreparedStatement preparedStatement = PreparedStatementImpl.create(response, new HashMap<>()); - ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create(preparedStatement.bind().build(), future); + PreparedStatement preparedStatement = SqlProtoFactory.preparedStatement(protoMetadata); + ExecuteQueryCallContext fakeCallContext = callContext(preparedStatement.bind().build(), future); return ResultSetImpl.create(SqlServerStreamImpl.create(future, stream.call(fakeCallContext))); } @@ -322,12 +321,10 @@ public void getMetadata_unwrapsExecutionExceptions() { SettableApiFuture metadataFuture = SettableApiFuture.create(); ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>(Collections.emptyList()); - PrepareResponse prepareResponse = - PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))); PreparedStatement preparedStatement = - PreparedStatementImpl.create(prepareResponse, new HashMap<>()); + SqlProtoFactory.preparedStatement(metadata(columnMetadata("foo", stringType()))); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); + callContext(preparedStatement.bind().build(), metadataFuture); ResultSet rs = ResultSetImpl.create( SqlServerStreamImpl.create(metadataFuture, stream.call(fakeCallContext))); @@ -341,12 +338,10 @@ public void getMetadata_returnsNonRuntimeExecutionExceptionsWrapped() { SettableApiFuture metadataFuture = SettableApiFuture.create(); ServerStreamingStashCallable stream = new ServerStreamingStashCallable<>(Collections.emptyList()); - PrepareResponse prepareResponse = - PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))); PreparedStatement preparedStatement = - PreparedStatementImpl.create(prepareResponse, new HashMap<>()); + SqlProtoFactory.preparedStatement(metadata(columnMetadata("foo", stringType()))); ExecuteQueryCallContext fakeCallContext = - ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); + callContext(preparedStatement.bind().build(), metadataFuture); ResultSet rs = ResultSetImpl.create( SqlServerStreamImpl.create(metadataFuture, stream.call(fakeCallContext))); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java index f3a37f99d5..c089138286 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatementTest.java @@ -31,6 +31,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Value; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.nullValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.timestampType; @@ -45,7 +46,6 @@ import com.google.cloud.Date; import com.google.cloud.Timestamp; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.protobuf.ByteString; import java.time.Duration; @@ -82,7 +82,7 @@ public static BoundStatement.Builder boundStatementBuilder(ColumnMetadata... par } // This doesn't impact bound statement, but set it so it looks like a real response Instant expiry = Instant.now().plus(Duration.ofMinutes(1)); - return PreparedStatementImpl.create( + return preparedStatement( PrepareResponse.fromProto( PrepareQueryResponse.newBuilder() .setPreparedQuery(EXPECTED_PREPARED_QUERY) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index 930f91b488..7bcc4650b3 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -18,6 +18,8 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; @@ -70,7 +72,6 @@ import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; @@ -91,6 +92,7 @@ import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider; import com.google.cloud.bigtable.data.v2.stub.sql.ExecuteQueryCallable; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; @@ -159,12 +161,11 @@ public class EnhancedBigtableStubTest { private static final Duration WATCHDOG_CHECK_DURATION = Duration.ofMillis(100); private static final PrepareResponse PREPARE_RESPONSE = PrepareResponse.fromProto( - PrepareQueryResponse.newBuilder() - .setPreparedQuery(ByteString.copyFromUtf8(WAIT_TIME_QUERY)) - .setMetadata(metadata(columnMetadata("foo", stringType()))) - .build()); + prepareResponse( + ByteString.copyFromUtf8(WAIT_TIME_QUERY), + metadata(columnMetadata("foo", stringType())))); private static final PreparedStatement WAIT_TIME_PREPARED_STATEMENT = - PreparedStatementImpl.create(PREPARE_RESPONSE, new HashMap<>()); + preparedStatement(PREPARE_RESPONSE, new HashMap<>()); private Server server; private MetadataInterceptor metadataInterceptor; @@ -901,12 +902,9 @@ public void testCreateExecuteQueryCallable() throws InterruptedException { ExecuteQueryCallable streamingCallable = enhancedBigtableStub.createExecuteQueryCallable(); PrepareResponse prepareResponse = PrepareResponse.fromProto( - PrepareQueryResponse.newBuilder() - .setPreparedQuery(ByteString.copyFromUtf8("abc")) - .setMetadata(metadata(columnMetadata("foo", stringType()))) - .build()); - PreparedStatement preparedStatement = - PreparedStatementImpl.create(prepareResponse, new HashMap<>()); + SqlProtoFactory.prepareResponse( + ByteString.copyFromUtf8("abc"), metadata(columnMetadata("foo", stringType())))); + PreparedStatement preparedStatement = preparedStatement(prepareResponse, new HashMap<>()); SqlServerStream sqlServerStream = streamingCallable.call(preparedStatement.bind().build()); ExecuteQueryRequest expectedRequest = ExecuteQueryRequest.newBuilder() diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java index 7c5d190fbf..995c1f0dbd 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/HeadersTest.java @@ -17,7 +17,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; -import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; @@ -42,8 +42,6 @@ import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Mutation; import com.google.cloud.bigtable.data.v2.models.Query; @@ -175,10 +173,7 @@ public void readModifyWriteTest() { @Test public void executeQueryTest() { PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse(metadata(columnMetadata("foo", stringType())))), - new HashMap<>()); + preparedStatement(metadata(columnMetadata("foo", stringType()))); client.executeQuery(preparedStatement.bind().build()); verifyHeaderSent(true); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java index f688d96669..14cfd25f66 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContextTest.java @@ -19,24 +19,34 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.DeadlineExceededException; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryData; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatementRefreshTimeoutException; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.FakePreparedStatement; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; +import io.grpc.Deadline; +import io.grpc.Status.Code; import java.util.Map; import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -49,17 +59,18 @@ public class ExecuteQueryCallContextTest { private static final Map> PARAM_TYPES = ImmutableMap.of("foo", SqlType.string()); private static final PreparedStatement PREPARED_STATEMENT = - PreparedStatementImpl.create( + preparedStatement( PrepareResponse.fromProto(prepareResponse(PREPARED_QUERY, METADATA)), PARAM_TYPES); @Test public void testToRequest() { ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create( + SqlProtoFactory.callContext( PREPARED_STATEMENT.bind().setStringParam("foo", "val").build(), SettableApiFuture.create()); RequestContext requestContext = RequestContext.create("project", "instance", "profile"); - ExecuteQueryRequest request = callContext.toRequest(requestContext); + ExecuteQueryRequest request = + callContext.buildRequestWithDeadline(requestContext, Deadline.after(1, TimeUnit.MINUTES)); assertThat(request.getPreparedQuery()).isEqualTo(PREPARED_QUERY); assertThat(request.getAppProfileId()).isEqualTo("profile"); @@ -73,7 +84,7 @@ public void testToRequest() { public void testFirstResponseReceived() throws ExecutionException, InterruptedException { SettableApiFuture mdFuture = SettableApiFuture.create(); ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create( + SqlProtoFactory.callContext( PREPARED_STATEMENT.bind().setStringParam("foo", "val").build(), mdFuture); callContext.finalizeMetadata(); @@ -85,7 +96,7 @@ public void testFirstResponseReceived() throws ExecutionException, InterruptedEx public void testSetMetadataException() { SettableApiFuture mdFuture = SettableApiFuture.create(); ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create( + SqlProtoFactory.callContext( PREPARED_STATEMENT.bind().setStringParam("foo", "val").build(), mdFuture); callContext.setMetadataException(new RuntimeException("test")); @@ -93,4 +104,83 @@ public void testSetMetadataException() { ExecutionException e = assertThrows(ExecutionException.class, mdFuture::get); assertThat(e.getCause()).isInstanceOf(RuntimeException.class); } + + @Test + public void testBuildRequestAttemptDeadline() { + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + SettableApiFuture mdFuture = SettableApiFuture.create(); + PreparedQueryData initialPlan = PreparedQueryData.create(SettableApiFuture.create()); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create( + new FakePreparedStatement() + // Reuse the same plan since we wont call refresh + .withUpdatedPlans(initialPlan, initialPlan) + .bind() + .build(), + mdFuture); + + assertThrows( + PreparedStatementRefreshTimeoutException.class, + () -> + callContext.buildRequestWithDeadline( + requestContext, Deadline.after(2, TimeUnit.MILLISECONDS))); + } + + @Test + public void testHardRefreshUpdatesPreparedQuery() { + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + SettableApiFuture mdFuture = SettableApiFuture.create(); + ExecuteQueryCallContext callContext = + SqlProtoFactory.callContext(new FakePreparedStatement().bind().build(), mdFuture); + + callContext.triggerImmediateRefreshOfPreparedQuery(); + ExecuteQueryRequest updatedRequest = + callContext.buildRequestWithDeadline( + requestContext, Deadline.after(10, TimeUnit.MILLISECONDS)); + assertThat(updatedRequest.getPreparedQuery()) + .isEqualTo(ByteString.copyFromUtf8("refreshedPlan")); + } + + @Test + public void testResumeToken() { + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + SettableApiFuture mdFuture = SettableApiFuture.create(); + ExecuteQueryCallContext callContext = + SqlProtoFactory.callContext(new FakePreparedStatement().bind().build(), mdFuture); + callContext.setLatestResumeToken(ByteString.copyFromUtf8("token")); + + assertThat(callContext.hasResumeToken()).isTrue(); + assertThat( + callContext + .buildRequestWithDeadline( + requestContext, Deadline.after(100, TimeUnit.MILLISECONDS)) + .getResumeToken()) + .isEqualTo(ByteString.copyFromUtf8("token")); + } + + @Test + public void testPrepareExceptionIsRetryable() { + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + SettableApiFuture mdFuture = SettableApiFuture.create(); + ExecuteQueryCallContext callContext = + SqlProtoFactory.callContext( + new FakePreparedStatement() + .withUpdatedPlans( + PreparedQueryData.create( + ApiFutures.immediateFailedFuture( + new DeadlineExceededException( + null, GrpcStatusCode.of(Code.DEADLINE_EXCEEDED), false))), + null) + .bind() + .build(), + mdFuture); + + ApiException e = + assertThrows( + ApiException.class, + () -> + callContext.buildRequestWithDeadline( + requestContext, Deadline.after(10, TimeUnit.MILLISECONDS))); + assertThat(e.isRetryable()).isTrue(); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java index 62286368a7..8643dc24da 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java @@ -33,12 +33,10 @@ import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.ProtoSqlRow; -import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.internal.SqlRow; -import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; -import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement.Builder; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; @@ -61,25 +59,26 @@ @RunWith(JUnit4.class) public class ExecuteQueryCallableTest { - private static final class FakePreparedStatement implements PreparedStatement { + private static final class FakePreparedStatement extends PreparedStatementImpl { - @Override - public Builder bind() { - return new BoundStatement.Builder(this, new HashMap<>()); + public FakePreparedStatement() { + super( + PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("foo", stringType())))), + new HashMap<>(), + null, + null); } @Override - public PrepareResponse getPrepareResponse() { - return PrepareResponse.fromProto( - prepareResponse(metadata(columnMetadata("foo", stringType())))); + public PreparedQueryData markExpiredAndStartRefresh( + PreparedQueryVersion expiredPreparedQueryVersion) { + return getLatestPrepareResponse(); } @Override - public void close() throws Exception {} + public void assertUsingSameStub(EnhancedBigtableStub stub) {} } - private static final RequestContext REQUEST_CONTEXT = - RequestContext.create("fake-project", "fake-instance", "fake-profile"); private static final PreparedStatement PREPARED_STATEMENT = new FakePreparedStatement(); private Server server; @@ -113,7 +112,7 @@ public void testCallContextAndServerStreamSetup() { Collections.singletonList(stringValue("foo"))); ServerStreamingStashCallable innerCallable = new ServerStreamingStashCallable<>(Collections.singletonList(row)); - ExecuteQueryCallable callable = new ExecuteQueryCallable(innerCallable, REQUEST_CONTEXT); + ExecuteQueryCallable callable = new ExecuteQueryCallable(innerCallable); SqlServerStream stream = callable.call(PREPARED_STATEMENT.bind().build()); assertThat(stream.metadataFuture()) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java index 7bd860115b..d42529209f 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryResumptionStrategyTest.java @@ -19,7 +19,7 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; -import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.extensions.proto.ProtoTruth.assertThat; @@ -27,13 +27,13 @@ import com.google.api.core.SettableApiFuture; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; import com.google.protobuf.ByteString; -import java.util.HashMap; +import io.grpc.Deadline; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -42,16 +42,14 @@ public class ExecuteQueryResumptionStrategyTest { @Test - public void tracksResumeToken() { + public void tracksResumeToken() throws ExecutionException, InterruptedException { ExecuteQueryResumptionStrategy resumptionStrategy = new ExecuteQueryResumptionStrategy(); PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto(prepareResponse(metadata(columnMetadata("s", stringType())))), - new HashMap<>()); + preparedStatement(metadata(columnMetadata("s", stringType()))); SettableApiFuture mdFuture = SettableApiFuture.create(); ExecuteQueryCallContext callContext = - ExecuteQueryCallContext.create(preparedStatement.bind().build(), mdFuture); + SqlProtoFactory.callContext(preparedStatement.bind().build(), mdFuture); resumptionStrategy.processResponse( partialResultSetWithToken(ByteString.copyFromUtf8("token"), stringValue("s"))); @@ -60,7 +58,9 @@ public void tracksResumeToken() { ExecuteQueryCallContext updatedCallContext = resumptionStrategy.getResumeRequest(callContext); assertThat( - updatedCallContext.toRequest(RequestContext.create("project", "instance", "profile"))) + updatedCallContext.buildRequestWithDeadline( + RequestContext.create("project", "instance", "profile"), + Deadline.after(1, TimeUnit.MINUTES))) .isEqualTo( ExecuteQueryRequest.newBuilder() .setInstanceName(NameUtil.formatInstanceName("project", "instance")) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java index ae295c1720..734fbe14a1 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java @@ -15,11 +15,14 @@ */ package com.google.cloud.bigtable.data.v2.stub.sql; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesType; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.bytesValue; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSets; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.planRefreshError; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; @@ -29,40 +32,36 @@ import com.google.api.gax.core.NoCredentialsProvider; import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.FailedPreconditionException; import com.google.api.gax.rpc.FixedTransportChannelProvider; import com.google.api.gax.rpc.StatusCode; -import com.google.bigtable.v2.BigtableGrpc; -import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.ResultSetMetadata; import com.google.bigtable.v2.Value; import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; -import com.google.cloud.bigtable.data.v2.internal.NameUtil; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatementRefreshTimeoutException; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; +import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.ExecuteRpcExpectation; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.PrepareRpcExpectation; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.TestBigtableSqlService; import com.google.cloud.bigtable.gaxx.reframing.IncompleteStreamException; import com.google.common.collect.ImmutableMap; -import com.google.common.truth.Truth; import com.google.protobuf.ByteString; -import io.grpc.Status; import io.grpc.Status.Code; -import io.grpc.stub.StreamObserver; import io.grpc.testing.GrpcServerRule; import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; +import java.lang.ref.WeakReference; +import java.time.Duration; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Queue; -import java.util.concurrent.LinkedBlockingDeque; -import javax.annotation.Nullable; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -72,28 +71,21 @@ @RunWith(JUnit4.class) public class ExecuteQueryRetryTest { - private static final String PROJECT_ID = "fake-project"; - private static final String INSTANCE_ID = "fake-instance"; - private static final String APP_PROFILE_ID = "fake-app-profile"; private static final ByteString PREPARED_QUERY = ByteString.copyFromUtf8("foo"); private static final ResultSetMetadata DEFAULT_METADATA = metadata(columnMetadata("strCol", stringType())); @Rule public GrpcServerRule serverRule = new GrpcServerRule(); - private TestBigtableService service; + private TestBigtableSqlService service; private BigtableDataClient client; private PreparedStatement preparedStatement; - @Before - public void setUp() throws IOException { - service = new TestBigtableService(); - serverRule.getServiceRegistry().addService(service); - + public static BigtableDataSettings.Builder defaultSettings(GrpcServerRule serverRule) { BigtableDataSettings.Builder settings = BigtableDataSettings.newBuilder() - .setProjectId(PROJECT_ID) - .setInstanceId(INSTANCE_ID) - .setAppProfileId(APP_PROFILE_ID) + .setProjectId(TestBigtableSqlService.DEFAULT_PROJECT_ID) + .setInstanceId(TestBigtableSqlService.DEFAULT_INSTANCE_ID) + .setAppProfileId(TestBigtableSqlService.DEFAULT_APP_PROFILE_ID) .setCredentialsProvider(NoCredentialsProvider.create()); settings @@ -104,12 +96,23 @@ public void setUp() throws IOException { // Refreshing channel doesn't work with FixedTransportChannelProvider .setRefreshingChannel(false) .build(); + // Remove log noise from client side metrics + settings.setMetricsProvider(NoopMetricsProvider.INSTANCE); + return settings; + } - client = BigtableDataClient.create(settings.build()); - preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto(prepareResponse(PREPARED_QUERY, DEFAULT_METADATA)), - new HashMap<>()); + @Before + public void setUp() throws IOException { + service = new TestBigtableSqlService(); + serverRule.getServiceRegistry().addService(service); + client = BigtableDataClient.create(defaultSettings(serverRule).build()); + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith(prepareResponse(PREPARED_QUERY, DEFAULT_METADATA))); + preparedStatement = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + // Reset the count of RPCs + service.prepareCount--; } @After @@ -122,7 +125,7 @@ public void tearDown() { @Test public void testAllSuccesses() { service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .respondWith( partialResultSetWithoutToken(stringValue("foo")), partialResultSetWithoutToken(stringValue("bar")), @@ -147,16 +150,16 @@ public void testRetryOnInitialError() { // - First attempt immediately fails // - Second attempt returns 'foo', w a token, and succeeds // Expect result to be 'foo' - service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation(ExecuteRpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create().respondWith(partialResultSetWithToken(stringValue("foo")))); + ExecuteRpcExpectation.create().respondWith(partialResultSetWithToken(stringValue("foo")))); ResultSet rs = client.executeQuery(preparedStatement.bind().build()); assertThat(rs.next()).isTrue(); assertThat(rs.getString("strCol")).isEqualTo("foo"); assertThat(rs.next()).isFalse(); rs.close(); - assertThat(service.requestCount).isEqualTo(2); + assertThat(service.executeCount).isEqualTo(2); } @Test @@ -168,18 +171,18 @@ public void testResumptionToken() { // and then succeeds // We expect the results to contain all of the returned data (no reset batches) service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .respondWith( partialResultSetWithToken(ByteString.copyFromUtf8("token1"), stringValue("foo"))) .respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .withResumeToken(ByteString.copyFromUtf8("token1")) .respondWith( partialResultSetWithToken(ByteString.copyFromUtf8("token2"), stringValue("bar"))) .respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .withResumeToken(ByteString.copyFromUtf8("token2")) .respondWith( partialResultSetWithToken(ByteString.copyFromUtf8("final"), stringValue("baz")))); @@ -193,7 +196,7 @@ public void testResumptionToken() { assertThat(rs.getString("strCol")).isEqualTo("baz"); assertThat(rs.next()).isFalse(); rs.close(); - assertThat(service.requestCount).isEqualTo(3); + assertThat(service.executeCount).isEqualTo(3); } @Test @@ -205,7 +208,7 @@ public void testResetOnResumption() { // - Third attempt should resume w 'token1', we return 'baz' w reset & a token, succeed // Expect the results to be 'foo' and 'baz' service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .respondWith( partialResultSetWithToken(ByteString.copyFromUtf8("token1"), stringValue("foo")), // This is after the token so should be dropped @@ -219,13 +222,13 @@ public void testResetOnResumption() { stringValue("longerStringDiscard"), stringValue("discard")); service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .withResumeToken(ByteString.copyFromUtf8("token1")) // Skip the last response, so we don't send a new token .respondWith(chunkedResponses.get(0), chunkedResponses.get(1)) .respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .withResumeToken(ByteString.copyFromUtf8("token1")) .respondWith( partialResultSets(1, true, ByteString.copyFromUtf8("final"), stringValue("baz")) @@ -238,7 +241,7 @@ public void testResetOnResumption() { assertThat(rs.getString("strCol")).isEqualTo("baz"); assertThat(rs.next()).isFalse(); rs.close(); - assertThat(service.requestCount).isEqualTo(3); + assertThat(service.executeCount).isEqualTo(3); } @Test @@ -247,7 +250,7 @@ public void testErrorAfterFinalData() { // - Second attempt uses 'finalToken' and succeeds // Expect results to be 'foo', 'bar', 'baz' service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .respondWith( partialResultSetWithoutToken(stringValue("foo")), partialResultSetWithoutToken(stringValue("bar")), @@ -255,7 +258,7 @@ public void testErrorAfterFinalData() { ByteString.copyFromUtf8("finalToken"), stringValue("baz"))) .respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create().withResumeToken(ByteString.copyFromUtf8("finalToken"))); + ExecuteRpcExpectation.create().withResumeToken(ByteString.copyFromUtf8("finalToken"))); ResultSet rs = client.executeQuery(preparedStatement.bind().build()); assertThat(rs.getMetadata().getColumns()).hasSize(1); assertThat(rs.getMetadata().getColumns().get(0).name()).isEqualTo("strCol"); @@ -271,11 +274,9 @@ public void testErrorAfterFinalData() { rs.close(); } - // TODO test changing metadata when plan refresh is implemented - @Test public void permanentErrorPropagatesToMetadata() { - service.addExpectation(RpcExpectation.create().respondWithStatus(Code.INVALID_ARGUMENT)); + service.addExpectation(ExecuteRpcExpectation.create().respondWithStatus(Code.INVALID_ARGUMENT)); ResultSet rs = client.executeQuery(preparedStatement.bind().build()); ApiException e = assertThrows(ApiException.class, rs::getMetadata); @@ -291,7 +292,7 @@ public void exhaustedRetriesPropagatesToMetadata() throws IOException { .getMaxAttempts(); assertThat(attempts).isGreaterThan(1); for (int i = 0; i < attempts; i++) { - service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation(ExecuteRpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); } ResultSet rs = client.executeQuery(preparedStatement.bind().build()); @@ -301,10 +302,11 @@ public void exhaustedRetriesPropagatesToMetadata() throws IOException { @Test public void retryableErrorWithSuccessfulRetryDoesNotPropagateToMetadata() { - service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); - service.addExpectation(RpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation(ExecuteRpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation(ExecuteRpcExpectation.create().respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create().respondWith(tokenOnlyResultSet(ByteString.copyFromUtf8("t")))); + ExecuteRpcExpectation.create() + .respondWith(tokenOnlyResultSet(ByteString.copyFromUtf8("t")))); ResultSet rs = client.executeQuery(preparedStatement.bind().build()); assertThat(rs.getMetadata().getColumns()).hasSize(1); } @@ -313,20 +315,18 @@ public void retryableErrorWithSuccessfulRetryDoesNotPropagateToMetadata() { public void preservesParamsOnRetry() { Map> paramTypes = ImmutableMap.of("strParam", SqlType.string()); PreparedStatement preparedStatementWithParams = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse(metadata(columnMetadata("strCol", stringType())))), - paramTypes); + SqlProtoFactory.preparedStatement( + metadata(columnMetadata("strCol", stringType())), paramTypes); Map params = ImmutableMap.of("strParam", stringValue("foo").toBuilder().setType(stringType()).build()); service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .withParams(params) .respondWith( partialResultSetWithToken(ByteString.copyFromUtf8("token1"), stringValue("foo"))) .respondWithStatus(Code.UNAVAILABLE)); service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .withParams(params) .withResumeToken(ByteString.copyFromUtf8("token1")) .respondWith( @@ -347,95 +347,454 @@ public void failsOnCompleteWithOpenPartialBatch() { // Return 'foo' with no token, followed by ok // This should throw an error, as the backend has violated its contract service.addExpectation( - RpcExpectation.create() + ExecuteRpcExpectation.create() .respondWith(partialResultSetWithoutToken(stringValue("foo"))) .respondWithStatus(Code.OK)); ResultSet rs = client.executeQuery(preparedStatement.bind().build()); assertThrows(IncompleteStreamException.class, rs::next); } - private static class TestBigtableService extends BigtableGrpc.BigtableImplBase { - Queue expectations = new LinkedBlockingDeque<>(); - int requestCount = 0; + @Test + public void retryOnExpiredPlan() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("bytesCol", bytesType()))))); + // change the schema on refresh (this can happen for SELECT * queries for example) + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("baz"), + metadata(columnMetadata("strCol", stringType()))))); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("baz")) + .respondWith(partialResultSetWithToken(stringValue("foo")))); + + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isFalse(); + assertThat(service.executeCount).isEqualTo(2); + assertThat(service.prepareCount).isEqualTo(2); + } - void addExpectation(RpcExpectation expectation) { - expectations.add(expectation); - } + @Test + public void planRefreshAfterInitialPartialBatch() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("bytesCol", bytesType()))))); + // change the schema on refresh (this can happen for SELECT * queries for example) + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("baz"), + metadata(columnMetadata("strCol", stringType()))))); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + .respondWith(partialResultSetWithoutToken(bytesValue("b"))) + .respondWithStatus(Code.UNAVAILABLE)); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + // This creates one response w reset=true and a token + List singleResponseBatch = partialResultSets(1, stringValue("foo")); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("baz")) + .respondWith(singleResponseBatch.get(0))); - @Override - public void executeQuery( - ExecuteQueryRequest request, StreamObserver responseObserver) { - RpcExpectation expectedRpc = expectations.poll(); - requestCount++; - int requestIndex = requestCount - 1; - - Truth.assertWithMessage("Unexpected request#" + requestIndex + ":" + request.toString()) - .that(expectedRpc) - .isNotNull(); - Truth.assertWithMessage("Unexpected request#" + requestIndex) - .that(request) - .isEqualTo(expectedRpc.getExpectedRequest()); - - for (ExecuteQueryResponse response : expectedRpc.responses) { - responseObserver.onNext(response); - } - if (expectedRpc.statusCode.toStatus().isOk()) { - responseObserver.onCompleted(); - } else if (expectedRpc.exception != null) { - responseObserver.onError(expectedRpc.exception); - } else { - responseObserver.onError(expectedRpc.statusCode.toStatus().asRuntimeException()); - } - } + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("foo"); + assertThat(rs.next()).isFalse(); + assertThat(rs.getMetadata().getColumnType("strCol")).isEqualTo(SqlType.string()); + assertThat(service.executeCount).isEqualTo(3); + assertThat(service.prepareCount).isEqualTo(2); } - private static class RpcExpectation { - ExecuteQueryRequest.Builder request; - Status.Code statusCode; - @Nullable ApiException exception; - List responses; - - private RpcExpectation() { - this.request = ExecuteQueryRequest.newBuilder(); - this.request.setPreparedQuery(PREPARED_QUERY); - this.request.setInstanceName(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID)); - this.request.setAppProfileId(APP_PROFILE_ID); - this.statusCode = Code.OK; - this.responses = new ArrayList<>(); - } + @Test + public void planRefreshErrorAfterFirstTokenCausesError() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("bytesCol", bytesType()))))); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + .respondWith(partialResultSetWithToken(bytesValue("b"))) + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); - static RpcExpectation create() { - return new RpcExpectation(); - } + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + // We received a token so the client yields the data + assertThat(rs.getBytes("bytesCol").toStringUtf8()).isEqualTo("b"); + IllegalStateException e = assertThrows(IllegalStateException.class, rs::next); + assertThat(e.getCause()).isInstanceOf(FailedPreconditionException.class); + } - RpcExpectation withResumeToken(ByteString resumeToken) { - this.request.setResumeToken(resumeToken); - return this; + @Test + public void preparedStatementCanBeGarbageCollected() throws InterruptedException { + // Check for memory leaks since the PreparedStatement handles background refresh + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + service.addExpectation( + ExecuteRpcExpectation.create().respondWith(partialResultSetWithToken(stringValue("s")))); + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + WeakReference prepareWeakRef = new WeakReference<>(ps); + ResultSet rs = client.executeQuery(ps.bind().build()); + WeakReference resultSetWeakRef = new WeakReference<>(rs); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("s"); + assertThat(rs.next()).isFalse(); + rs.close(); + // Note that the result set holds a reference to the ResultSetMetadata that lives in + // the PreparedStatement. So prepare won't be gc'd until the ResultSet is null. + rs = null; + ps = null; + for (int i = 0; i < 5; i++) { + // This isn't guaranteed to run GC, so call it a few times. Testing has shown that this + // is enough to prevent any flakes in 1000 runs + System.gc(); + Thread.sleep(10); } + assertThat(resultSetWeakRef.get()).isNull(); + assertThat(prepareWeakRef.get()).isNull(); + } - RpcExpectation withParams(Map params) { - this.request.putAllParams(params); - return this; - } + @Test + public void planRefreshRespectsExecuteTotalTimeout() throws IOException { + BigtableDataSettings.Builder settings = defaultSettings(serverRule); + settings + .stubSettings() + .executeQuerySettings() + .setRetrySettings( + RetrySettings.newBuilder() + .setMaxAttempts(10) + .setTotalTimeoutDuration(Duration.ofMillis(30)) + .build()) + .build(); + settings.stubSettings().build(); + BigtableDataClient clientWithTimeout = BigtableDataClient.create(settings.build()); - RpcExpectation respondWithStatus(Status.Code code) { - this.statusCode = code; - return this; - } + // Initially return a prepare response without delay + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + // Trigger plan refresh + service.addExpectation( + ExecuteRpcExpectation.create() + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + service.addExpectation( + PrepareRpcExpectation.create() + .withDelay(Duration.ofSeconds(2)) + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("strCol", stringType()))))); + + PreparedStatement ps = + clientWithTimeout.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = clientWithTimeout.executeQuery(ps.bind().build()); + assertThrows(PreparedStatementRefreshTimeoutException.class, rs::next); + assertThat(service.prepareCount).isEqualTo(2); + } - RpcExpectation respondWithException(Status.Code code, ApiException exception) { - this.statusCode = code; - this.exception = exception; - return this; - } + @Test + public void planRefreshRespectsAttemptTimeout() throws IOException { + BigtableDataSettings.Builder settings = defaultSettings(serverRule); + settings + .stubSettings() + .executeQuerySettings() + .setRetrySettings( + RetrySettings.newBuilder() + // First attempt triggers plan refresh retry. + // Second should time out + .setMaxAttempts(2) + .setInitialRpcTimeoutDuration(Duration.ofMillis(10)) + .setMaxRpcTimeoutDuration(Duration.ofMinutes(10)) + .setTotalTimeoutDuration(Duration.ZERO) + .build()) + .build(); + settings.stubSettings().build(); + BigtableDataClient clientWithTimeout = BigtableDataClient.create(settings.build()); - RpcExpectation respondWith(ExecuteQueryResponse... responses) { - this.responses = Arrays.asList(responses); - return this; - } + // Initially return a prepare response without delay + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + // Trigger plan refresh + service.addExpectation( + ExecuteRpcExpectation.create() + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + // called after failed precondition + service.addExpectation( + PrepareRpcExpectation.create() + .withDelay(Duration.ofSeconds(2)) + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("strCol", stringType()))))); + + PreparedStatement ps = + clientWithTimeout.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = clientWithTimeout.executeQuery(ps.bind().build()); + assertThrows(PreparedStatementRefreshTimeoutException.class, rs::next); + assertThat(service.prepareCount).isEqualTo(2); + } - ExecuteQueryRequest getExpectedRequest() { - return this.request.build(); - } + @Test + public void executeRetriesPlanRefreshErrors() throws IOException { + // Initially return a prepare response without delay + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + // Trigger plan refresh + service.addExpectation( + ExecuteRpcExpectation.create() + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWithStatus(Code.UNAVAILABLE)); + // called after unavailable + service.addExpectation( + PrepareRpcExpectation.create() + .withDelay(Duration.ofSeconds(2)) + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("strCol", stringType()))))); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + .respondWith(partialResultSetWithToken(stringValue("s")))); + + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("s"); + assertThat(rs.next()).isFalse(); + assertThat(service.executeCount).isEqualTo(2); + assertThat(service.prepareCount).isEqualTo(3); + } + + @Test + public void prepareFailuresBurnExecuteAttempts() throws IOException { + BigtableDataSettings.Builder settings = defaultSettings(serverRule); + settings + .stubSettings() + .executeQuerySettings() + .setRetrySettings( + RetrySettings.newBuilder() + .setMaxAttempts(4) + .setInitialRpcTimeoutDuration(Duration.ofMinutes(10)) + .setMaxRpcTimeoutDuration(Duration.ofMinutes(10)) + .setTotalTimeoutDuration(Duration.ofMinutes(50)) + .build()) + .build(); + settings.stubSettings().build(); + BigtableDataClient clientWithTimeout = BigtableDataClient.create(settings.build()); + + // Initially return a prepare response without delay + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + // Attempt 1 - Trigger plan refresh + service.addExpectation( + ExecuteRpcExpectation.create() + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + // Attempt 2 + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWithStatus(Code.INTERNAL)); + // Attempt 3 + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWithStatus(Code.INTERNAL)); + // Attempt 4 + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWithStatus(Code.INTERNAL)); + // This is triggered by the failure in attempt 4. It succeeds + // but isn't used bc execute stops retrying + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("strCol", stringType()))))); + + PreparedStatement ps = + clientWithTimeout.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = clientWithTimeout.executeQuery(ps.bind().build()); + assertThrows(ApiException.class, rs::next); + // initial success plus 3 refresh failures, plus the refresh triggered by the final failure + assertThat(service.prepareCount).isEqualTo(5); + } + + @Test + public void canRetryAfterRefreshAttemptTimeout() throws IOException { + BigtableDataSettings.Builder settings = defaultSettings(serverRule); + settings + .stubSettings() + .executeQuerySettings() + .setRetrySettings( + RetrySettings.newBuilder() + // First attempt triggers plan refresh retry. + // Second should time out, third should succeed + .setMaxAttempts(3) + .setInitialRpcTimeoutDuration(Duration.ofMillis(10)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(10)) + .setTotalTimeoutDuration(Duration.ofMinutes(50)) + .build()) + .build(); + settings.stubSettings().build(); + BigtableDataClient clientWithTimeout = BigtableDataClient.create(settings.build()); + + // Initially return a prepare response without delay + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + // Attempt 1 - Trigger plan refresh + service.addExpectation( + ExecuteRpcExpectation.create() + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + // Attempt 2 + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + // first refresh attempt times out, but then it succeeds + .withDelay(Duration.ofMillis(15)) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("strCol", stringType()))))); + + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + .respondWith(partialResultSetWithToken(stringValue("s")))); + + PreparedStatement ps = + clientWithTimeout.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = clientWithTimeout.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("s"); + assertThat(rs.next()).isFalse(); + assertThat(service.executeCount).isEqualTo(2); + assertThat(service.prepareCount).isEqualTo(2); + } + + @Test + public void prepareRefreshTimeIsFactoredIntoExecuteAttemptTimeout() throws IOException { + BigtableDataSettings.Builder settings = defaultSettings(serverRule); + settings + .stubSettings() + .executeQuerySettings() + .setRetrySettings( + RetrySettings.newBuilder() + // First attempt triggers plan refresh retry. + // Second should time out, third should succeed + .setMaxAttempts(2) + .setInitialRpcTimeoutDuration(Duration.ofMillis(30)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(30)) + .setTotalTimeoutDuration(Duration.ofMinutes(30)) + .build()) + .build(); + settings.stubSettings().build(); + BigtableDataClient clientWithTimeout = BigtableDataClient.create(settings.build()); + // Initially return a prepare response without delay + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + // Attempt 1 - Trigger plan refresh + service.addExpectation( + ExecuteRpcExpectation.create() + .respondWithException(Code.FAILED_PRECONDITION, planRefreshError())); + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + // Burn most of the execute attempt timeout and succeed + .withDelay(Duration.ofMillis(20)) + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("bar"), + metadata(columnMetadata("strCol", stringType()))))); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("bar")) + // Should timeout bc we used 20 ms on prepare refresh and have 30ms timeout + .withDelay(Duration.ofMillis(20)) + .respondWith(partialResultSetWithToken(stringValue("s")))); + + PreparedStatement ps = + clientWithTimeout.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = clientWithTimeout.executeQuery(ps.bind().build()); + ApiException e = assertThrows(ApiException.class, rs::next); + assertThat(e.getStatusCode().getCode()).isEqualTo(StatusCode.Code.DEADLINE_EXCEEDED); + // initial success plus one refresh + assertThat(service.prepareCount).isEqualTo(2); + // refresh error plus timed out req + assertThat(service.executeCount).isEqualTo(2); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java index 77ec69da9d..9312d3ffe5 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/MetadataErrorHandlingCallableTest.java @@ -18,14 +18,12 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; -import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; import com.google.api.core.SettableApiFuture; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.SqlRow; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; @@ -34,7 +32,6 @@ import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; -import java.util.HashMap; import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import org.junit.Before; @@ -53,14 +50,10 @@ public class MetadataErrorHandlingCallableTest { public void setUp() { metadataFuture = SettableApiFuture.create(); PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse( - metadata( - columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), - new HashMap<>()); + preparedStatement( + metadata(columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))); - callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); + callContext = SqlProtoFactory.callContext(preparedStatement.bind().build(), metadataFuture); outerObserver = new MockResponseObserver<>(true); observer = new MetadataErrorHandlingObserver(outerObserver, callContext); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java index 2ab050c573..9fc1e7a2ea 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java @@ -21,31 +21,50 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.planRefreshError; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import com.google.api.core.ApiClock; import com.google.api.core.SettableApiFuture; +import com.google.api.gax.core.FakeApiClock; +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.FailedPreconditionException; +import com.google.api.gax.rpc.ResponseObserver; import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; -import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryData; import com.google.cloud.bigtable.data.v2.internal.ProtoResultSetMetadata; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata; -import com.google.cloud.bigtable.data.v2.stub.sql.PlanRefreshingCallable.MetadataObserver; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; +import com.google.cloud.bigtable.data.v2.stub.sql.PlanRefreshingCallable.PlanRefreshingObserver; +import com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.FakePreparedStatement; import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockResponseObserver; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockStreamController; +import com.google.protobuf.ByteString; +import io.grpc.Deadline; +import io.grpc.Status.Code; +import java.time.Duration; import java.util.Collections; -import java.util.HashMap; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -63,22 +82,23 @@ public class PlanRefreshingCallableTest { ExecuteQueryCallContext callContext; MockResponseObserver outerObserver; SettableApiFuture metadataFuture; - PlanRefreshingCallable.MetadataObserver observer; + PlanRefreshingObserver observer; + RetrySettings retrySettings; + ApiClock clock; @Before public void setUp() { metadataFuture = SettableApiFuture.create(); PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse( - metadata( - columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), - new HashMap<>()); + preparedStatement( + metadata(columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))); + retrySettings = + EnhancedBigtableStubSettings.newBuilder().executeQuerySettings().retrySettings().build(); + clock = new FakeApiClock(System.nanoTime()); callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); outerObserver = new MockResponseObserver<>(true); - observer = new MetadataObserver(outerObserver, callContext); + observer = new PlanRefreshingObserver(outerObserver, callContext); } @Test @@ -151,12 +171,8 @@ public void testCallable() throws ExecutionException, InterruptedException { MockResponseObserver outerObserver = new MockResponseObserver<>(true); SettableApiFuture metadataFuture = SettableApiFuture.create(); PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse( - metadata( - columnMetadata("foo", stringType()), columnMetadata("bar", int64Type())))), - new HashMap<>()); + preparedStatement( + metadata(columnMetadata("foo", stringType()), columnMetadata("bar", int64Type()))); ExecuteQueryCallContext callContext = ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); @@ -169,4 +185,100 @@ public void testCallable() throws ExecutionException, InterruptedException { assertThat(outerObserver.isDone()).isTrue(); assertThat(outerObserver.getFinalError()).isNull(); } + + @Test + public void testPlanRefreshError() { + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + MockServerStreamingCallable innerCallable = + new MockServerStreamingCallable<>(); + PlanRefreshingCallable planRefreshingCallable = + new PlanRefreshingCallable(innerCallable, requestContext); + MockResponseObserver outerObserver = new MockResponseObserver<>(true); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(new FakePreparedStatement().bind().build(), metadataFuture); + + planRefreshingCallable.call(callContext, outerObserver); + innerCallable.popLastCall().getController().getObserver().onError(planRefreshError()); + ApiException e = (ApiException) outerObserver.getFinalError(); + + assertThat(e.isRetryable()).isTrue(); + assertThat(callContext.resultSetMetadataFuture().isDone()).isFalse(); + ExecuteQueryRequest nextRequest = + callContext.buildRequestWithDeadline( + requestContext, Deadline.after(1, TimeUnit.MILLISECONDS)); + assertThat(nextRequest.getPreparedQuery()).isEqualTo(ByteString.copyFromUtf8("refreshedPlan")); + } + + @Test + public void testPlanRefreshErrorAfterToken() { + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + MockServerStreamingCallable innerCallable = + new MockServerStreamingCallable<>(); + PlanRefreshingCallable planRefreshingCallable = + new PlanRefreshingCallable(innerCallable, requestContext); + MockResponseObserver outerObserver = new MockResponseObserver<>(true); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(new FakePreparedStatement().bind().build(), metadataFuture); + + planRefreshingCallable.call(callContext, outerObserver); + ResponseObserver innerObserver = + innerCallable.popLastCall().getController().getObserver(); + innerObserver.onResponse(partialResultSetWithToken(stringValue("foo"))); + innerObserver.onError(planRefreshError()); + + Throwable t = outerObserver.getFinalError(); + assertThat(t).isInstanceOf(IllegalStateException.class); + } + + @Test + public void testIsPlanRefreshError() { + assertThat(PlanRefreshingCallable.isPlanRefreshError(planRefreshError())).isTrue(); + assertFalse( + PlanRefreshingCallable.isPlanRefreshError( + new FailedPreconditionException( + "A different failed precondition", + null, + GrpcStatusCode.of(Code.FAILED_PRECONDITION), + false))); + } + + @Test + public void planRefreshDelayIsFactoredIntoExecuteTimeout() { + MockServerStreamingCallable innerCallable = + new MockServerStreamingCallable<>(); + RequestContext requestContext = RequestContext.create("project", "instance", "profile"); + PlanRefreshingCallable callable = new PlanRefreshingCallable(innerCallable, requestContext); + MockResponseObserver outerObserver = new MockResponseObserver<>(true); + SettableApiFuture metadataFuture = SettableApiFuture.create(); + SettableApiFuture prepareFuture = SettableApiFuture.create(); + PreparedStatement preparedStatement = + new FakePreparedStatement().withUpdatedPlans(PreparedQueryData.create(prepareFuture), null); + ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + ExecuteQueryCallContext callContext = + ExecuteQueryCallContext.create(preparedStatement.bind().build(), metadataFuture); + + Duration originalAttemptTimeout = Duration.ofMillis(100); + scheduler.schedule( + () -> { + prepareFuture.set( + PrepareResponse.fromProto( + prepareResponse( + ByteString.copyFromUtf8("initialPlan"), + metadata(columnMetadata("strCol", stringType()))))); + }, + 50, + TimeUnit.MILLISECONDS); + ApiCallContext context = + GrpcCallContext.createDefault().withTimeoutDuration(originalAttemptTimeout); + // prepare takes 50 ms to resolve. Despite that the execute timeout should be around 100ms from + // now (w padding) + Deadline paddedDeadlineAtStartOfCall = + Deadline.after(originalAttemptTimeout.toMillis() + 5, TimeUnit.MILLISECONDS); + callable.call(callContext, outerObserver, context); + scheduler.shutdown(); + GrpcCallContext grpcCallContext = + (GrpcCallContext) innerCallable.popLastCall().getApiCallContext(); + Deadline executeDeadline = grpcCallContext.getCallOptions().getDeadline(); + assertThat(executeDeadline.isBefore(paddedDeadlineAtStartOfCall)).isTrue(); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java index 23fc8bbef3..25858cd9f7 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactory.java @@ -15,10 +15,19 @@ */ package com.google.cloud.bigtable.data.v2.stub.sql; +import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ErrorDetails; +import com.google.api.gax.rpc.FailedPreconditionException; import com.google.bigtable.v2.ArrayValue; +import com.google.bigtable.v2.BigtableGrpc; import com.google.bigtable.v2.ColumnMetadata; +import com.google.bigtable.v2.ExecuteQueryRequest; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.PartialResultSet; +import com.google.bigtable.v2.PrepareQueryRequest; import com.google.bigtable.v2.PrepareQueryResponse; import com.google.bigtable.v2.ProtoRows; import com.google.bigtable.v2.ProtoRowsBatch; @@ -27,34 +36,121 @@ import com.google.bigtable.v2.Type; import com.google.bigtable.v2.Type.Struct.Field; import com.google.bigtable.v2.Value; +import com.google.cloud.bigtable.data.v2.internal.NameUtil; +import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryData; +import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl.PreparedQueryVersion; +import com.google.cloud.bigtable.data.v2.internal.QueryParamUtil; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.hash.HashFunction; import com.google.common.hash.Hashing; +import com.google.common.truth.Truth; +import com.google.protobuf.Any; import com.google.protobuf.ByteString; import com.google.protobuf.Timestamp; +import com.google.rpc.PreconditionFailure; +import com.google.rpc.PreconditionFailure.Violation; import com.google.type.Date; +import io.grpc.Metadata; +import io.grpc.Status; +import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.StreamObserver; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; import java.util.Arrays; - +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.LinkedBlockingDeque; +import javax.annotation.Nullable; + +// TODO rename this to SqlApiTestUtils /** Utilities for creating sql proto objects in tests */ public class SqlProtoFactory { private static final HashFunction CRC32C = Hashing.crc32c(); + private static final Metadata.Key ERROR_DETAILS_KEY = + Metadata.Key.of("grpc-status-details-bin", Metadata.BINARY_BYTE_MARSHALLER); private SqlProtoFactory() {} + public static ApiException planRefreshError() { + Metadata trailers = new Metadata(); + PreconditionFailure failure = + PreconditionFailure.newBuilder() + .addViolations(Violation.newBuilder().setType("PREPARED_QUERY_EXPIRED").build()) + .build(); + ErrorDetails refreshErrorDetails = + ErrorDetails.builder().setRawErrorMessages(ImmutableList.of(Any.pack(failure))).build(); + byte[] status = + com.google.rpc.Status.newBuilder().addDetails(Any.pack(failure)).build().toByteArray(); + // This needs to be in trailers in order to round trip + trailers.put(ERROR_DETAILS_KEY, status); + + // This is not initially retryable, the PlanRefreshingCallable overrides this. + return new FailedPreconditionException( + new StatusRuntimeException(Status.FAILED_PRECONDITION, trailers), + GrpcStatusCode.of(Code.FAILED_PRECONDITION), + false, + refreshErrorDetails); + } + public static PrepareQueryResponse prepareResponse( - ByteString preparedQuery, ResultSetMetadata metadata) { + ByteString preparedQuery, ResultSetMetadata metadata, Instant validUntil) { return PrepareQueryResponse.newBuilder() .setPreparedQuery(preparedQuery) - .setValidUntil(Timestamp.newBuilder().setSeconds(1000).setNanos(1000).build()) + // set validUntil a year in the future so these plans never expire in test runs + .setValidUntil( + Timestamp.newBuilder() + .setSeconds(validUntil.getEpochSecond()) + .setNanos(validUntil.getNano()) + .build()) .setMetadata(metadata) .build(); } + public static PrepareQueryResponse prepareResponse( + ByteString preparedQuery, ResultSetMetadata metadata) { + return prepareResponse(preparedQuery, metadata, Instant.now().plus(Duration.ofDays(365))); + } + public static PrepareQueryResponse prepareResponse(ResultSetMetadata metadata) { return prepareResponse(ByteString.copyFromUtf8("foo"), metadata); } + public static PreparedStatementImpl preparedStatement(ResultSetMetadata metadata) { + return preparedStatement(metadata, new HashMap<>()); + } + + public static PreparedStatementImpl preparedStatement( + ResultSetMetadata metadata, Map> paramTypes) { + // We never expire the test prepare response so it's safe to null the stub and request + return preparedStatement(PrepareResponse.fromProto(prepareResponse(metadata)), paramTypes); + } + + public static PreparedStatementImpl preparedStatement( + PrepareResponse response, Map> paramTypes) { + return new FakePreparedStatement(response, paramTypes); + } + + public static ExecuteQueryCallContext callContext(BoundStatement boundStatement) { + return callContext(boundStatement, SettableApiFuture.create()); + } + + public static ExecuteQueryCallContext callContext( + BoundStatement boundStatement, + SettableApiFuture mdFuture) { + return ExecuteQueryCallContext.create(boundStatement, mdFuture); + } + public static ColumnMetadata columnMetadata(String name, Type type) { return ColumnMetadata.newBuilder().setName(name).setType(type).build(); } @@ -257,4 +353,267 @@ public static ResultSetMetadata metadata(ColumnMetadata... columnMetadata) { public static int checksum(ByteString bytes) { return CRC32C.hashBytes(bytes.toByteArray()).asInt(); } + + /** Used to test ExecuteQuery and PrepareQuery APIs using the RpcExpectations below */ + public static class TestBigtableSqlService extends BigtableGrpc.BigtableImplBase { + public static final String DEFAULT_PROJECT_ID = "fake-project"; + public static final String DEFAULT_INSTANCE_ID = "fake-instance"; + public static final String DEFAULT_APP_PROFILE_ID = "fake-app-profile"; + public static final ByteString DEFAULT_PREPARED_QUERY = ByteString.copyFromUtf8("foo"); + Queue executeExpectations = new LinkedBlockingDeque<>(); + Queue prepareExpectations = new LinkedBlockingDeque<>(); + int executeCount = 0; + public int prepareCount = 0; + + public void addExpectation(ExecuteRpcExpectation expectation) { + executeExpectations.add(expectation); + } + + public void addExpectation(PrepareRpcExpectation expectation) { + prepareExpectations.add(expectation); + } + + @Override + public void executeQuery( + ExecuteQueryRequest request, StreamObserver responseObserver) { + ExecuteRpcExpectation expectedRpc = executeExpectations.poll(); + executeCount++; + int requestIndex = executeCount - 1; + + Truth.assertWithMessage("Unexpected request#" + requestIndex + ":" + request.toString()) + .that(expectedRpc) + .isNotNull(); + Truth.assertWithMessage("Unexpected request#" + requestIndex) + .that(request) + .isEqualTo(expectedRpc.getExpectedRequest()); + + try { + Thread.sleep(expectedRpc.delay.toMillis()); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + for (ExecuteQueryResponse response : expectedRpc.responses) { + responseObserver.onNext(response); + } + if (expectedRpc.statusCode.toStatus().isOk()) { + responseObserver.onCompleted(); + } else if (expectedRpc.exception != null) { + responseObserver.onError(expectedRpc.exception); + } else { + responseObserver.onError(expectedRpc.statusCode.toStatus().asRuntimeException()); + } + } + + @Override + public void prepareQuery( + PrepareQueryRequest request, StreamObserver responseObserver) { + PrepareRpcExpectation expectedRpc = prepareExpectations.poll(); + prepareCount++; + int requestIndex = prepareCount - 1; + + Truth.assertWithMessage("Unexpected request#" + requestIndex + ":" + request.toString()) + .that(expectedRpc) + .isNotNull(); + Truth.assertWithMessage("Unexpected request#" + requestIndex) + .that(request) + .isEqualTo(expectedRpc.getExpectedRequest()); + + try { + Thread.sleep(expectedRpc.delay.toMillis()); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + if (expectedRpc.statusCode == Code.OK) { + responseObserver.onNext(expectedRpc.response); + responseObserver.onCompleted(); + } else { + responseObserver.onError(expectedRpc.statusCode.toStatus().asRuntimeException()); + } + } + } + + public static class ExecuteRpcExpectation { + ExecuteQueryRequest.Builder request; + Status.Code statusCode; + @Nullable ApiException exception; + List responses; + Duration delay; + + private ExecuteRpcExpectation() { + this.request = ExecuteQueryRequest.newBuilder(); + this.request.setPreparedQuery(TestBigtableSqlService.DEFAULT_PREPARED_QUERY); + this.request.setInstanceName( + NameUtil.formatInstanceName( + TestBigtableSqlService.DEFAULT_PROJECT_ID, + TestBigtableSqlService.DEFAULT_INSTANCE_ID)); + this.request.setAppProfileId(TestBigtableSqlService.DEFAULT_APP_PROFILE_ID); + this.statusCode = Code.OK; + this.responses = new ArrayList<>(); + this.delay = Duration.ZERO; + } + + public static ExecuteRpcExpectation create() { + return new ExecuteRpcExpectation(); + } + + public ExecuteRpcExpectation withResumeToken(ByteString resumeToken) { + this.request.setResumeToken(resumeToken); + return this; + } + + public ExecuteRpcExpectation withDelay(Duration delay) { + this.delay = delay; + return this; + } + + public ExecuteRpcExpectation withParams(Map params) { + this.request.putAllParams(params); + return this; + } + + public ExecuteRpcExpectation withPreparedQuery(ByteString preparedQuery) { + this.request.setPreparedQuery(preparedQuery); + return this; + } + + public ExecuteRpcExpectation respondWithStatus(Status.Code code) { + this.statusCode = code; + return this; + } + + public ExecuteRpcExpectation respondWithException(Status.Code code, ApiException exception) { + this.statusCode = code; + this.exception = exception; + return this; + } + + public ExecuteRpcExpectation respondWith(ExecuteQueryResponse... responses) { + this.responses = Arrays.asList(responses); + return this; + } + + ExecuteQueryRequest getExpectedRequest() { + return this.request.build(); + } + } + + public static class PrepareRpcExpectation { + PrepareQueryRequest.Builder request; + Status.Code statusCode; + PrepareQueryResponse response; + Duration delay; + + private PrepareRpcExpectation() { + this.request = PrepareQueryRequest.newBuilder(); + this.request.setInstanceName( + NameUtil.formatInstanceName( + TestBigtableSqlService.DEFAULT_PROJECT_ID, + TestBigtableSqlService.DEFAULT_INSTANCE_ID)); + this.request.setAppProfileId(TestBigtableSqlService.DEFAULT_APP_PROFILE_ID); + this.statusCode = Code.OK; + this.delay = Duration.ZERO; + } + + public static PrepareRpcExpectation create() { + return new PrepareRpcExpectation(); + } + + public PrepareRpcExpectation withSql(String sqlQuery) { + this.request.setQuery(sqlQuery); + return this; + } + + public PrepareRpcExpectation withParamTypes(Map> paramTypes) { + Map protoParamTypes = new HashMap<>(); + for (Map.Entry> entry : paramTypes.entrySet()) { + Type proto = QueryParamUtil.convertToQueryParamProto(entry.getValue()); + protoParamTypes.put(entry.getKey(), proto); + } + this.request.putAllParamTypes(protoParamTypes); + return this; + } + + public PrepareRpcExpectation respondWithStatus(Status.Code code) { + this.statusCode = code; + return this; + } + + public PrepareRpcExpectation respondWith(PrepareQueryResponse res) { + this.response = res; + return this; + } + + public PrepareRpcExpectation withDelay(Duration delay) { + this.delay = delay; + return this; + } + + PrepareQueryRequest getExpectedRequest() { + return this.request.build(); + } + } + + /** + * Fake prepared statement for testing. Note that the schema changes on calls to hard refresh. + * This is used to test plan updates propagate. + */ + public static final class FakePreparedStatement extends PreparedStatementImpl { + private static final PrepareResponse DEFAULT_INITIAL_RESPONSE = + PrepareResponse.fromProto( + prepareResponse( + ByteString.copyFromUtf8("initialPlan"), + metadata(columnMetadata("strCol", stringType())))); + private static final PreparedQueryData DEFAULT_INITIAL_PLAN = + PreparedQueryData.create(ApiFutures.immediateFuture(DEFAULT_INITIAL_RESPONSE)); + private static final PreparedQueryData DEFAULT_PLAN_ON_REFRESH = + PreparedQueryData.create( + ApiFutures.immediateFuture( + PrepareResponse.fromProto( + prepareResponse( + ByteString.copyFromUtf8("refreshedPlan"), + metadata(columnMetadata("bytesColl", bytesType())))))); + + private PreparedQueryData initialPlan; + private PreparedQueryData planOnRefresh; + private Map> paramTypes; + + public FakePreparedStatement() { + super(DEFAULT_INITIAL_RESPONSE, new HashMap<>(), null, null); + this.initialPlan = DEFAULT_INITIAL_PLAN; + this.planOnRefresh = DEFAULT_PLAN_ON_REFRESH; + this.paramTypes = new HashMap<>(); + } + + public FakePreparedStatement( + PrepareResponse prepareResponse, Map> paramTypes) { + super(prepareResponse, paramTypes, null, null); + this.initialPlan = PreparedQueryData.create(ApiFutures.immediateFuture(prepareResponse)); + // Don't expect an refresh using this configuration + this.planOnRefresh = null; + this.paramTypes = paramTypes; + } + + FakePreparedStatement withUpdatedPlans( + PreparedQueryData initialPlan, PreparedQueryData planOnRefresh) { + this.initialPlan = initialPlan; + this.planOnRefresh = planOnRefresh; + return this; + } + + @Override + public PreparedQueryData getLatestPrepareResponse() { + Preconditions.checkState( + initialPlan != null, "Trying to refresh FakePreparedStatement without planOnRefresh set"); + return initialPlan; + } + + @Override + public PreparedQueryData markExpiredAndStartRefresh( + PreparedQueryVersion expiredPreparedQueryVersion) { + return planOnRefresh; + } + + @Override + public void assertUsingSameStub(EnhancedBigtableStub stub) {} + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactoryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactoryTest.java index 2a3eb9e404..cb2c068939 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactoryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlProtoFactoryTest.java @@ -18,6 +18,9 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.common.truth.Truth.assertThat; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ErrorDetails; +import com.google.api.gax.rpc.StatusCode.Code; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.PartialResultSet; import com.google.bigtable.v2.ProtoRows; @@ -44,4 +47,16 @@ public void serializedProtoRows_canRoundTrip() throws InvalidProtocolBufferExcep assertThat(protoRows.getValuesList().get(1).getBytesValue()) .isEqualTo(ByteString.copyFromUtf8("bytes")); } + + @Test + public void testPlanRefreshError() { + ApiException planRefreshError = SqlProtoFactory.planRefreshError(); + assertThat(planRefreshError.getStatusCode().getCode()).isEqualTo(Code.FAILED_PRECONDITION); + ErrorDetails details = planRefreshError.getErrorDetails(); + assertThat(details.getPreconditionFailure()).isNotNull(); + assertThat(details.getPreconditionFailure().getViolationsList()).isNotEmpty(); + assertThat(details.getPreconditionFailure().getViolationsList().get(0).getType()) + .isEqualTo("PREPARED_QUERY_EXPIRED"); + assertThat(PlanRefreshingCallable.isPlanRefreshError(planRefreshError)).isTrue(); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java index 65cb680599..fd6f0e2302 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMergingCallableTest.java @@ -17,13 +17,14 @@ import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.arrayType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.arrayValue; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.callContext; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.columnMetadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Type; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.int64Value; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.metadata; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithToken; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.partialResultSetWithoutToken; -import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.prepareResponse; +import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.preparedStatement; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringType; import static com.google.cloud.bigtable.data.v2.stub.sql.SqlProtoFactory.stringValue; import static com.google.common.truth.Truth.assertThat; @@ -32,7 +33,6 @@ import com.google.api.core.SettableApiFuture; import com.google.api.gax.rpc.ServerStream; import com.google.bigtable.v2.ExecuteQueryResponse; -import com.google.cloud.bigtable.data.v2.internal.PrepareResponse; import com.google.cloud.bigtable.data.v2.internal.PreparedStatementImpl; import com.google.cloud.bigtable.data.v2.internal.ProtoSqlRow; import com.google.cloud.bigtable.data.v2.internal.SqlRow; @@ -42,8 +42,8 @@ import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable; import com.google.common.collect.Lists; import java.util.Arrays; -import java.util.HashMap; import java.util.List; +import java.util.concurrent.ExecutionException; import java.util.stream.Collectors; import org.junit.Test; import org.junit.runner.RunWith; @@ -58,7 +58,7 @@ public class SqlRowMergingCallableTest { @Test - public void testMerging() { + public void testMerging() throws ExecutionException, InterruptedException { ServerStreamingStashCallable inner = new ServerStreamingStashCallable<>( Lists.newArrayList( @@ -68,22 +68,19 @@ public void testMerging() { arrayValue(stringValue("foo"), stringValue("bar"))), partialResultSetWithToken(stringValue("test"), int64Value(10), arrayValue()))); - PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse( - metadata( - columnMetadata("stringCol", stringType()), - columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType()))))), - new HashMap<>()); + PreparedStatementImpl preparedStatement = + preparedStatement( + metadata( + columnMetadata("stringCol", stringType()), + columnMetadata("intCol", int64Type()), + columnMetadata("arrayCol", arrayType(stringType())))); BoundStatement boundStatement = preparedStatement.bind().build(); - ResultSetMetadata metadata = preparedStatement.getPrepareResponse().resultSetMetadata(); + ResultSetMetadata metadata = + preparedStatement.getLatestPrepareResponse().prepareFuture().get().resultSetMetadata(); SettableApiFuture mdFuture = SettableApiFuture.create(); mdFuture.set(metadata); SqlRowMergingCallable rowMergingCallable = new SqlRowMergingCallable(inner); - ServerStream results = - rowMergingCallable.call(ExecuteQueryCallContext.create(boundStatement, mdFuture)); + ServerStream results = rowMergingCallable.call(callContext(boundStatement, mdFuture)); List resultsList = results.stream().collect(Collectors.toList()); assertThat(resultsList) .containsExactly( @@ -98,16 +95,13 @@ public void testMerging() { } @Test - public void testError() { - PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse( - metadata( - columnMetadata("stringCol", stringType()), - columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType()))))), - new HashMap<>()); + public void testError() throws ExecutionException, InterruptedException { + PreparedStatementImpl preparedStatement = + preparedStatement( + metadata( + columnMetadata("stringCol", stringType()), + columnMetadata("intCol", int64Type()), + columnMetadata("arrayCol", arrayType(stringType())))); BoundStatement boundStatement = preparedStatement.bind().build(); // empty response is invalid @@ -117,9 +111,9 @@ public void testError() { SqlRowMergingCallable rowMergingCallable = new SqlRowMergingCallable(inner); SettableApiFuture mdFuture = SettableApiFuture.create(); - mdFuture.set(preparedStatement.getPrepareResponse().resultSetMetadata()); - ServerStream results = - rowMergingCallable.call(ExecuteQueryCallContext.create(boundStatement, mdFuture)); + mdFuture.set( + preparedStatement.getLatestPrepareResponse().prepareFuture().get().resultSetMetadata()); + ServerStream results = rowMergingCallable.call(callContext(boundStatement)); assertThrows(IllegalStateException.class, () -> results.iterator().next()); } @@ -127,14 +121,11 @@ public void testError() { @Test public void testMetdataFutureError() { PreparedStatement preparedStatement = - PreparedStatementImpl.create( - PrepareResponse.fromProto( - prepareResponse( - metadata( - columnMetadata("stringCol", stringType()), - columnMetadata("intCol", int64Type()), - columnMetadata("arrayCol", arrayType(stringType()))))), - new HashMap<>()); + preparedStatement( + metadata( + columnMetadata("stringCol", stringType()), + columnMetadata("intCol", int64Type()), + columnMetadata("arrayCol", arrayType(stringType())))); BoundStatement boundStatement = preparedStatement.bind().build(); // empty response is invalid @@ -145,8 +136,7 @@ public void testMetdataFutureError() { SqlRowMergingCallable rowMergingCallable = new SqlRowMergingCallable(inner); SettableApiFuture mdFuture = SettableApiFuture.create(); mdFuture.setException(new RuntimeException("test")); - ServerStream results = - rowMergingCallable.call(ExecuteQueryCallContext.create(boundStatement, mdFuture)); + ServerStream results = rowMergingCallable.call(callContext(boundStatement, mdFuture)); assertThrows(RuntimeException.class, () -> results.iterator().next()); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/testing/MockStreamingApi.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/testing/MockStreamingApi.java index 4ecca917ed..f82f1fed45 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/testing/MockStreamingApi.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/testing/MockStreamingApi.java @@ -25,6 +25,7 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import javax.annotation.Nullable; public class MockStreamingApi { public static class MockServerStreamingCallable @@ -36,7 +37,7 @@ public static class MockServerStreamingCallable public void call( RequestT request, ResponseObserver responseObserver, ApiCallContext context) { MockStreamController controller = new MockStreamController<>(responseObserver); - calls.add(new MockServerStreamingCall<>(request, controller)); + calls.add(new MockServerStreamingCall<>(request, controller, context)); responseObserver.onStart(controller); } @@ -52,10 +53,15 @@ public MockServerStreamingCall popLastCall() { public static class MockServerStreamingCall { private final RequestT request; private final MockStreamController controller; + private final ApiCallContext apiCallContext; - public MockServerStreamingCall(RequestT request, MockStreamController controller) { + public MockServerStreamingCall( + RequestT request, + MockStreamController controller, + @Nullable ApiCallContext apiCallContext) { this.request = request; this.controller = controller; + this.apiCallContext = apiCallContext; } public RequestT getRequest() { @@ -65,6 +71,10 @@ public RequestT getRequest() { public MockStreamController getController() { return controller; } + + public ApiCallContext getApiCallContext() { + return apiCallContext; + } } public static class MockStreamController implements StreamController { diff --git a/test-proxy/pom.xml b/test-proxy/pom.xml index a0d8ac7290..c68092bb56 100644 --- a/test-proxy/pom.xml +++ b/test-proxy/pom.xml @@ -71,9 +71,9 @@ protobuf-maven-plugin 0.6.1 - com.google.protobuf:protoc:3.22.3:exe:${os.detected.classifier} + com.google.protobuf:protoc:3.22.3:exe:osx-x86_64 grpc-java - io.grpc:protoc-gen-grpc-java:1.24.0:exe:${os.detected.classifier} + io.grpc:protoc-gen-grpc-java:1.24.0:exe:osx-x86_64 From c16aa732315a2c52ed6abd8698bb134d2bb745b6 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Mon, 17 Mar 2025 17:24:42 -0400 Subject: [PATCH 08/11] Add ConvertExceptionCallable to ExecuteQuery call chain Change-Id: I22cdeb1560a0975c8f2ff380ca30f3e09070417d --- .../data/v2/stub/EnhancedBigtableStub.java | 9 +- .../v2/stub/sql/ExecuteQueryRetryTest.java | 99 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index db75050a4f..160f0b6b2c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -1175,6 +1175,13 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { ServerStreamingCallable withPlanRefresh = new PlanRefreshingCallable(withStatsHeaders, requestContext); + // Sometimes ExecuteQuery connections are disconnected via an RST frame. This error is transient + // and should be treated similar to UNAVAILABLE. However, this exception has an INTERNAL error + // code which by default is not retryable. Convert the exception, so it can be retried in the + // client. + ServerStreamingCallable convertException = + new ConvertExceptionCallable<>(withPlanRefresh); + ServerStreamingCallSettings retrySettings = ServerStreamingCallSettings.newBuilder() .setResumptionStrategy(new ExecuteQueryResumptionStrategy()) @@ -1189,7 +1196,7 @@ public Map extract(ExecuteQueryRequest executeQueryRequest) { // attempt stream will have reset set to true, so any unyielded data from the previous // attempt will be reset properly ServerStreamingCallable retries = - withRetries(withPlanRefresh, retrySettings); + withRetries(convertException, retrySettings); ServerStreamingCallable merging = new SqlRowMergingCallable(retries); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java index 734fbe14a1..e2297911b4 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java @@ -31,11 +31,13 @@ import static org.junit.Assert.assertThrows; import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.GrpcStatusCode; import com.google.api.gax.grpc.GrpcTransportChannel; import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.FailedPreconditionException; import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.api.gax.rpc.InternalException; import com.google.api.gax.rpc.StatusCode; import com.google.bigtable.v2.ExecuteQueryResponse; import com.google.bigtable.v2.ResultSetMetadata; @@ -54,7 +56,9 @@ import com.google.cloud.bigtable.gaxx.reframing.IncompleteStreamException; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; +import io.grpc.Status; import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; import io.grpc.testing.GrpcServerRule; import java.io.IOException; import java.lang.ref.WeakReference; @@ -797,4 +801,99 @@ public void prepareRefreshTimeIsFactoredIntoExecuteAttemptTimeout() throws IOExc // refresh error plus timed out req assertThat(service.executeCount).isEqualTo(2); } + + @Test + public void retriesRstStreamError() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + ApiException rstStreamException = + new InternalException( + new StatusRuntimeException( + Status.INTERNAL.withDescription( + "INTERNAL: HTTP/2 error code: INTERNAL_ERROR\nReceived Rst Stream")), + GrpcStatusCode.of(Status.Code.INTERNAL), + false); + service.addExpectation( + ExecuteRpcExpectation.create().respondWithException(Code.INTERNAL, rstStreamException)); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("foo")) + .respondWith(partialResultSetWithToken(stringValue("s")))); + + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("s"); + assertThat(rs.next()).isFalse(); + assertThat(service.executeCount).isEqualTo(2); + assertThat(service.prepareCount).isEqualTo(1); + } + + @Test + public void retriesRetriableAuthException() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + ApiException authException = + new InternalException( + new StatusRuntimeException( + Status.INTERNAL.withDescription( + "Authentication backend internal server error. Please retry")), + GrpcStatusCode.of(Status.Code.INTERNAL), + false); + service.addExpectation( + ExecuteRpcExpectation.create().respondWithException(Code.INTERNAL, authException)); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("foo")) + .respondWith(partialResultSetWithToken(stringValue("s")))); + + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("s"); + assertThat(rs.next()).isFalse(); + assertThat(service.executeCount).isEqualTo(2); + assertThat(service.prepareCount).isEqualTo(1); + } + + @Test + public void retriesGoAwayException() { + service.addExpectation( + PrepareRpcExpectation.create() + .withSql("SELECT * FROM table") + .respondWith( + prepareResponse( + ByteString.copyFromUtf8("foo"), + metadata(columnMetadata("strCol", stringType()))))); + ApiException authException = + new InternalException( + new StatusRuntimeException( + Status.INTERNAL.withDescription("Stream closed before write could take place")), + GrpcStatusCode.of(Status.Code.INTERNAL), + false); + service.addExpectation( + ExecuteRpcExpectation.create().respondWithException(Code.INTERNAL, authException)); + service.addExpectation( + ExecuteRpcExpectation.create() + .withPreparedQuery(ByteString.copyFromUtf8("foo")) + .respondWith(partialResultSetWithToken(stringValue("s")))); + + PreparedStatement ps = client.prepareStatement("SELECT * FROM table", new HashMap<>()); + ResultSet rs = client.executeQuery(ps.bind().build()); + assertThat(rs.next()).isTrue(); + assertThat(rs.getString("strCol")).isEqualTo("s"); + assertThat(rs.next()).isFalse(); + assertThat(service.executeCount).isEqualTo(2); + assertThat(service.prepareCount).isEqualTo(1); + } } From 54a049a0c2691de16ee089e3c4042701ee5db12b Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Fri, 28 Feb 2025 17:55:35 -0500 Subject: [PATCH 09/11] Update test proxy for prepare, remove BetaApi annotations, add clirr exclusions for changed beta api This also temporarily disables the ExecuteQuery conformance tests so they can be rewritten to use Prepare. Updated local copies of the tests are passing Also fixes BoundStatement docs I missed in previous PR Change-Id: I9024d50fff3c076064b8277fd4a5aa426e0b5de4 --- .../clirr-ignored-differences.xml | 90 ++++++++++++++++--- .../bigtable/data/v2/BigtableDataClient.java | 10 ++- .../data/v2/models/sql/BoundStatement.java | 24 ++--- .../data/v2/models/sql/ColumnMetadata.java | 2 - .../data/v2/models/sql/PreparedStatement.java | 2 - .../data/v2/models/sql/ResultSet.java | 2 - .../data/v2/models/sql/ResultSetMetadata.java | 2 - .../bigtable/data/v2/models/sql/SqlType.java | 2 - .../bigtable/data/v2/models/sql/Struct.java | 2 - .../data/v2/models/sql/StructReader.java | 2 - test-proxy/known_failures.txt | 1 + ...r.java => BoundStatementDeserializer.java} | 48 +++++----- .../bigtable/testproxy/CbtTestProxy.java | 26 +++++- 13 files changed, 139 insertions(+), 74 deletions(-) rename test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/{StatementDeserializer.java => BoundStatementDeserializer.java} (77%) diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml index 303ecc3a06..a9734b96d4 100644 --- a/google-cloud-bigtable/clirr-ignored-differences.xml +++ b/google-cloud-bigtable/clirr-ignored-differences.xml @@ -282,20 +282,6 @@ *getTimestamp(*) java.time.Instant - - - 7006 - com/google/cloud/bigtable/data/v2/models/sql/StructReader - *getTimestamp(*) - java.time.Instant - - - - 7005 - com/google/cloud/bigtable/data/v2/models/sql/Statement$Builder - *setTimestampParam(java.lang.String, org.threeten.bp.Instant) - *setTimestampParam(java.lang.String, java.time.Instant) - 7013 @@ -320,4 +306,80 @@ com/google/cloud/bigtable/data/v2/stub/metrics/BigtableCloudMonitoringExporter * + + + 7005 + com/google/cloud/bigtable/data/v2/BigtableDataClient + *executeQuery* + * + + + + 8001 + com/google/cloud/bigtable/data/v2/models/sql/Statement + * + + + + 8001 + com/google/cloud/bigtable/data/v2/models/sql/Statement$Builder + * + + + + 8001 + com/google/cloud/bigtable/data/v2/models/sql/Statement$Builder + * + + + + 7004 + com/google/cloud/bigtable/data/v2/internal/SqlRowMergerUtil + * + + + + 7004 + com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext + *ExecuteQueryCallContext* + + + + 7009 + com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext + *ExecuteQueryCallContext* + + + + 7005 + com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallContext + *create* + * + + + + 7004 + com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable + *ExecuteQueryCallable* + * + + + + 7005 + com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallable + *call* + * + + + + 8001 + com/google/cloud/bigtable/data/v2/stub/sql/MetadataResolvingCallable + + + + 7004 + com/google/cloud/bigtable/data/v2/stub/sql/SqlRowMerger + * + * + diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 7f16645060..889d36e383 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -2737,7 +2737,6 @@ public void readChangeStreamAsync( * * @see {@link PreparedStatement} & {@link BoundStatement} for query options. */ - @BetaApi public ResultSet executeQuery(BoundStatement boundStatement) { boundStatement.assertUsingSameStub(stub); SqlServerStream stream = stub.createExecuteQueryCallable().call(boundStatement); @@ -2748,12 +2747,19 @@ public ResultSet executeQuery(BoundStatement boundStatement) { * Prepares a query for execution. If possible this should be called once and reused across * requests. This will amortize the cost of query preparation. * + *

      A parameterized query should contain placeholders in the form of {@literal @} followed by + * the parameter name. Parameter names may consist of any combination of letters, numbers, and + * underscores. + * + *

      Parameters can appear anywhere that a literal value is expected. The same parameter name can + * be used more than once, for example: {@code WHERE cf["qualifier1"] = @value OR cf["qualifier2"] + * = @value } + * * @param query sql query string to prepare * @param paramTypes a Map of the parameter names and the corresponding {@link SqlType} for all * query parameters in 'query' * @return {@link PreparedStatement} which is used to create {@link BoundStatement}s to execute */ - @BetaApi public PreparedStatement prepareStatement(String query, Map> paramTypes) { PrepareQueryRequest request = PrepareQueryRequest.create(query, paramTypes); PrepareResponse response = stub.prepareQueryCallable().call(request); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java index 3165edb23c..82c1084afd 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/BoundStatement.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.bigtable.v2.ArrayValue; import com.google.bigtable.v2.ExecuteQueryRequest; @@ -39,33 +38,22 @@ import java.util.Map; import javax.annotation.Nullable; -// TODO update doc /** - * A SQL statement that can be executed by calling {@link + * A bound SQL statement that can be executed by calling {@link * com.google.cloud.bigtable.data.v2.BigtableDataClient#executeQuery(BoundStatement)}. * - *

      A statement contains a SQL string and optional parameters. A parameterized query should - * contain placeholders in the form of {@literal @} followed by the parameter name. Parameter names - * may consist of any combination of letters, numbers, and underscores. + *

      It is an error to bind a statement with unset parameters. * - *

      Parameters can appear anywhere that a literal value is expected. The same parameter name can - * be used more than once, for example: {@code WHERE cf["qualifier1"] = @value OR cf["qualifier2"] - * = @value } - * - *

      It is an error to execute an SQL query with placeholders for unset parameters. - * - *

      Parameterized Statements are constructed using a {@link Builder} and calling - * setTypeParam(String paramName, Type value) for the appropriate type. For example: + *

      BoundStatements are constructed using a {@link Builder} and calling setTypeParam(String + * paramName, Type value) for the appropriate type. For example: * * 

      {@code
- * Statement statement = Statement
- *     .newBuilder("SELECT cf[@qualifer] FROM table WHERE _key=@key")
+ * BoundStatementt boundStatement = preparedStatement.bind()
  *     .setBytesParam("qualifier", ByteString.copyFromUtf8("test"))
  *     .setBytesParam("key", ByteString.copyFromUtf8("testKey"))
  *     .build();
  * }
      */ -@BetaApi public class BoundStatement { private final PreparedStatementImpl preparedStatement; @@ -104,7 +92,7 @@ public Builder(PreparedStatementImpl preparedStatement, Map> this.params = new HashMap<>(); } - /** Builds a {@code Statement} from the builder */ + /** Builds a {@link BoundStatement} from the builder */ public BoundStatement build() { for (Map.Entry> paramType : paramTypes.entrySet()) { String paramName = paramType.getKey(); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java index 0a722a914d..e27ca6ea39 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java @@ -15,10 +15,8 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; /** Represents the metadata for a column in a {@link ResultSet} */ -@BetaApi public interface ColumnMetadata { /** The name of the column. Returns Empty string if the column has no name */ String name(); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java index 46c0ec59b7..a45c513284 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; /** * The results of query preparation that can be used to create {@link BoundStatement}s to execute @@ -24,7 +23,6 @@ *

      Whenever possible this should be shared across different instances of the same query, in order * to amortize query preparation costs. */ -@BetaApi public interface PreparedStatement { /** diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java index 807e995712..cf9e26d421 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; /** * A set of SQL data, generated as the result of an ExecuteQuery request. @@ -38,7 +37,6 @@ *

      {@code ResultSet} implementations are not required to be thread-safe: the thread that asked * for a ResultSet must be the one that interacts with it. */ -@BetaApi public interface ResultSet extends StructReader, AutoCloseable { /** diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSetMetadata.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSetMetadata.java index 23e7155e67..3ebabf9d03 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSetMetadata.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSetMetadata.java @@ -15,11 +15,9 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; import java.util.List; /** Provides information about the schema of a {@link ResultSet}. */ -@BetaApi public interface ResultSetMetadata { /** @return full list of {@link ColumnMetadata} for each column in the {@link ResultSet}. */ diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/SqlType.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/SqlType.java index d4d3261dcf..7191ec9f3c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/SqlType.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/SqlType.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.cloud.Date; import com.google.cloud.bigtable.common.Type; @@ -34,7 +33,6 @@ * * @param the corresponding java type */ -@BetaApi public interface SqlType extends Serializable { /* Enumeration of the types */ diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Struct.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Struct.java index 23b113f9f7..a043e714f0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Struct.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/Struct.java @@ -15,12 +15,10 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; import java.io.Serializable; /** * The representation of a SQL Struct type. Data can be accessed using the methods from the {@code * StructReader} interface. */ -@BetaApi public interface Struct extends StructReader, Serializable {} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/StructReader.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/StructReader.java index f127b6b54c..76ecfb1ef9 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/StructReader.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/StructReader.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; -import com.google.api.core.BetaApi; import com.google.cloud.Date; import com.google.protobuf.ByteString; import java.time.Instant; @@ -36,7 +35,6 @@ * a {@code NullPointerException}; {@link #isNull(int)} & {@link #isNull(String)} can be used to * check for null values. */ -@BetaApi public interface StructReader { /** * @param columnIndex index of the column diff --git a/test-proxy/known_failures.txt b/test-proxy/known_failures.txt index c33137957e..ac6a05ed90 100644 --- a/test-proxy/known_failures.txt +++ b/test-proxy/known_failures.txt @@ -1 +1,2 @@ TestFeatureGap/(traffic_director_enabled|direct_access_requested) +TestExecuteQuery diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/StatementDeserializer.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/BoundStatementDeserializer.java similarity index 77% rename from test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/StatementDeserializer.java rename to test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/BoundStatementDeserializer.java index 4eb5f47e3a..43da147274 100644 --- a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/StatementDeserializer.java +++ b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/BoundStatementDeserializer.java @@ -18,92 +18,94 @@ import com.google.bigtable.v2.Value; import com.google.bigtable.v2.Value.KindCase; import com.google.cloud.Date; +import com.google.cloud.bigtable.data.v2.models.sql.BoundStatement; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.protobuf.Timestamp; import java.time.Instant; import java.util.ArrayList; import java.util.List; import java.util.Map; -public class StatementDeserializer { +public class BoundStatementDeserializer { - static Statement toStatement(ExecuteQueryRequest request) { - Statement.Builder statementBuilder = Statement.newBuilder(request.getRequest().getQuery()); + static BoundStatement toBoundStatement( + PreparedStatement preparedStatement, ExecuteQueryRequest request) { + BoundStatement.Builder boundStatementBuilder = preparedStatement.bind(); for (Map.Entry paramEntry : request.getRequest().getParamsMap().entrySet()) { String name = paramEntry.getKey(); Value value = paramEntry.getValue(); switch (value.getType().getKindCase()) { case BYTES_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setBytesParam(name, null); + boundStatementBuilder.setBytesParam(name, null); } else if (value.getKindCase().equals(KindCase.BYTES_VALUE)) { - statementBuilder.setBytesParam(name, value.getBytesValue()); + boundStatementBuilder.setBytesParam(name, value.getBytesValue()); } else { throw new IllegalArgumentException("Unexpected bytes value: " + value); } break; case STRING_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setStringParam(name, null); + boundStatementBuilder.setStringParam(name, null); } else if (value.getKindCase().equals(KindCase.STRING_VALUE)) { - statementBuilder.setStringParam(name, value.getStringValue()); + boundStatementBuilder.setStringParam(name, value.getStringValue()); } else { throw new IllegalArgumentException("Malformed string value: " + value); } break; case INT64_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setLongParam(name, null); + boundStatementBuilder.setLongParam(name, null); } else if (value.getKindCase().equals(KindCase.INT_VALUE)) { - statementBuilder.setLongParam(name, value.getIntValue()); + boundStatementBuilder.setLongParam(name, value.getIntValue()); } else { throw new IllegalArgumentException("Malformed int64 value: " + value); } break; case FLOAT32_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setFloatParam(name, null); + boundStatementBuilder.setFloatParam(name, null); } else if (value.getKindCase().equals(KindCase.FLOAT_VALUE)) { - statementBuilder.setFloatParam(name, (float) value.getFloatValue()); + boundStatementBuilder.setFloatParam(name, (float) value.getFloatValue()); } else { throw new IllegalArgumentException("Malformed float32 value: " + value); } break; case FLOAT64_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setDoubleParam(name, null); + boundStatementBuilder.setDoubleParam(name, null); } else if (value.getKindCase().equals(KindCase.FLOAT_VALUE)) { - statementBuilder.setDoubleParam(name, value.getFloatValue()); + boundStatementBuilder.setDoubleParam(name, value.getFloatValue()); } else { throw new IllegalArgumentException("Malformed float64 value: " + value); } break; case BOOL_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setBooleanParam(name, null); + boundStatementBuilder.setBooleanParam(name, null); } else if (value.getKindCase().equals(KindCase.BOOL_VALUE)) { - statementBuilder.setBooleanParam(name, value.getBoolValue()); + boundStatementBuilder.setBooleanParam(name, value.getBoolValue()); } else { throw new IllegalArgumentException("Malformed boolean value: " + value); } break; case TIMESTAMP_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setTimestampParam(name, null); + boundStatementBuilder.setTimestampParam(name, null); } else if (value.getKindCase().equals(KindCase.TIMESTAMP_VALUE)) { Timestamp ts = value.getTimestampValue(); - statementBuilder.setTimestampParam(name, toInstant(ts)); + boundStatementBuilder.setTimestampParam(name, toInstant(ts)); } else { throw new IllegalArgumentException("Malformed timestamp value: " + value); } break; case DATE_TYPE: if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setDateParam(name, null); + boundStatementBuilder.setDateParam(name, null); } else if (value.getKindCase().equals(KindCase.DATE_VALUE)) { com.google.type.Date protoDate = value.getDateValue(); - statementBuilder.setDateParam(name, fromProto(protoDate)); + boundStatementBuilder.setDateParam(name, fromProto(protoDate)); } else { throw new IllegalArgumentException("Malformed boolean value: " + value); } @@ -111,13 +113,13 @@ static Statement toStatement(ExecuteQueryRequest request) { case ARRAY_TYPE: SqlType.Array sqlType = (SqlType.Array) SqlType.fromProto(value.getType()); if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { - statementBuilder.setListParam(name, null, sqlType); + boundStatementBuilder.setListParam(name, null, sqlType); } else if (value.getKindCase().equals(KindCase.ARRAY_VALUE)) { List array = new ArrayList<>(); for (Value elem : value.getArrayValue().getValuesList()) { array.add(decodeArrayElement(elem, sqlType.getElementType())); } - statementBuilder.setListParam(name, array, sqlType); + boundStatementBuilder.setListParam(name, array, sqlType); } else { throw new IllegalArgumentException("Malformed array value: " + value); } @@ -126,7 +128,7 @@ static Statement toStatement(ExecuteQueryRequest request) { throw new IllegalArgumentException("Unexpected query param type in param: " + value); } } - return statementBuilder.build(); + return boundStatementBuilder.build(); } static Object decodeArrayElement(Value value, SqlType elemType) { diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java index 05731cf9c5..da205c3d3d 100644 --- a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java +++ b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java @@ -32,6 +32,7 @@ import com.google.bigtable.v2.Column; import com.google.bigtable.v2.Family; import com.google.bigtable.v2.Row; +import com.google.bigtable.v2.Value; import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.models.BulkMutation; @@ -42,7 +43,9 @@ import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.RowCell; import com.google.cloud.bigtable.data.v2.models.RowMutation; +import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; +import com.google.cloud.bigtable.data.v2.models.sql.SqlType; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; import com.google.cloud.bigtable.testproxy.CloudBigtableV2TestProxyGrpc.CloudBigtableV2TestProxyImplBase; import com.google.common.base.Preconditions; @@ -61,6 +64,7 @@ import java.io.Closeable; import java.io.IOException; import java.time.Duration; +import java.util.HashMap; import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; @@ -124,6 +128,8 @@ private static BigtableDataSettings.Builder overrideTimeoutSetting( settingsBuilder.stubSettings().sampleRowKeysSettings().retrySettings(), newTimeout); updateTimeout( settingsBuilder.stubSettings().executeQuerySettings().retrySettings(), newTimeout); + updateTimeout( + settingsBuilder.stubSettings().prepareQuerySettings().retrySettings(), newTimeout); return settingsBuilder; } @@ -687,8 +693,19 @@ public void executeQuery( responseObserver.onError(e); return; } - try (ResultSet resultSet = - client.dataClient().executeQuery(StatementDeserializer.toStatement(request))) { + ResultSet resultSet = null; + try { + Map> paramTypes = new HashMap<>(); + for (Map.Entry entry : request.getRequest().getParamsMap().entrySet()) { + paramTypes.put(entry.getKey(), SqlType.fromProto(entry.getValue().getType())); + } + PreparedStatement preparedStatement = + client.dataClient().prepareStatement(request.getRequest().getQuery(), paramTypes); + resultSet = + client + .dataClient() + .executeQuery( + BoundStatementDeserializer.toBoundStatement(preparedStatement, request)); responseObserver.onNext(ResultSetSerializer.toExecuteQueryResult(resultSet)); } catch (InterruptedException e) { responseObserver.onError(e); @@ -730,9 +747,12 @@ public void executeQuery( .build()); responseObserver.onCompleted(); return; + } finally { + if (resultSet != null) { + resultSet.close(); + } } responseObserver.onCompleted(); - return; } @Override From c2b25ef24411fbc633b08051e04cbbcd742d3477 Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Thu, 20 Mar 2025 10:10:31 -0400 Subject: [PATCH 10/11] Formatting fixes & deflake time based tests Change-Id: I5f73a4af1588d015e7e0edebf3ec73e188bf216e --- .../bigtable/data/v2/models/sql/ColumnMetadata.java | 1 - .../data/v2/models/sql/PreparedStatement.java | 1 - .../cloud/bigtable/data/v2/models/sql/ResultSet.java | 1 - .../data/v2/stub/sql/ExecuteQueryRetryTest.java | 12 ++++++------ 4 files changed, 6 insertions(+), 9 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java index e27ca6ea39..20d063922c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ColumnMetadata.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; - /** Represents the metadata for a column in a {@link ResultSet} */ public interface ColumnMetadata { /** The name of the column. Returns Empty string if the column has no name */ diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java index a45c513284..e54c86953b 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/PreparedStatement.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; - /** * The results of query preparation that can be used to create {@link BoundStatement}s to execute * queries. diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java index cf9e26d421..a149c03728 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/sql/ResultSet.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.models.sql; - /** * A set of SQL data, generated as the result of an ExecuteQuery request. * diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java index e2297911b4..6179ce3ccc 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java @@ -684,8 +684,8 @@ public void prepareFailuresBurnExecuteAttempts() throws IOException { clientWithTimeout.prepareStatement("SELECT * FROM table", new HashMap<>()); ResultSet rs = clientWithTimeout.executeQuery(ps.bind().build()); assertThrows(ApiException.class, rs::next); - // initial success plus 3 refresh failures, plus the refresh triggered by the final failure - assertThat(service.prepareCount).isEqualTo(5); + // initial success plus 3 refresh failures, plus (maybe) refresh triggered by the final failure + assertThat(service.prepareCount).isGreaterThan(3); } @Test @@ -699,9 +699,9 @@ public void canRetryAfterRefreshAttemptTimeout() throws IOException { // First attempt triggers plan refresh retry. // Second should time out, third should succeed .setMaxAttempts(3) - .setInitialRpcTimeoutDuration(Duration.ofMillis(10)) - .setMaxRpcTimeoutDuration(Duration.ofMillis(10)) - .setTotalTimeoutDuration(Duration.ofMinutes(50)) + .setInitialRpcTimeoutDuration(Duration.ofMillis(100)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(100)) + .setTotalTimeoutDuration(Duration.ofMinutes(500)) .build()) .build(); settings.stubSettings().build(); @@ -724,7 +724,7 @@ public void canRetryAfterRefreshAttemptTimeout() throws IOException { PrepareRpcExpectation.create() .withSql("SELECT * FROM table") // first refresh attempt times out, but then it succeeds - .withDelay(Duration.ofMillis(15)) + .withDelay(Duration.ofMillis(150)) .respondWith( prepareResponse( ByteString.copyFromUtf8("bar"), From 4b1ac321a039f130f9b211f5bd043f2a7911c0da Mon Sep 17 00:00:00 2001 From: Jack Dingilian Date: Thu, 20 Mar 2025 11:14:52 -0400 Subject: [PATCH 11/11] Remove accidental change to testproxy build and deflake tests Change-Id: Iea243f2c89872f947b5b79e37fa507e403f5acea --- .../internal/PreparedStatementImplTest.java | 1 + .../v2/stub/sql/ExecuteQueryRetryTest.java | 24 +++++++++---------- .../stub/sql/PlanRefreshingCallableTest.java | 3 ++- test-proxy/known_failures.txt | 3 +-- test-proxy/pom.xml | 4 ++-- 5 files changed, 18 insertions(+), 17 deletions(-) diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java index 28775af904..06f52598bc 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/PreparedStatementImplTest.java @@ -140,6 +140,7 @@ public void testBackgroundRefresh() throws InterruptedException, ExecutionExcept do { Thread.sleep(10); } while (service.prepareCount < 2); + Thread.sleep(50); PreparedQueryData updatedPlan = preparedStatement.getLatestPrepareResponse(); PrepareResponse updatedResponse = updatedPlan.prepareFuture().get(); assertThat(updatedPlan.version()).isNotEqualTo(initialPlan.version()); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java index 6179ce3ccc..a348fc9e35 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryRetryTest.java @@ -548,8 +548,8 @@ public void planRefreshRespectsAttemptTimeout() throws IOException { // First attempt triggers plan refresh retry. // Second should time out .setMaxAttempts(2) - .setInitialRpcTimeoutDuration(Duration.ofMillis(10)) - .setMaxRpcTimeoutDuration(Duration.ofMinutes(10)) + .setInitialRpcTimeoutDuration(Duration.ofMillis(500)) + .setMaxRpcTimeoutDuration(Duration.ofMinutes(500)) .setTotalTimeoutDuration(Duration.ZERO) .build()) .build(); @@ -699,9 +699,9 @@ public void canRetryAfterRefreshAttemptTimeout() throws IOException { // First attempt triggers plan refresh retry. // Second should time out, third should succeed .setMaxAttempts(3) - .setInitialRpcTimeoutDuration(Duration.ofMillis(100)) - .setMaxRpcTimeoutDuration(Duration.ofMillis(100)) - .setTotalTimeoutDuration(Duration.ofMinutes(500)) + .setInitialRpcTimeoutDuration(Duration.ofSeconds(1)) + .setMaxRpcTimeoutDuration(Duration.ofSeconds(1)) + .setTotalTimeoutDuration(Duration.ofSeconds(5)) .build()) .build(); settings.stubSettings().build(); @@ -724,7 +724,7 @@ public void canRetryAfterRefreshAttemptTimeout() throws IOException { PrepareRpcExpectation.create() .withSql("SELECT * FROM table") // first refresh attempt times out, but then it succeeds - .withDelay(Duration.ofMillis(150)) + .withDelay(Duration.ofMillis(1500)) .respondWith( prepareResponse( ByteString.copyFromUtf8("bar"), @@ -756,9 +756,9 @@ public void prepareRefreshTimeIsFactoredIntoExecuteAttemptTimeout() throws IOExc // First attempt triggers plan refresh retry. // Second should time out, third should succeed .setMaxAttempts(2) - .setInitialRpcTimeoutDuration(Duration.ofMillis(30)) - .setMaxRpcTimeoutDuration(Duration.ofMillis(30)) - .setTotalTimeoutDuration(Duration.ofMinutes(30)) + .setInitialRpcTimeoutDuration(Duration.ofMillis(500)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofMinutes(500)) .build()) .build(); settings.stubSettings().build(); @@ -779,7 +779,7 @@ public void prepareRefreshTimeIsFactoredIntoExecuteAttemptTimeout() throws IOExc PrepareRpcExpectation.create() .withSql("SELECT * FROM table") // Burn most of the execute attempt timeout and succeed - .withDelay(Duration.ofMillis(20)) + .withDelay(Duration.ofMillis(350)) .respondWith( prepareResponse( ByteString.copyFromUtf8("bar"), @@ -787,8 +787,8 @@ public void prepareRefreshTimeIsFactoredIntoExecuteAttemptTimeout() throws IOExc service.addExpectation( ExecuteRpcExpectation.create() .withPreparedQuery(ByteString.copyFromUtf8("bar")) - // Should timeout bc we used 20 ms on prepare refresh and have 30ms timeout - .withDelay(Duration.ofMillis(20)) + // Should timeout bc we used 350 ms on prepare refresh and have 500ms timeout + .withDelay(Duration.ofMillis(350)) .respondWith(partialResultSetWithToken(stringValue("s")))); PreparedStatement ps = diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java index 9fc1e7a2ea..12ca22e31c 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java @@ -243,7 +243,7 @@ public void testIsPlanRefreshError() { } @Test - public void planRefreshDelayIsFactoredIntoExecuteTimeout() { + public void planRefreshDelayIsFactoredIntoExecuteTimeout() throws InterruptedException { MockServerStreamingCallable innerCallable = new MockServerStreamingCallable<>(); RequestContext requestContext = RequestContext.create("project", "instance", "profile"); @@ -276,6 +276,7 @@ public void planRefreshDelayIsFactoredIntoExecuteTimeout() { Deadline.after(originalAttemptTimeout.toMillis() + 5, TimeUnit.MILLISECONDS); callable.call(callContext, outerObserver, context); scheduler.shutdown(); + scheduler.awaitTermination(30, TimeUnit.SECONDS); GrpcCallContext grpcCallContext = (GrpcCallContext) innerCallable.popLastCall().getApiCallContext(); Deadline executeDeadline = grpcCallContext.getCallOptions().getDeadline(); diff --git a/test-proxy/known_failures.txt b/test-proxy/known_failures.txt index ac6a05ed90..cd890631db 100644 --- a/test-proxy/known_failures.txt +++ b/test-proxy/known_failures.txt @@ -1,2 +1 @@ -TestFeatureGap/(traffic_director_enabled|direct_access_requested) -TestExecuteQuery +TestExecuteQuery|TestFeatureGap/(traffic_director_enabled|direct_access_requested) diff --git a/test-proxy/pom.xml b/test-proxy/pom.xml index c68092bb56..a0d8ac7290 100644 --- a/test-proxy/pom.xml +++ b/test-proxy/pom.xml @@ -71,9 +71,9 @@ protobuf-maven-plugin 0.6.1 - com.google.protobuf:protoc:3.22.3:exe:osx-x86_64 + com.google.protobuf:protoc:3.22.3:exe:${os.detected.classifier} grpc-java - io.grpc:protoc-gen-grpc-java:1.24.0:exe:osx-x86_64 + io.grpc:protoc-gen-grpc-java:1.24.0:exe:${os.detected.classifier}