Skip to content

Commit db62a44

Browse files
miraleungmiraleung
authored
[ggj][codegen] feat: add initial batching descriptor field to ServiceStubSettings (#262)
* feat: parse batching descriptor fields * feat: add Field.isMap and map-parsing test * feat: add initial batching descriptor field to ServiceStubSettings
1 parent 9dc2038 commit db62a44

5 files changed

Lines changed: 550 additions & 0 deletions

File tree

‎src/main/java/com/google/api/generator/gapic/composer/BatchingDescriptorComposer.java‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
// Copyright 2020 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package com.google.api.generator.gapic.composer;
16+
17+
import com.google.api.gax.batching.PartitionKey;
18+
import com.google.api.gax.batching.RequestBuilder;
19+
import com.google.api.gax.rpc.BatchedRequestIssuer;
20+
import com.google.api.gax.rpc.BatchingDescriptor;
21+
import com.google.api.generator.engine.ast.AnonymousClassExpr;
22+
import com.google.api.generator.engine.ast.AssignmentExpr;
23+
import com.google.api.generator.engine.ast.ConcreteReference;
24+
import com.google.api.generator.engine.ast.Expr;
25+
import com.google.api.generator.engine.ast.ExprStatement;
26+
import com.google.api.generator.engine.ast.IfStatement;
27+
import com.google.api.generator.engine.ast.MethodDefinition;
28+
import com.google.api.generator.engine.ast.MethodInvocationExpr;
29+
import com.google.api.generator.engine.ast.NewObjectExpr;
30+
import com.google.api.generator.engine.ast.Reference;
31+
import com.google.api.generator.engine.ast.ScopeNode;
32+
import com.google.api.generator.engine.ast.TypeNode;
33+
import com.google.api.generator.engine.ast.Variable;
34+
import com.google.api.generator.engine.ast.VariableExpr;
35+
import com.google.api.generator.gapic.model.GapicBatchingSettings;
36+
import com.google.api.generator.gapic.model.Message;
37+
import com.google.api.generator.gapic.model.Method;
38+
import com.google.api.generator.gapic.utils.JavaStyle;
39+
import com.google.common.base.Preconditions;
40+
import java.util.ArrayList;
41+
import java.util.Arrays;
42+
import java.util.List;
43+
import java.util.Map;
44+
import java.util.Objects;
45+
import java.util.stream.Collectors;
46+
47+
public class BatchingDescriptorComposer {
48+
private static final String BATCHING_DESC_PATTERN = "%s_BATCHING_DESC";
49+
50+
private static final Reference BATCHING_DESCRIPTOR_REF =
51+
ConcreteReference.withClazz(BatchingDescriptor.class);
52+
private static final Reference REQUEST_BUILDER_REF =
53+
ConcreteReference.withClazz(RequestBuilder.class);
54+
private static final Reference BATCHED_REQUEST_ISSUER_REF =
55+
ConcreteReference.withClazz(BatchedRequestIssuer.class);
56+
57+
private static final TypeNode PARTITION_KEY_TYPE = toType(PartitionKey.class);
58+
59+
private static final String ADD_ALL_METHOD_PATTERN = "addAll%s";
60+
private static final String GET_LIST_METHOD_PATTERN = "get%sList";
61+
62+
public static Expr createBatchingDescriptorFieldDeclExpr(
63+
Method method, GapicBatchingSettings batchingSettings, Map<String, Message> messageTypes) {
64+
List<MethodDefinition> javaMethods = new ArrayList<>();
65+
javaMethods.add(createGetBatchPartitionKeyMethod(method, batchingSettings, messageTypes));
66+
javaMethods.add(createGetRequestBuilderMethod(method, batchingSettings));
67+
68+
TypeNode batchingDescriptorType =
69+
toType(BATCHING_DESCRIPTOR_REF, method.inputType(), method.outputType());
70+
AnonymousClassExpr batchingDescriptorClassExpr =
71+
AnonymousClassExpr.builder()
72+
.setType(batchingDescriptorType)
73+
.setMethods(javaMethods)
74+
.build();
75+
76+
String varName =
77+
String.format(BATCHING_DESC_PATTERN, JavaStyle.toUpperSnakeCase(method.name()));
78+
return AssignmentExpr.builder()
79+
.setVariableExpr(
80+
VariableExpr.builder()
81+
.setVariable(
82+
Variable.builder().setType(batchingDescriptorType).setName(varName).build())
83+
.setIsDecl(true)
84+
.setScope(ScopeNode.PRIVATE)
85+
.setIsStatic(true)
86+
.setIsFinal(true)
87+
.build())
88+
.setValueExpr(batchingDescriptorClassExpr)
89+
.build();
90+
}
91+
92+
private static MethodDefinition createGetBatchPartitionKeyMethod(
93+
Method method, GapicBatchingSettings batchingSettings, Map<String, Message> messageTypes) {
94+
String methodInputTypeName = method.inputType().reference().name();
95+
Message inputMessage = messageTypes.get(methodInputTypeName);
96+
Preconditions.checkNotNull(
97+
inputMessage,
98+
String.format(
99+
"Message %s not found for RPC method %s", methodInputTypeName, method.name()));
100+
101+
VariableExpr requestVarExpr =
102+
VariableExpr.withVariable(
103+
Variable.builder().setType(method.inputType()).setName("request").build());
104+
105+
List<Expr> partitionKeyArgExprs = new ArrayList<>();
106+
for (String discriminatorFieldName : batchingSettings.discriminatorFieldNames()) {
107+
Preconditions.checkNotNull(
108+
inputMessage.fieldMap().get(discriminatorFieldName),
109+
String.format(
110+
"Batching discriminator field %s not found in message %s",
111+
discriminatorFieldName, inputMessage.name()));
112+
String getterMethodName =
113+
String.format("get%s", JavaStyle.toUpperCamelCase(discriminatorFieldName));
114+
partitionKeyArgExprs.add(
115+
MethodInvocationExpr.builder()
116+
.setExprReferenceExpr(requestVarExpr)
117+
.setMethodName(getterMethodName)
118+
.build());
119+
}
120+
Expr returnExpr =
121+
NewObjectExpr.builder()
122+
.setType(PARTITION_KEY_TYPE)
123+
.setArguments(partitionKeyArgExprs)
124+
.build();
125+
126+
return MethodDefinition.builder()
127+
.setIsOverride(true)
128+
.setScope(ScopeNode.PUBLIC)
129+
.setReturnType(PARTITION_KEY_TYPE)
130+
.setName("getBatchPartitionKey")
131+
.setArguments(requestVarExpr.toBuilder().setIsDecl(true).build())
132+
.setReturnExpr(returnExpr)
133+
.build();
134+
}
135+
136+
private static MethodDefinition createGetRequestBuilderMethod(
137+
Method method, GapicBatchingSettings batchingSettings) {
138+
TypeNode builderType = toType(REQUEST_BUILDER_REF, method.inputType());
139+
VariableExpr builderVarExpr =
140+
VariableExpr.withVariable(
141+
Variable.builder().setType(builderType).setName("builder").build());
142+
VariableExpr requestVarExpr =
143+
VariableExpr.withVariable(
144+
Variable.builder().setType(method.inputType()).setName("request").build());
145+
146+
Expr toBuilderExpr =
147+
AssignmentExpr.builder()
148+
.setVariableExpr(builderVarExpr)
149+
.setValueExpr(
150+
MethodInvocationExpr.builder()
151+
.setExprReferenceExpr(requestVarExpr)
152+
.setMethodName("toBuilder")
153+
.setReturnType(builderType)
154+
.build())
155+
.build();
156+
157+
String upperBatchedFieldName = JavaStyle.toUpperCamelCase(batchingSettings.batchedFieldName());
158+
String getFooListMethodName = String.format(GET_LIST_METHOD_PATTERN, upperBatchedFieldName);
159+
Expr getFooListExpr =
160+
MethodInvocationExpr.builder()
161+
.setExprReferenceExpr(requestVarExpr)
162+
.setMethodName(getFooListMethodName)
163+
.build();
164+
165+
String addAllMethodName = String.format(ADD_ALL_METHOD_PATTERN, upperBatchedFieldName);
166+
Expr addAllExpr =
167+
MethodInvocationExpr.builder()
168+
.setExprReferenceExpr(builderVarExpr)
169+
.setMethodName(addAllMethodName)
170+
.setArguments(getFooListExpr)
171+
.build();
172+
173+
MethodDefinition appendRequestMethod =
174+
MethodDefinition.builder()
175+
.setIsOverride(true)
176+
.setScope(ScopeNode.PUBLIC)
177+
.setReturnType(TypeNode.VOID)
178+
.setName("appendRequest")
179+
.setArguments(requestVarExpr.toBuilder().setIsDecl(true).build())
180+
.setBody(
181+
Arrays.asList(
182+
IfStatement.builder()
183+
.setConditionExpr(
184+
MethodInvocationExpr.builder()
185+
.setStaticReferenceType(toType(Objects.class))
186+
.setMethodName("isNull")
187+
.setArguments(builderVarExpr)
188+
.setReturnType(TypeNode.BOOLEAN)
189+
.build())
190+
.setBody(Arrays.asList(ExprStatement.withExpr(toBuilderExpr)))
191+
.setElseBody(Arrays.asList(ExprStatement.withExpr(addAllExpr)))
192+
.build()))
193+
.build();
194+
195+
MethodDefinition buildMethod =
196+
MethodDefinition.builder()
197+
.setIsOverride(true)
198+
.setScope(ScopeNode.PUBLIC)
199+
.setReturnType(method.inputType())
200+
.setName("build")
201+
.setReturnExpr(
202+
MethodInvocationExpr.builder()
203+
.setExprReferenceExpr(builderVarExpr)
204+
.setMethodName("build")
205+
.setReturnType(method.inputType())
206+
.build())
207+
.build();
208+
209+
AnonymousClassExpr requestBuilderAnonClassExpr =
210+
AnonymousClassExpr.builder()
211+
.setType(builderType)
212+
.setStatements(
213+
Arrays.asList(
214+
ExprStatement.withExpr(
215+
builderVarExpr
216+
.toBuilder()
217+
.setIsDecl(true)
218+
.setScope(ScopeNode.PRIVATE)
219+
.build())))
220+
.setMethods(Arrays.asList(appendRequestMethod, buildMethod))
221+
.build();
222+
223+
return MethodDefinition.builder()
224+
.setIsOverride(true)
225+
.setScope(ScopeNode.PUBLIC)
226+
.setReturnType(builderType)
227+
.setName("getRequestBuilder")
228+
.setReturnExpr(requestBuilderAnonClassExpr)
229+
.build();
230+
}
231+
232+
private static TypeNode toType(Class clazz) {
233+
return TypeNode.withReference(ConcreteReference.withClazz(clazz));
234+
}
235+
236+
private static TypeNode toType(Reference reference, TypeNode... types) {
237+
return TypeNode.withReference(
238+
reference.copyAndSetGenerics(
239+
Arrays.asList(types).stream().map(t -> t.reference()).collect(Collectors.toList())));
240+
}
241+
}

‎src/main/java/com/google/api/generator/gapic/composer/ServiceStubSettingsClassComposer.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,17 @@ private static List<Statement> createClassStatements(
283283

284284
memberVarExprs.addAll(
285285
createPagingStaticAssignExprs(service, serviceConfig, messageTypes, types));
286+
287+
for (Method method : service.methods()) {
288+
Optional<GapicBatchingSettings> batchingSettingOpt =
289+
serviceConfig.getBatchingSetting(service, method);
290+
if (batchingSettingOpt.isPresent()) {
291+
memberVarExprs.add(
292+
BatchingDescriptorComposer.createBatchingDescriptorFieldDeclExpr(
293+
method, batchingSettingOpt.get(), messageTypes));
294+
}
295+
}
296+
286297
return memberVarExprs.stream().map(e -> ExprStatement.withExpr(e)).collect(Collectors.toList());
287298
}
288299

‎src/test/java/com/google/api/generator/gapic/composer/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package(default_visibility = ["//visibility:public"])
22

33
TESTS = [
4+
"BatchingDescriptorComposerTest",
45
"ComposerTest",
56
"GrpcServiceCallableFactoryClassComposerTest",
67
"GrpcServiceStubClassComposerTest",

0 commit comments

Comments
 (0)