[ggj][codegen] feat: integrate batching with retry settings parsing

src/main/java/com/google/api/generator/gapic/model/GapicBatchingSettings.java

@@ -45,6 +45,12 @@ public enum FlowControlLimitExceededBehavior {
 
   public abstract FlowControlLimitExceededBehavior flowControlLimitExceededBehavior();
 
+  public boolean matches(Service service, Method method) {
+    return protoPakkage().equals(service.protoPakkage())
+        && serviceName().equals(service.name())
+        && methodName().equals(method.name());
+  }
+
   public static Builder builder() {
     return new AutoValue_GapicBatchingSettings.Builder()
         .setFlowControlLimitExceededBehavior(FlowControlLimitExceededBehavior.IGNORE);

src/main/java/com/google/api/generator/gapic/model/GapicServiceConfig.java

@@ -42,14 +42,19 @@ public class GapicServiceConfig {
 
   private final List<MethodConfig> methodConfigs;
   private final Map<MethodConfig.Name, Integer> methodConfigTable;
+  private final Map<MethodConfig.Name, GapicBatchingSettings> batchingSettingsTable;
 
   private GapicServiceConfig(
-      List<MethodConfig> methodConfigs, Map<MethodConfig.Name, Integer> methodConfigTable) {
+      List<MethodConfig> methodConfigs,
+      Map<MethodConfig.Name, Integer> methodConfigTable,
+      Map<MethodConfig.Name, GapicBatchingSettings> batchingSettingsTable) {
     this.methodConfigs = methodConfigs;
     this.methodConfigTable = methodConfigTable;
+    this.batchingSettingsTable = batchingSettingsTable;
   }
 
-  public static GapicServiceConfig create(ServiceConfig serviceConfig) {
+  public static GapicServiceConfig create(
+      ServiceConfig serviceConfig, Optional<List<GapicBatchingSettings>> batchingSettingsOpt) {
     // Keep this  processing logic out of the constructor.
     Map<MethodConfig.Name, Integer> methodConfigTable = new HashMap<>();
     List<MethodConfig> methodConfigs = serviceConfig.getMethodConfigList();
@@ -60,36 +65,29 @@ public static GapicServiceConfig create(ServiceConfig serviceConfig) {
       }
     }
 
-    return new GapicServiceConfig(methodConfigs, methodConfigTable);
+    Map<MethodConfig.Name, GapicBatchingSettings> batchingSettingsTable = new HashMap<>();
+    if (batchingSettingsOpt.isPresent()) {
+      for (GapicBatchingSettings batchingSetting : batchingSettingsOpt.get()) {
+        batchingSettingsTable.put(
+            MethodConfig.Name.newBuilder()
+                .setService(
+                    String.format(
+                        "%s.%s", batchingSetting.protoPakkage(), batchingSetting.serviceName()))
+                .setMethod(batchingSetting.methodName())
+                .build(),
+            batchingSetting);
+      }
+    }
+
+    return new GapicServiceConfig(methodConfigs, methodConfigTable, batchingSettingsTable);
   }
 
   public Map<String, GapicRetrySettings> getAllGapicRetrySettings(Service service) {
     return service.methods().stream()
         .collect(
             Collectors.toMap(
                 m -> getRetryParamsName(service, m),
-                m -> {
-                  GapicRetrySettings.Kind kind = GapicRetrySettings.Kind.FULL;
-                  Optional<Integer> retryPolicyIndexOpt = retryPolicyIndexLookup(service, m);
-                  if (!retryPolicyIndexOpt.isPresent()) {
-                    kind = GapicRetrySettings.Kind.NONE;
-                  } else {
-                    MethodConfig methodConfig = methodConfigs.get(retryPolicyIndexOpt.get());
-                    if (!methodConfig.hasTimeout() && !methodConfig.hasRetryPolicy()) {
-                      kind = GapicRetrySettings.Kind.NONE;
-                    } else {
-                      kind =
-                          methodConfig.hasRetryPolicy()
-                              ? GapicRetrySettings.Kind.FULL
-                              : GapicRetrySettings.Kind.NO_RETRY;
-                    }
-                  }
-                  return GapicRetrySettings.builder()
-                      .setTimeout(timeoutLookup(service, m))
-                      .setRetryPolicy(retryPolicyLookup(service, m))
-                      .setKind(kind)
-                      .build();
-                },
+                m -> toGapicRetrySettings(service, m),
                 (r1, r2) -> r2,
                 LinkedHashMap::new));
   }
@@ -120,6 +118,40 @@ public String getRetryParamsName(Service service, Method method) {
     return NO_RETRY_PARAMS_NAME;
   }
 
+  public boolean hasBatchingSetting(Service service, Method method) {
+    return batchingSettingsTable.containsKey(toName(service, method));
+  }
+
+  public Optional<GapicBatchingSettings> getBatchingSetting(Service service, Method method) {
+    return hasBatchingSetting(service, method)
+        ? Optional.of(batchingSettingsTable.get(toName(service, method)))
+        : Optional.empty();
+  }
+
+  private GapicRetrySettings toGapicRetrySettings(Service service, Method method) {
+    GapicRetrySettings.Kind kind = GapicRetrySettings.Kind.FULL;
+    Optional<Integer> retryPolicyIndexOpt = retryPolicyIndexLookup(service, method);
+    if (!retryPolicyIndexOpt.isPresent()) {
+      kind = GapicRetrySettings.Kind.NONE;
+    } else {
+      MethodConfig methodConfig = methodConfigs.get(retryPolicyIndexOpt.get());
+      if (!methodConfig.hasTimeout() && !methodConfig.hasRetryPolicy()) {
+        kind = GapicRetrySettings.Kind.NONE;
+      } else {
+        kind =
+            methodConfig.hasRetryPolicy()
+                ? GapicRetrySettings.Kind.FULL
+                : GapicRetrySettings.Kind.NO_RETRY;
+      }
+    }
+
+    return GapicRetrySettings.builder()
+        .setTimeout(timeoutLookup(service, method))
+        .setRetryPolicy(retryPolicyLookup(service, method))
+        .setKind(kind)
+        .build();
+  }
+
   private List<Code> retryCodesLookup(Service service, Method method) {
     RetryPolicy retryPolicy = retryPolicyLookup(service, method);
     if (retryPolicy.equals(EMPTY_RETRY_POLICY)) {

src/main/java/com/google/api/generator/gapic/protoparser/BatchingSettingsConfigParser.java

@@ -45,7 +45,14 @@ public class BatchingSettingsConfigParser {
   private static String YAML_KEY_BATCHING_FLOW_CONTROL_LIMIT_EXCEEDED_BEHAVIOR =
       "flow_control_limit_exceeded_behavior";
 
-  public static Optional<List<GapicBatchingSettings>> parse(String gapicYamlConfigFilePath) {
+  public static Optional<List<GapicBatchingSettings>> parse(
+      Optional<String> gapicYamlConfigFilePathOpt) {
+    return gapicYamlConfigFilePathOpt.isPresent()
+        ? parse(gapicYamlConfigFilePathOpt.get())
+        : Optional.empty();
+  }
+
+  static Optional<List<GapicBatchingSettings>> parse(String gapicYamlConfigFilePath) {
     if (Strings.isNullOrEmpty(gapicYamlConfigFilePath)
         || !(new File(gapicYamlConfigFilePath)).exists()) {
       return Optional.empty();

src/main/java/com/google/api/generator/gapic/protoparser/Parser.java

@@ -20,6 +20,7 @@
 import com.google.api.generator.engine.ast.TypeNode;
 import com.google.api.generator.engine.ast.VaporReference;
 import com.google.api.generator.gapic.model.Field;
+import com.google.api.generator.gapic.model.GapicBatchingSettings;
 import com.google.api.generator.gapic.model.GapicContext;
 import com.google.api.generator.gapic.model.GapicServiceConfig;
 import com.google.api.generator.gapic.model.LongrunningOperation;
@@ -70,15 +71,15 @@ public GapicParserException(String errorMessage) {
   }
 
   public static GapicContext parse(CodeGeneratorRequest request) {
-    Optional<String> serviceConfigPathOpt = PluginArgumentParser.parseJsonConfigPath(request);
-    String serviceConfigPath = serviceConfigPathOpt.isPresent() ? serviceConfigPathOpt.get() : null;
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(serviceConfigPath);
-
-    // TODO(miraleung): Actually pars the yaml file.
     Optional<String> gapicYamlConfigPathOpt =
         PluginArgumentParser.parseGapicYamlConfigPath(request);
-    String gapicYamlConfigPath =
-        gapicYamlConfigPathOpt.isPresent() ? gapicYamlConfigPathOpt.get() : null;
+    Optional<List<GapicBatchingSettings>> batchingSettingsOpt =
+        BatchingSettingsConfigParser.parse(gapicYamlConfigPathOpt);
+
+    Optional<String> serviceConfigPathOpt = PluginArgumentParser.parseJsonConfigPath(request);
+    String serviceConfigPath = serviceConfigPathOpt.isPresent() ? serviceConfigPathOpt.get() : null;
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(serviceConfigPath, batchingSettingsOpt);
 
     // Keep message and resource name parsing separate for cleaner logic.
     // While this takes an extra pass through the protobufs, the extra time is relatively trivial

src/main/java/com/google/api/generator/gapic/protoparser/ServiceConfigParser.java

@@ -14,6 +14,7 @@
 
 package com.google.api.generator.gapic.protoparser;
 
+import com.google.api.generator.gapic.model.GapicBatchingSettings;
 import com.google.api.generator.gapic.model.GapicServiceConfig;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
@@ -22,15 +23,17 @@
 import java.io.FileNotFoundException;
 import java.io.FileReader;
 import java.io.IOException;
+import java.util.List;
 import java.util.Optional;
 
 public class ServiceConfigParser {
-  public static Optional<GapicServiceConfig> parse(String serviceConfigFilePath) {
+  public static Optional<GapicServiceConfig> parse(
+      String serviceConfigFilePath, Optional<List<GapicBatchingSettings>> batchingSettingsOpt) {
     Optional<ServiceConfig> rawConfigOpt = parseFile(serviceConfigFilePath);
     if (!rawConfigOpt.isPresent()) {
       return Optional.empty();
     }
-    return Optional.of(GapicServiceConfig.create(rawConfigOpt.get()));
+    return Optional.of(GapicServiceConfig.create(rawConfigOpt.get(), batchingSettingsOpt));
   }
 
   @VisibleForTesting

src/test/java/com/google/api/generator/gapic/composer/RetrySettingsComposerTest.java

@@ -81,7 +81,8 @@ public void paramDefinitionsBlock_noConfigsFound() {
 
     String jsonFilename = "retrying_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -117,7 +118,8 @@ public void paramDefinitionsBlock_basic() {
 
     String jsonFilename = "showcase_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -157,7 +159,8 @@ public void codesDefinitionsBlock_noConfigsFound() {
 
     String jsonFilename = "retrying_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -193,7 +196,8 @@ public void codesDefinitionsBlock_basic() {
 
     String jsonFilename = "showcase_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -232,7 +236,8 @@ public void simpleBuilderExpr_basic() {
 
     String jsonFilename = "showcase_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -313,7 +318,8 @@ public void lroBuilderExpr() {
 
     String jsonFilename = "showcase_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 

src/test/java/com/google/api/generator/gapic/composer/ServiceStubSettingsClassComposerTest.java

@@ -62,7 +62,8 @@ public void generateServiceClasses() {
 
     String jsonFilename = "showcase_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> configOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<GapicServiceConfig> configOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), Optional.empty());
     assertTrue(configOpt.isPresent());
     GapicServiceConfig config = configOpt.get();
 

src/test/java/com/google/api/generator/gapic/model/GapicServiceConfigTest.java

@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
 import com.google.api.generator.gapic.protoparser.Parser;
@@ -28,6 +29,7 @@
 import io.grpc.serviceconfig.MethodConfig;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -48,7 +50,9 @@ public void serviceConfig_noConfigsFound() {
 
     String jsonFilename = "retrying_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<List<GapicBatchingSettings>> batchingSettingsOpt = Optional.empty();
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), batchingSettingsOpt);
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -77,7 +81,9 @@ public void serviceConfig_basic() {
 
     String jsonFilename = "showcase_grpc_service_config.json";
     Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
-    Optional<GapicServiceConfig> serviceConfigOpt = ServiceConfigParser.parse(jsonPath.toString());
+    Optional<List<GapicBatchingSettings>> batchingSettingsOpt = Optional.empty();
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), batchingSettingsOpt);
     assertTrue(serviceConfigOpt.isPresent());
     GapicServiceConfig serviceConfig = serviceConfigOpt.get();
 
@@ -124,6 +130,78 @@ public void serviceConfig_basic() {
     assertEquals(GapicServiceConfig.EMPTY_RETRY_CODES, retryCodes.get(retryCodeName));
   }
 
+  @Test
+  public void serviceConfig_withBatchingSettings() {
+    FileDescriptor echoFileDescriptor = EchoOuterClass.getDescriptor();
+    ServiceDescriptor echoServiceDescriptor = echoFileDescriptor.getServices().get(0);
+    Service service = parseService(echoFileDescriptor);
+
+    String jsonFilename = "showcase_grpc_service_config.json";
+    Path jsonPath = Paths.get(JSON_DIRECTORY, jsonFilename);
+
+    GapicBatchingSettings origBatchingSetting =
+        GapicBatchingSettings.builder()
+            .setProtoPakkage("google.showcase.v1beta1")
+            .setServiceName("Echo")
+            .setMethodName("Echo")
+            .setElementCountThreshold(1000)
+            .setRequestByteThreshold(2000)
+            .setDelayThresholdMillis(3000)
+            .build();
+    Optional<List<GapicBatchingSettings>> batchingSettingsOpt =
+        Optional.of(Arrays.asList(origBatchingSetting));
+
+    Optional<GapicServiceConfig> serviceConfigOpt =
+        ServiceConfigParser.parse(jsonPath.toString(), batchingSettingsOpt);
+    assertTrue(serviceConfigOpt.isPresent());
+    GapicServiceConfig serviceConfig = serviceConfigOpt.get();
+
+    Map<String, GapicRetrySettings> retrySettings = serviceConfig.getAllGapicRetrySettings(service);
+    assertEquals(2, retrySettings.size());
+    Map<String, List<Code>> retryCodes = serviceConfig.getAllRetryCodes(service);
+    assertEquals(2, retryCodes.size());
+
+    // Echo method has an explicitly-defined setting.
+    Method method = findMethod(service, "Echo");
+    assertThat(method).isNotNull();
+
+    // No change to the retry settings.
+    String retryParamsName = serviceConfig.getRetryParamsName(service, method);
+    assertEquals("retry_policy_1_params", retryParamsName);
+    GapicRetrySettings settings = retrySettings.get(retryParamsName);
+    assertThat(settings).isNotNull();
+    assertEquals(10, settings.timeout().getSeconds());
+    assertEquals(GapicRetrySettings.Kind.FULL, settings.kind());
+
+    // No changge to the retry codes.
+    String retryCodeName = serviceConfig.getRetryCodeName(service, method);
+    assertEquals("retry_policy_1_codes", retryCodeName);
+    List<Code> retryCode = retryCodes.get(retryCodeName);
+    assertThat(retryCode).containsExactly(Code.UNAVAILABLE, Code.UNKNOWN);
+
+    // Check batching settings.
+    assertTrue(serviceConfig.hasBatchingSetting(service, method));
+    Optional<GapicBatchingSettings> batchingSettingOpt =
+        serviceConfig.getBatchingSetting(service, method);
+    assertTrue(batchingSettingOpt.isPresent());
+    GapicBatchingSettings batchingSetting = batchingSettingOpt.get();
+    assertEquals(
+        origBatchingSetting.elementCountThreshold(), batchingSetting.elementCountThreshold());
+    assertEquals(
+        origBatchingSetting.requestByteThreshold(), batchingSetting.requestByteThreshold());
+    assertEquals(
+        origBatchingSetting.delayThresholdMillis(), batchingSetting.delayThresholdMillis());
+
+    // Chat method defaults to the service-defined setting.
+    method = findMethod(service, "Chat");
+    assertThat(method).isNotNull();
+    retryParamsName = serviceConfig.getRetryParamsName(service, method);
+    assertEquals("no_retry_0_params", retryParamsName);
+    retryCodeName = serviceConfig.getRetryCodeName(service, method);
+    assertEquals("no_retry_0_codes", retryCodeName);
+    assertFalse(serviceConfig.hasBatchingSetting(service, method));
+  }
+
   private static Service parseService(FileDescriptor fileDescriptor) {
     Map<String, Message> messageTypes = Parser.parseMessages(fileDescriptor);
     Map<String, ResourceName> resourceNames = Parser.parseResourceNames(fileDescriptor);