Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit 494185b

Browse files
trollyxiatrollyxia
committed
text-proxy: Allow testproxy to build its own proto registry
Change-Id: Ie930064363d92d61daad498e8380dd87ab6722dd
1 parent 6e60c01 commit 494185b

3 files changed

Lines changed: 177 additions & 8 deletions

File tree

‎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: 171 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,155 @@
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+
private final int value;
59+
private final EnumDescriptor descriptor;
60+
61+
private DummyEnum(int value, EnumDescriptor descriptor) {
62+
this.value = value;
63+
this.descriptor = descriptor;
64+
}
65+
66+
@Override
67+
public int getNumber() {
68+
return value;
69+
}
70+
71+
@Override
72+
public com.google.protobuf.Descriptors.EnumValueDescriptor getValueDescriptor() {
73+
return descriptor.findValueByNumber(value);
74+
}
75+
76+
@Override
77+
public com.google.protobuf.Descriptors.EnumDescriptor getDescriptorForType() {
78+
return descriptor;
79+
}
80+
}
81+
82+
/**
83+
* A map of all known message descriptors, keyed by their fully qualified name (e.g.,
84+
* "my.package.MyMessage").
85+
*/
86+
private final java.util.Map<String, Descriptor> messageDescriptorMap;
87+
88+
/**
89+
* A map of all known enum descriptors, keyed by their fully qualified name (e.g.,
90+
* "my.package.MyEnum").
91+
*/
92+
private final java.util.Map<String, EnumDescriptor> enumDescriptorMap;
93+
94+
/**
95+
* Helper function to recursively adds a message descriptor and all its nested types to the map.
96+
*/
97+
private void populateDescriptorMapsRecursively(Descriptor descriptor) {
98+
messageDescriptorMap.put(descriptor.getFullName(), descriptor);
99+
100+
for (EnumDescriptor nestedEnum : descriptor.getEnumTypes()) {
101+
enumDescriptorMap.put(nestedEnum.getFullName(), nestedEnum);
102+
}
103+
for (Descriptor nestedMessage : descriptor.getNestedTypes()) {
104+
populateDescriptorMapsRecursively(nestedMessage);
105+
}
106+
}
107+
108+
/**
109+
* Creates a serializer with a descriptor cache built from the provided FileDescriptorSet. This is
110+
* useful for handling PROTO or ENUM types that require schema lookup.
111+
*
112+
* @param descriptorSet A set containing one or more .proto file definitions and all their
113+
* non-standard dependencies.
114+
* @throws IllegalArgumentException if the descriptorSet contains unresolvable dependencies.
115+
*/
116+
public ResultSetSerializer(FileDescriptorSet descriptorSet) throws IllegalArgumentException {
117+
this.messageDescriptorMap = new HashMap<>();
118+
this.enumDescriptorMap = new HashMap<>();
119+
120+
// Java's dynamic descriptor loading requires manual dependency resolution.
121+
// We must build a DAG and build files only after their dependencies are built.
122+
java.util.Map<String, FileDescriptorProto> pendingProtos = new HashMap<>();
123+
java.util.Map<String, FileDescriptor> builtDescriptors = new HashMap<>();
124+
for (FileDescriptorProto fileDescriptorProto : descriptorSet.getFileList()) {
125+
pendingProtos.put(fileDescriptorProto.getName(), fileDescriptorProto);
126+
}
127+
128+
// Keep looping as long as we are making progress.
129+
int filesBuiltThisPass;
130+
do {
131+
filesBuiltThisPass = 0;
132+
// Use a copy of the values, to avoid modifying the pendingProtos map while iterating it.
133+
List<FileDescriptorProto> protosToBuild = new ArrayList<>(pendingProtos.values());
134+
135+
for (FileDescriptorProto proto : protosToBuild) {
136+
boolean allDependenciesMet = true;
137+
List<FileDescriptor> dependencies = new ArrayList<>();
138+
139+
for (String dependencyName : proto.getDependencyList()) {
140+
FileDescriptor dependency = builtDescriptors.get(dependencyName);
141+
if (dependency != null) {
142+
// Dependency is already built, add it.
143+
dependencies.add(dependency);
144+
} else if (pendingProtos.containsKey(dependencyName)) {
145+
// Dependency is not yet built, we must wait.
146+
allDependenciesMet = false;
147+
break;
148+
} else {
149+
// Dependency is not in our set. We assume it's a well-known type (e.g.,
150+
// google/protobuf/timestamp.proto) that buildFrom() can find and link automatically.
151+
}
152+
}
153+
154+
if (allDependenciesMet) {
155+
try {
156+
// All dependencies are met, we can build this file.
157+
FileDescriptor fileDescriptor =
158+
FileDescriptor.buildFrom(proto, dependencies.toArray(new FileDescriptor[0]));
159+
160+
builtDescriptors.put(fileDescriptor.getName(), fileDescriptor);
161+
pendingProtos.remove(proto.getName());
162+
filesBuiltThisPass++;
163+
164+
// Now, populate both message and enum maps with all messages/enums in this file.
165+
for (EnumDescriptor enumDescriptor : fileDescriptor.getEnumTypes()) {
166+
enumDescriptorMap.put(enumDescriptor.getFullName(), enumDescriptor);
167+
}
168+
for (Descriptor messageDescriptor : fileDescriptor.getMessageTypes()) {
169+
populateDescriptorMapsRecursively(messageDescriptor);
170+
}
171+
} catch (DescriptorValidationException e) {
172+
throw new IllegalArgumentException(
173+
"Failed to build descriptor for " + proto.getName(), e);
174+
}
175+
}
176+
}
177+
} while (filesBuiltThisPass > 0);
178+
179+
// Throw if we finished looping but still have pending protos.
180+
if (!pendingProtos.isEmpty()) {
181+
throw new IllegalArgumentException(
182+
"Unresolvable or circular dependencies found for: " + pendingProtos.keySet());
183+
}
184+
}
185+
186+
public ExecuteQueryResult toExecuteQueryResult(ResultSet resultSet)
45187
throws ExecutionException, InterruptedException {
46188
ExecuteQueryResult.Builder resultBuilder = ExecuteQueryResult.newBuilder();
47189
for (ColumnMetadata columnMetadata : resultSet.getMetadata().getColumns()) {
@@ -64,24 +206,28 @@ public static ExecuteQueryResult toExecuteQueryResult(ResultSet resultSet)
64206
return resultBuilder.build();
65207
}
66208

67-
private static Value toProtoValue(Object value, SqlType<?> type) {
209+
private Value toProtoValue(Object value, SqlType<?> type) {
68210
if (value == null) {
69211
return Value.getDefaultInstance();
70212
}
71213

72214
Value.Builder valueBuilder = Value.newBuilder();
73215
switch (type.getCode()) {
74216
case BYTES:
75-
case PROTO:
76217
valueBuilder.setBytesValue((ByteString) value);
77218
break;
219+
case PROTO:
220+
valueBuilder.setBytesValue(((AbstractMessage) value).toByteString());
221+
break;
78222
case STRING:
79223
valueBuilder.setStringValue((String) value);
80224
break;
81225
case INT64:
82-
case ENUM:
83226
valueBuilder.setIntValue((Long) value);
84227
break;
228+
case ENUM:
229+
valueBuilder.setIntValue(((ProtocolMessageEnum) value).getNumber());
230+
break;
85231
case FLOAT32:
86232
valueBuilder.setFloatValue((Float) value);
87233
break;
@@ -151,7 +297,7 @@ private static Value toProtoValue(Object value, SqlType<?> type) {
151297
return valueBuilder.build();
152298
}
153299

154-
private static Object getColumn(StructReader struct, int fieldIndex, SqlType<?> fieldType) {
300+
private Object getColumn(StructReader struct, int fieldIndex, SqlType<?> fieldType) {
155301
if (struct.isNull(fieldIndex)) {
156302
return null;
157303
}
@@ -162,17 +308,35 @@ private static Object getColumn(StructReader struct, int fieldIndex, SqlType<?>
162308
case BOOL:
163309
return struct.getBoolean(fieldIndex);
164310
case BYTES:
165-
case PROTO:
166311
return struct.getBytes(fieldIndex);
312+
case PROTO:
313+
SchemalessProto protoType = (SchemalessProto) fieldType;
314+
Descriptor descriptor = messageDescriptorMap.get(protoType.getMessageName());
315+
if (descriptor == null) {
316+
throw new IllegalArgumentException(
317+
"Descriptor for message " + protoType.getMessageName() + " could not be found");
318+
}
319+
return struct.getProtoMessage(fieldIndex, DynamicMessage.getDefaultInstance(descriptor));
167320
case DATE:
168321
return struct.getDate(fieldIndex);
169322
case FLOAT32:
170323
return struct.getFloat(fieldIndex);
171324
case FLOAT64:
172325
return struct.getDouble(fieldIndex);
173326
case INT64:
174-
case ENUM:
175327
return struct.getLong(fieldIndex);
328+
case ENUM:
329+
SchemalessEnum enumType = (SchemalessEnum) fieldType;
330+
EnumDescriptor enumDescriptor = enumDescriptorMap.get(enumType.getEnumName());
331+
if (enumDescriptor == null) {
332+
throw new IllegalArgumentException(
333+
"Descriptor for enum " + enumType.getEnumName() + " could not be found");
334+
}
335+
// We need to extract the integer value of the enum. `getProtoEnum` is the only
336+
// available method, but it is designed for static enum types. To work around this,
337+
// we can pass a lambda that constructs our DummyEnum with the captured integer value
338+
// and the descriptor from the outer scope.
339+
return struct.getProtoEnum(fieldIndex, number -> new DummyEnum(number, enumDescriptor));
176340
case MAP:
177341
return struct.getMap(fieldIndex, (SqlType.Map<?, ?>) fieldType);
178342
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)