Skip to content

Commit 86deeb8

Browse files
jackdingilianPatrick Wrobel
jackdingilian
and
Patrick Wrobel
authored
test: Add test proxy implementation for ExecuteQuery api (googleapis#2360)
* test: Test proxy support SSL backend and ExecuteQuery Change-Id: I39f81dcf098b93eff79973648c57515e8908fde3 * test: Support query params in test proxy Change-Id: I53122fdf0301fb41d4f24881dc8de65d174dee35 * Remove security options for a separate PR Change-Id: Ibe2fb899f83744d4e38debc58379e9942350a83f --------- Co-authored-by: Patrick Wrobel <pwrobel@google.com>
1 parent a59572a commit 86deeb8

4 files changed

Lines changed: 504 additions & 2 deletions

File tree

‎java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java‎

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2022 Google LLC
2+
* Copyright 2024 Google LLC
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -41,6 +41,7 @@
4141
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
4242
import com.google.cloud.bigtable.data.v2.models.RowCell;
4343
import com.google.cloud.bigtable.data.v2.models.RowMutation;
44+
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
4445
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
4546
import com.google.cloud.bigtable.testproxy.CloudBigtableV2TestProxyGrpc.CloudBigtableV2TestProxyImplBase;
4647
import com.google.common.base.Preconditions;
@@ -50,6 +51,7 @@
5051
import io.grpc.ManagedChannelBuilder;
5152
import io.grpc.Status;
5253
import io.grpc.StatusException;
54+
import io.grpc.StatusRuntimeException;
5355
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
5456
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
5557
import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext;
@@ -65,6 +67,7 @@
6567
import java.util.List;
6668
import java.util.Map;
6769
import java.util.concurrent.ConcurrentHashMap;
70+
import java.util.concurrent.ExecutionException;
6871
import java.util.logging.Logger;
6972
import java.util.regex.Matcher;
7073
import java.util.regex.Pattern;
@@ -159,6 +162,8 @@ private static BigtableDataSettings.Builder overrideTimeoutSetting(
159162
settingsBuilder.stubSettings().readModifyWriteRowSettings().retrySettings(), newTimeout);
160163
updateTimeout(
161164
settingsBuilder.stubSettings().sampleRowKeysSettings().retrySettings(), newTimeout);
165+
updateTimeout(
166+
settingsBuilder.stubSettings().executeQuerySettings().retrySettings(), newTimeout);
162167

163168
return settingsBuilder;
164169
}
@@ -698,6 +703,64 @@ public void readModifyWriteRow(
698703
responseObserver.onCompleted();
699704
}
700705

