Skip to content

Commit 9cbfd50

Browse files
trollyxiamutianf
trollyxia
and
mutianf
authored
textproxy: Allow testproxy to build its own proto registry (#2682)
* textproxy: Allow testproxy to build its own proto registry Change-Id: Ie930064363d92d61daad498e8380dd87ab6722dd * testproxy: Remove manual dependency resolution from ResultSetSerializer We now require the clients to provide FileDescriptorSet with files in dependency order Change-Id: I240cfe8f4d499ba7053dfab1f119972de1810fce --------- Co-authored-by: Mattie Fu <[email protected]>
1 parent 5c616df commit 9cbfd50

File tree

3 files changed

+146
-8
lines changed

3 files changed

+146
-8
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -706,7 +706,8 @@ public void executeQuery(
706706
.dataClient()
707707
.executeQuery(
708708
BoundStatementDeserializer.toBoundStatement(preparedStatement, request));
709-
responseObserver.onNext(ResultSetSerializer.toExecuteQueryResult(resultSet));
709+
responseObserver.onNext(
710+
new ResultSetSerializer(request.getProtoDescriptors()).toExecuteQueryResult(resultSet));
710711
} catch (InterruptedException e) {
711712
responseObserver.onError(e);
712713
return;

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

Lines changed: 140 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,124 @@
3535
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
3636
import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
3737
import com.google.cloud.bigtable.data.v2.models.sql.StructReader;
38+
import com.google.protobuf.AbstractMessage;
3839
import com.google.protobuf.ByteString;
40+
import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
41+
import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
42+
import com.google.protobuf.Descriptors.Descriptor;
43+
import com.google.protobuf.Descriptors.DescriptorValidationException;
44+
import com.google.protobuf.Descriptors.EnumDescriptor;
45+
import com.google.protobuf.Descriptors.FileDescriptor;
46+
import com.google.protobuf.DynamicMessage;
47+
import com.google.protobuf.ProtocolMessageEnum;
3948
import java.time.Instant;
49+
import java.util.ArrayList;
50+
import java.util.HashMap;
4051
import java.util.List;
4152
import java.util.concurrent.ExecutionException;
4253

4354
public class ResultSetSerializer {
44-
public static ExecuteQueryResult toExecuteQueryResult(ResultSet resultSet)
55+
56+
// This is a helper enum to satisfy the type constraints of {@link StructReader#getProtoEnum}.
57+
private static class DummyEnum implements ProtocolMessageEnum {
58+
59+
private final int value;
60+
private final EnumDescriptor descriptor;
61+
62+
private DummyEnum(int value, EnumDescriptor descriptor) {
63+
this.value = value;
64+
this.descriptor = descriptor;
65+
}
66+
67+
@Override
68+
public int getNumber() {
69+
return value;
70+
}
71+
72+
@Override
73+
public com.google.protobuf.Descriptors.EnumValueDescriptor getValueDescriptor() {
74+
return descriptor.findValueByNumber(value);
75+
}
76+
77+
@Override
78+
public com.google.protobuf.Descriptors.EnumDescriptor getDescriptorForType() {
79+
return descriptor;
80+
}
81+
}
82+
83+
/**
84+
* A map of all known message descriptors, keyed by their fully qualified name (e.g.,
85+
* "my.package.MyMessage").
86+
*/
87+
private final java.util.Map<String, Descriptor> messageDescriptorMap;
88+
89+
/**
90+
* A map of all known enum descriptors, keyed by their fully qualified name (e.g.,
91+
* "my.package.MyEnum").
92+
*/
93+
private final java.util.Map<String, EnumDescriptor> enumDescriptorMap;
94+
95+
/**
96+
* Helper function to recursively adds a message descriptor and all its nested types to the map.
97+
*/
98+
private void populateDescriptorMapsRecursively(Descriptor descriptor) {
99+
messageDescriptorMap.put(descriptor.getFullName(), descriptor);
100+
101+
for (EnumDescriptor nestedEnum : descriptor.getEnumTypes()) {
102+
enumDescriptorMap.put(nestedEnum.getFullName(), nestedEnum);
103+
}
104+
for (Descriptor nestedMessage : descriptor.getNestedTypes()) {
105+
populateDescriptorMapsRecursively(nestedMessage);
106+
}
107+
}
108+
109+
/**
110+
* Creates a serializer with a descriptor cache built from the provided FileDescriptorSet. This is
111+
* useful for handling PROTO or ENUM types that require schema lookup.
112+
*
113+
* @param descriptorSet A set containing one or more .proto file definitions and all their
114+
* non-standard dependencies. All .proto file must be provided in dependency order.
115+
* @throws IllegalArgumentException if the descriptorSet contains unresolvable dependencies.
116+
*/
117+
public ResultSetSerializer(FileDescriptorSet descriptorSet) throws IllegalArgumentException {
118+
this.messageDescriptorMap = new HashMap<>();
119+
this.enumDescriptorMap = new HashMap<>();
120+
java.util.Map<String, FileDescriptor> builtDescriptors = new HashMap<>();
121+
122+
for (FileDescriptorProto fileDescriptorProto : descriptorSet.getFileList()) {
123+
// Collect dependencies. This code require files inside the descriptor set to be sorted
124+
// according to the dependency order.
125+
List<FileDescriptor> dependencies = new ArrayList<>();
126+
for (String dependencyName : fileDescriptorProto.getDependencyList()) {
127+
FileDescriptor dependency = builtDescriptors.get(dependencyName);
128+
if (dependency != null) {
129+
// Dependency is already built, add it.
130+
dependencies.add(dependency);
131+
}
132+
// Dependency is not in our set. We assume it's a well-known type (e.g.,
133+
// google/protobuf/timestamp.proto) that buildFrom() can find and link automatically.
134+
}
135+
136+
try {
137+
FileDescriptor fileDescriptor =
138+
FileDescriptor.buildFrom(
139+
fileDescriptorProto, dependencies.toArray(new FileDescriptor[0]));
140+
builtDescriptors.put(fileDescriptor.getName(), fileDescriptor);
141+
// Now, populate both message and enum maps with all messages/enums in this file.
142+
for (EnumDescriptor enumDescriptor : fileDescriptor.getEnumTypes()) {
143+
enumDescriptorMap.put(enumDescriptor.getFullName(), enumDescriptor);
144+
}
145+
for (Descriptor messageDescriptor : fileDescriptor.getMessageTypes()) {
146+
populateDescriptorMapsRecursively(messageDescriptor);
147+
}
148+
} catch (DescriptorValidationException e) {
149+
throw new IllegalArgumentException(
150+
"Failed to build descriptor for " + fileDescriptorProto.getName(), e);
151+
}
152+
}
153+
}
154+
155+
public ExecuteQueryResult toExecuteQueryResult(ResultSet resultSet)
45156
throws ExecutionException, InterruptedException {
46157
ExecuteQueryResult.Builder resultBuilder = ExecuteQueryResult.newBuilder();
47158
for (ColumnMetadata columnMetadata : resultSet.getMetadata().getColumns()) {
@@ -64,24 +175,28 @@ public static ExecuteQueryResult toExecuteQueryResult(ResultSet resultSet)
64175
return resultBuilder.build();
65176
}
66177

67-
private static Value toProtoValue(Object value, SqlType<?> type) {
178+
private Value toProtoValue(Object value, SqlType<?> type) {
68179
if (value == null) {
69180
return Value.getDefaultInstance();
70181
}
71182

72183
Value.Builder valueBuilder = Value.newBuilder();
73184
switch (type.getCode()) {
74185
case BYTES:
75-
case PROTO:
76186
valueBuilder.setBytesValue((ByteString) value);
77187
break;
188+
case PROTO:
189+
valueBuilder.setBytesValue(((AbstractMessage) value).toByteString());
190+
break;
78191
case STRING:
79192
valueBuilder.setStringValue((String) value);
80193
break;
81194
case INT64:
82-
case ENUM:
83195
valueBuilder.setIntValue((Long) value);
84196
break;
197+
case ENUM:
198+
valueBuilder.setIntValue(((ProtocolMessageEnum) value).getNumber());
199+
break;
85200
case FLOAT32:
86201
valueBuilder.setFloatValue((Float) value);
87202
break;
@@ -151,7 +266,7 @@ private static Value toProtoValue(Object value, SqlType<?> type) {
151266
return valueBuilder.build();
152267
}
153268

154-
private static Object getColumn(StructReader struct, int fieldIndex, SqlType<?> fieldType) {
269+
private Object getColumn(StructReader struct, int fieldIndex, SqlType<?> fieldType) {
155270
if (struct.isNull(fieldIndex)) {
156271
return null;
157272
}
@@ -162,17 +277,35 @@ private static Object getColumn(StructReader struct, int fieldIndex, SqlType<?>
162277
case BOOL:
163278
return struct.getBoolean(fieldIndex);
164279
case BYTES:
165-
case PROTO:
166280
return struct.getBytes(fieldIndex);
281+
case PROTO:
282+
SchemalessProto protoType = (SchemalessProto) fieldType;
283+
Descriptor descriptor = messageDescriptorMap.get(protoType.getMessageName());
284+
if (descriptor == null) {
285+
throw new IllegalArgumentException(
286+
"Descriptor for message " + protoType.getMessageName() + " could not be found");
287+
}
288+
return struct.getProtoMessage(fieldIndex, DynamicMessage.getDefaultInstance(descriptor));
167289
case DATE:
168290
return struct.getDate(fieldIndex);
169291
case FLOAT32:
170292
return struct.getFloat(fieldIndex);
171293
case FLOAT64:
172294
return struct.getDouble(fieldIndex);
173295
case INT64:
174-
case ENUM:
175296
return struct.getLong(fieldIndex);
297+
case ENUM:
298+
SchemalessEnum enumType = (SchemalessEnum) fieldType;
299+
EnumDescriptor enumDescriptor = enumDescriptorMap.get(enumType.getEnumName());
300+
if (enumDescriptor == null) {
301+
throw new IllegalArgumentException(
302+
"Descriptor for enum " + enumType.getEnumName() + " could not be found");
303+
}
304+
// We need to extract the integer value of the enum. `getProtoEnum` is the only
305+
// available method, but it is designed for static enum types. To work around this,
306+
// we can pass a lambda that constructs our DummyEnum with the captured integer value
307+
// and the descriptor from the outer scope.
308+
return struct.getProtoEnum(fieldIndex, number -> new DummyEnum(number, enumDescriptor));
176309
case MAP:
177310
return struct.getMap(fieldIndex, (SqlType.Map<?, ?>) fieldType);
178311
case STRING:

‎test-proxy/src/main/proto/test_proxy.proto‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import "google/api/client.proto";
2020
import "google/bigtable/v2/bigtable.proto";
2121
import "google/bigtable/v2/data.proto";
2222
import "google/protobuf/duration.proto";
23+
import "google/protobuf/descriptor.proto";
2324
import "google/rpc/status.proto";
2425

2526
option go_package = "./testproxypb";
@@ -256,6 +257,9 @@ message ExecuteQueryRequest {
256257

257258
// The raw request to the Bigtable server.
258259
google.bigtable.v2.ExecuteQueryRequest request = 2;
260+
261+
// The file descriptor set for the query.
262+
google.protobuf.FileDescriptorSet proto_descriptors = 3;
259263
}
260264

261265
// Response from test proxy service for ExecuteQueryRequest.

0 commit comments

Comments
 (0)