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 0859b2b50f..068fe54f76 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 @@ -43,7 +43,6 @@ 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.ResultSet; -import com.google.cloud.bigtable.data.v2.models.sql.Statement; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; import com.google.cloud.bigtable.testproxy.CloudBigtableV2TestProxyGrpc.CloudBigtableV2TestProxyImplBase; import com.google.common.base.Preconditions; @@ -53,6 +52,7 @@ import io.grpc.ManagedChannelBuilder; import io.grpc.Status; import io.grpc.StatusException; +import io.grpc.StatusRuntimeException; import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts; import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder; import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext; @@ -122,6 +122,8 @@ private static BigtableDataSettings.Builder overrideTimeoutSetting( settingsBuilder.stubSettings().readModifyWriteRowSettings().retrySettings(), newTimeout); updateTimeout( settingsBuilder.stubSettings().sampleRowKeysSettings().retrySettings(), newTimeout); + updateTimeout( + settingsBuilder.stubSettings().executeQuerySettings().retrySettings(), newTimeout); return settingsBuilder; } @@ -680,14 +682,15 @@ public void executeQuery( responseObserver.onError(e); return; } - try (ResultSet resultSet = - client.dataClient().executeQuery(Statement.of(request.getRequest().getQuery()))) { + client.dataClient().executeQuery(StatementDeserializer.toStatement(request))) { responseObserver.onNext(ResultSetSerializer.toExecuteQueryResult(resultSet)); } catch (InterruptedException e) { responseObserver.onError(e); + return; } catch (ExecutionException e) { responseObserver.onError(e); + return; } catch (ApiException e) { responseObserver.onNext( ExecuteQueryResult.newBuilder() @@ -697,12 +700,33 @@ public void executeQuery( .setMessage(e.getMessage()) .build()) .build()); + responseObserver.onCompleted(); + return; + } catch (StatusRuntimeException e) { + responseObserver.onNext( + ExecuteQueryResult.newBuilder() + .setStatus( + com.google.rpc.Status.newBuilder() + .setCode(e.getStatus().getCode().value()) + .setMessage(e.getStatus().getDescription()) + .build()) + .build()); + responseObserver.onCompleted(); + return; } catch (RuntimeException e) { - responseObserver.onError(e); - } finally { + // If client encounters problem, don't return any results. + responseObserver.onNext( + ExecuteQueryResult.newBuilder() + .setStatus( + com.google.rpc.Status.newBuilder() + .setCode(Code.INTERNAL.getNumber()) + .setMessage(e.getMessage()) + .build()) + .build()); responseObserver.onCompleted(); + return; } - + responseObserver.onCompleted(); return; } diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/ResultSetSerializer.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/ResultSetSerializer.java index 966e688cd8..c138c82a6b 100644 --- a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/ResultSetSerializer.java +++ b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/ResultSetSerializer.java @@ -71,26 +71,21 @@ private static Value toProtoValue(Object value, SqlType type) { case BYTES: valueBuilder.setBytesValue((ByteString) value); break; - case STRING: valueBuilder.setStringValue((String) value); break; case INT64: valueBuilder.setIntValue((Long) value); break; - case FLOAT32: valueBuilder.setFloatValue((Float) value); break; - case FLOAT64: valueBuilder.setFloatValue((Double) value); break; - case BOOL: valueBuilder.setBoolValue((Boolean) value); break; - case TIMESTAMP: Instant ts = (Instant) value; valueBuilder.setTimestampValue( @@ -99,7 +94,6 @@ private static Value toProtoValue(Object value, SqlType type) { .setNanos(ts.getNano()) .build()); break; - case DATE: Date date = (Date) value; valueBuilder.setDateValue( @@ -109,7 +103,6 @@ private static Value toProtoValue(Object value, SqlType type) { .setDay(date.getDayOfMonth()) .build()); break; - case ARRAY: SqlType elementType = ((SqlType.Array) type).getElementType(); ArrayValue.Builder arrayValue = ArrayValue.newBuilder(); @@ -118,7 +111,6 @@ private static Value toProtoValue(Object value, SqlType type) { } valueBuilder.setArrayValue(arrayValue.build()); break; - case MAP: SqlType.Map mapType = (SqlType.Map) type; SqlType mapKeyType = mapType.getKeyType(); @@ -137,7 +129,6 @@ private static Value toProtoValue(Object value, SqlType type) { .build()))); valueBuilder.setArrayValue(mapArrayValue.build()); break; - case STRUCT: StructReader structValue = (StructReader) value; SqlType.Struct structType = (SqlType.Struct) type; @@ -148,7 +139,6 @@ private static Value toProtoValue(Object value, SqlType type) { } valueBuilder.setArrayValue(structArrayValue); break; - default: throw new IllegalStateException("Unexpected Type: " + type); } @@ -208,7 +198,9 @@ private static Type toProtoType(SqlType type) { case TIMESTAMP: return Type.newBuilder().setTimestampType(Timestamp.getDefaultInstance()).build(); case DATE: - return Type.newBuilder().setBytesType(Bytes.getDefaultInstance()).build(); + return Type.newBuilder() + .setDateType(com.google.bigtable.v2.Type.Date.getDefaultInstance()) + .build(); case ARRAY: SqlType.Array arrayType = (SqlType.Array) type; return Type.newBuilder() 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/StatementDeserializer.java new file mode 100644 index 0000000000..ae3b50aa7f --- /dev/null +++ b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/StatementDeserializer.java @@ -0,0 +1,167 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.testproxy; + +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.SqlType; +import com.google.cloud.bigtable.data.v2.models.sql.Statement; +import com.google.protobuf.Timestamp; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.threeten.bp.Instant; + +public class StatementDeserializer { + + static Statement toStatement(ExecuteQueryRequest request) { + Statement.Builder statementBuilder = Statement.newBuilder(request.getRequest().getQuery()); + 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); + } else if (value.getKindCase().equals(KindCase.BYTES_VALUE)) { + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.STRING_VALUE)) { + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.INT_VALUE)) { + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.FLOAT_VALUE)) { + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.FLOAT_VALUE)) { + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.BOOL_VALUE)) { + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.TIMESTAMP_VALUE)) { + Timestamp ts = value.getTimestampValue(); + statementBuilder.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); + } else if (value.getKindCase().equals(KindCase.DATE_VALUE)) { + com.google.type.Date protoDate = value.getDateValue(); + statementBuilder.setDateParam(name, fromProto(protoDate)); + } else { + throw new IllegalArgumentException("Malformed boolean value: " + value); + } + break; + 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); + } 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); + } else { + throw new IllegalArgumentException("Malformed array value: " + value); + } + break; + default: + throw new IllegalArgumentException("Unexpected query param type in param: " + value); + } + } + return statementBuilder.build(); + } + + static Object decodeArrayElement(Value value, SqlType elemType) { + if (value.getKindCase().equals(KindCase.KIND_NOT_SET)) { + return null; + } + switch (elemType.getCode()) { + case BYTES: + return value.getBytesValue(); + case STRING: + return value.getStringValue(); + case INT64: + return value.getIntValue(); + case FLOAT64: + return value.getFloatValue(); + case FLOAT32: + // cast to float so we produce List, etc + return (float) value.getFloatValue(); + case BOOL: + return value.getBoolValue(); + case TIMESTAMP: + return toInstant(value.getTimestampValue()); + case DATE: + return fromProto(value.getDateValue()); + default: + // We should have already thrown an exception in the SqlRowMerger + throw new IllegalStateException("Unsupported array query param element type: " + elemType); + } + } + + private static Instant toInstant(Timestamp timestamp) { + return Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()); + } + + private static Date fromProto(com.google.type.Date proto) { + return Date.fromYearMonthDay(proto.getYear(), proto.getMonth(), proto.getDay()); + } +} diff --git a/test-proxy/src/main/proto/test_proxy.proto b/test-proxy/src/main/proto/test_proxy.proto index 0c9b64b8cd..b82354b08e 100644 --- a/test-proxy/src/main/proto/test_proxy.proto +++ b/test-proxy/src/main/proto/test_proxy.proto @@ -249,18 +249,22 @@ message ReadModifyWriteRowRequest { google.bigtable.v2.ReadModifyWriteRowRequest request = 2; } +// Request to test proxy service to execute a query. message ExecuteQueryRequest { // The ID of the target client object. string client_id = 1; + // The raw request to the Bigtable server. google.bigtable.v2.ExecuteQueryRequest request = 2; } +// Response from test proxy service for ExecuteQueryRequest. message ExecuteQueryResult { // The RPC status from the client binding. google.rpc.Status status = 1; - google.bigtable.v2.ResultSetMetadata result_set_metadata = 2; // deprecated + // deprecated + google.bigtable.v2.ResultSetMetadata result_set_metadata = 2; // Name and type information for the query result. ResultSetMetadata metadata = 4; @@ -269,10 +273,13 @@ message ExecuteQueryResult { repeated SqlRow rows = 3; } +// Schema information for the query result. message ResultSetMetadata { + // Column metadata for each column inthe query result. repeated google.bigtable.v2.ColumnMetadata columns = 1; } +// Representation of a single row in the query result. message SqlRow { // Columnar values returned by the query. repeated google.bigtable.v2.Value values = 1;