706+
@Override
707+
public void executeQuery(
708+
ExecuteQueryRequest request, StreamObserver<ExecuteQueryResult> responseObserver) {
709+
CbtClient client;
710+
try {
711+
client = getClient(request.getClientId());
712+
} catch (StatusException e) {
713+
responseObserver.onError(e);
714+
return;
715+
}
716+
try (ResultSet resultSet =
717+
client.dataClient().executeQuery(StatementDeserializer.toStatement(request))) {
718+
responseObserver.onNext(ResultSetSerializer.toExecuteQueryResult(resultSet));
719+
} catch (InterruptedException e) {
720+
responseObserver.onError(e);
721+
return;
722+
} catch (ExecutionException e) {
723+
responseObserver.onError(e);
724+
return;
725+
} catch (ApiException e) {
726+
responseObserver.onNext(
727+
ExecuteQueryResult.newBuilder()
728+
.setStatus(
729+
com.google.rpc.Status.newBuilder()
730+
.setCode(e.getStatusCode().getCode().ordinal())
731+
.setMessage(e.getMessage())
732+
.build())
733+
.build());
734+
responseObserver.onCompleted();
735+
return;
736+
} catch (StatusRuntimeException e) {
737+
responseObserver.onNext(
738+
ExecuteQueryResult.newBuilder()
739+
.setStatus(
740+
com.google.rpc.Status.newBuilder()
741+
.setCode(e.getStatus().getCode().value())
742+
.setMessage(e.getStatus().getDescription())
743+
.build())
744+
.build());
745+
responseObserver.onCompleted();
746+
return;
747+
} catch (RuntimeException e) {
748+
// If client encounters problem, don't return any results.
749+
responseObserver.onNext(
750+
ExecuteQueryResult.newBuilder()
751+
.setStatus(
752+
com.google.rpc.Status.newBuilder()
753+
.setCode(Code.INTERNAL.getNumber())
754+
.setMessage(e.getMessage())
755+
.build())
756+
.build());
757+
responseObserver.onCompleted();
758+
return;
759+
}
760+
responseObserver.onCompleted();
761+
return;
762+
}
763+
701764
@Override
702765
public synchronized void close() {
703766
Iterator<Map.Entry<String, CbtClient>> it = idClientMap.entrySet().iterator();

‎java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/ResultSetSerializer.java‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.testproxy;
17+
18+
import com.google.bigtable.v2.ArrayValue;
19+
import com.google.bigtable.v2.Type;
20+
import com.google.bigtable.v2.Type.Array;
21+
import com.google.bigtable.v2.Type.Bool;
22+
import com.google.bigtable.v2.Type.Bytes;
23+
import com.google.bigtable.v2.Type.Float32;
24+
import com.google.bigtable.v2.Type.Float64;
25+
import com.google.bigtable.v2.Type.Int64;
26+
import com.google.bigtable.v2.Type.Map;
27+
import com.google.bigtable.v2.Type.Struct;
28+
import com.google.bigtable.v2.Type.Timestamp;
29+
import com.google.bigtable.v2.Value;
30+
import com.google.cloud.Date;
31+
import com.google.cloud.bigtable.data.v2.models.sql.ColumnMetadata;
32+
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
33+
import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
34+
import com.google.cloud.bigtable.data.v2.models.sql.StructReader;
35+
import com.google.protobuf.ByteString;
36+
import java.util.List;
37+
import java.util.concurrent.ExecutionException;
38+
import org.threeten.bp.Instant;
39+
40+
public class ResultSetSerializer {
41+
public static ExecuteQueryResult toExecuteQueryResult(ResultSet resultSet)
42+
throws ExecutionException, InterruptedException {
43+
ExecuteQueryResult.Builder resultBuilder = ExecuteQueryResult.newBuilder();
44+
for (ColumnMetadata columnMetadata : resultSet.getMetadata().getColumns()) {
45+
resultBuilder
46+
.getMetadataBuilder()
47+
.addColumnsBuilder()
48+
.setName(columnMetadata.name())
49+
.setType(toProtoType(columnMetadata.type()));
50+
}
51+
52+
while (resultSet.next()) {
53+
SqlRow.Builder rowBuilder = resultBuilder.addRowsBuilder();
54+
55+
for (int i = 0; i < resultSet.getMetadata().getColumns().size(); i++) {
56+
SqlType<?> colType = resultSet.getMetadata().getColumnType(i);
57+
rowBuilder.addValues(toProtoValue(getColumn(resultSet, i, colType), colType));
58+
}
59+
}
60+
61+
return resultBuilder.build();
62+
}
63+
64+
private static Value toProtoValue(Object value, SqlType<?> type) {
65+
if (value == null) {
66+
return Value.getDefaultInstance();
67+
}
68+
69+
Value.Builder valueBuilder = Value.newBuilder();
70+
switch (type.getCode()) {
71+
case BYTES:
72+
valueBuilder.setBytesValue((ByteString) value);
73+
break;
74+
case STRING:
75+
valueBuilder.setStringValue((String) value);
76+
break;
77+
case INT64:
78+
valueBuilder.setIntValue((Long) value);
79+
break;
80+
case FLOAT32:
81+
valueBuilder.setFloatValue((Float) value);
82+
break;
83+
case FLOAT64:
84+
valueBuilder.setFloatValue((Double) value);
85+
break;
86+
case BOOL:
87+
valueBuilder.setBoolValue((Boolean) value);
88+
break;
89+
case TIMESTAMP:
90+
Instant ts = (Instant) value;
91+
valueBuilder.setTimestampValue(
92+
com.google.protobuf.Timestamp.newBuilder()
93+
.setSeconds(ts.getEpochSecond())
94+
.setNanos(ts.getNano())
95+
.build());
96+
break;
97+
case DATE:
98+
Date date = (Date) value;
99+
valueBuilder.setDateValue(
100+
com.google.type.Date.newBuilder()
101+
.setYear(date.getYear())
102+
.setMonth(date.getMonth())
103+
.setDay(date.getDayOfMonth())
104+
.build());
105+
break;
106+
case ARRAY:
107+
SqlType<?> elementType = ((SqlType.Array<?>) type).getElementType();
108+
ArrayValue.Builder arrayValue = ArrayValue.newBuilder();
109+
for (Object item : (List<?>) value) {
110+
arrayValue.addValues(toProtoValue(item, elementType));
111+
}
112+
valueBuilder.setArrayValue(arrayValue.build());
113+
break;
114+
case MAP:
115+
SqlType.Map<?, ?> mapType = (SqlType.Map<?, ?>) type;
116+
SqlType<?> mapKeyType = mapType.getKeyType();
117+
SqlType<?> mapValueType = mapType.getValueType();
118+
119+
ArrayValue.Builder mapArrayValue = ArrayValue.newBuilder();
120+
((java.util.Map<?, ?>) value)
121+
.forEach(
122+
(k, v) ->
123+
mapArrayValue.addValues(
124+
Value.newBuilder()
125+
.setArrayValue(
126+
ArrayValue.newBuilder()
127+
.addValues(toProtoValue(k, mapKeyType))
128+
.addValues(toProtoValue(v, mapValueType))
129+
.build())));
130+
valueBuilder.setArrayValue(mapArrayValue.build());
131+
break;
132+
case STRUCT:
133+
StructReader structValue = (StructReader) value;
134+
SqlType.Struct structType = (SqlType.Struct) type;
135+
ArrayValue.Builder structArrayValue = ArrayValue.newBuilder();
136+
for (int i = 0; i < structType.getFields().size(); ++i) {
137+
SqlType<?> fieldType = structType.getType(i);
138+
structArrayValue.addValues(toProtoValue(getColumn(structValue, i, fieldType), fieldType));
139+
}
140+
valueBuilder.setArrayValue(structArrayValue);
141+
break;
142+
default:
143+
throw new IllegalStateException("Unexpected Type: " + type);
144+
}
145+
146+
return valueBuilder.build();
147+
}
148+
149+
private static Object getColumn(StructReader struct, int fieldIndex, SqlType<?> fieldType) {
150+
if (struct.isNull(fieldIndex)) {
151+
return null;
152+
}
153+
154+
switch (fieldType.getCode()) {
155+
case ARRAY:
156+
return struct.getList(fieldIndex, (SqlType.Array<?>) fieldType);
157+
case BOOL:
158+
return struct.getBoolean(fieldIndex);
159+
case BYTES:
160+
return struct.getBytes(fieldIndex);
161+
case DATE:
162+
return struct.getDate(fieldIndex);
163+
case FLOAT32:
164+
return struct.getFloat(fieldIndex);
165+
case FLOAT64:
166+
return struct.getDouble(fieldIndex);
167+
case INT64:
168+
return struct.getLong(fieldIndex);
169+
case MAP:
170+
return struct.getMap(fieldIndex, (SqlType.Map<?, ?>) fieldType);
171+
case STRING:
172+
return struct.getString(fieldIndex);
173+
case STRUCT:
174+
return struct.getStruct(fieldIndex);
175+
case TIMESTAMP:
176+
return struct.getTimestamp(fieldIndex);
177+
default:
178+
throw new IllegalStateException("Unexpected Type: " + fieldType);
179+
}
180+
}
181+
182+
private static Type toProtoType(SqlType<?> type) {
183+
switch (type.getCode()) {
184+
case BYTES:
185+
return Type.newBuilder().setBytesType(Bytes.getDefaultInstance()).build();
186+
case STRING:
187+
return Type.newBuilder()
188+
.setStringType(com.google.bigtable.v2.Type.String.getDefaultInstance())
189+
.build();
190+
case INT64:
191+
return Type.newBuilder().setInt64Type(Int64.getDefaultInstance()).build();
192+
case FLOAT32:
193+
return Type.newBuilder().setFloat32Type(Float32.getDefaultInstance()).build();
194+
case FLOAT64:
195+
return Type.newBuilder().setFloat64Type(Float64.getDefaultInstance()).build();
196+
case BOOL:
197+
return Type.newBuilder().setBoolType(Bool.getDefaultInstance()).build();
198+
case TIMESTAMP:
199+
return Type.newBuilder().setTimestampType(Timestamp.getDefaultInstance()).build();
200+
case DATE:
201+
return Type.newBuilder()
202+
.setDateType(com.google.bigtable.v2.Type.Date.getDefaultInstance())
203+
.build();
204+
case ARRAY:
205+
SqlType.Array<?> arrayType = (SqlType.Array<?>) type;
206+
return Type.newBuilder()
207+
.setArrayType(
208+
Array.newBuilder().setElementType(toProtoType(arrayType.getElementType())))
209+
.build();
210+
case MAP:
211+
SqlType.Map<?, ?> mapType = (SqlType.Map<?, ?>) type;
212+
return Type.newBuilder()
213+
.setMapType(
214+
Map.newBuilder()
215+
.setKeyType(toProtoType(mapType.getKeyType()))
216+
.setValueType(toProtoType(mapType.getValueType())))
217+
.build();
218+
case STRUCT:
219+
SqlType.Struct structType = (SqlType.Struct) type;
220+
Struct.Builder structBuilder = Struct.newBuilder();
221+
for (SqlType.Struct.Field field : structType.getFields()) {
222+
structBuilder
223+
.addFieldsBuilder()
224+
.setFieldName(field.name())
225+
.setType(toProtoType(field.type()));
226+
}
227+
return Type.newBuilder().setStructType(structBuilder.build()).build();
228+
229+
default:
230+
throw new IllegalStateException("Unexpected Type: " + type);
231+
}
232+
}
233+
}

0 commit comments

Comments
 (0)