feat: parse batching descriptor fields (#258)

googleapis · Sep 2, 2020 · d146456 · d146456
1 parent e9de699
commit d146456
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 0 deletions.
diff --git a/src/main/java/com/google/api/generator/gapic/model/GapicBatchingSettings.java b/src/main/java/com/google/api/generator/gapic/model/GapicBatchingSettings.java
@@ -15,6 +15,8 @@
 package com.google.api.generator.gapic.model;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
+import java.util.List;
 import javax.annotation.Nullable;
 
 @AutoValue
@@ -25,6 +27,7 @@ public enum FlowControlLimitExceededBehavior {
     IGNORE
   };
 
+  // Threshold fields.
   public abstract String protoPakkage();
 
   public abstract String serviceName();
@@ -45,6 +48,14 @@ public enum FlowControlLimitExceededBehavior {
 
   public abstract FlowControlLimitExceededBehavior flowControlLimitExceededBehavior();
 
+  // Batch descriptor fields.
+  public abstract String batchedFieldName();
+
+  public abstract ImmutableList<String> discriminatorFieldNames();
+
+  @Nullable
+  public abstract String subresponseFieldName();
+
   public boolean matches(Service service, Method method) {
     return protoPakkage().equals(service.protoPakkage())
         && serviceName().equals(service.name())
@@ -77,6 +88,12 @@ public abstract static class Builder {
     public abstract Builder setFlowControlLimitExceededBehavior(
         FlowControlLimitExceededBehavior behavior);
 
+    public abstract Builder setBatchedFieldName(String batchedFieldName);
+
+    public abstract Builder setDiscriminatorFieldNames(List<String> discriminatorFieldNames);
+
+    public abstract Builder setSubresponseFieldName(String subresponseFieldName);
+
     public abstract GapicBatchingSettings build();
   }
 }
diff --git a/src/main/java/com/google/api/generator/gapic/protoparser/BatchingSettingsConfigParser.java b/src/main/java/com/google/api/generator/gapic/protoparser/BatchingSettingsConfigParser.java
@@ -38,6 +38,8 @@ public class BatchingSettingsConfigParser {
   private static String YAML_KEY_METHODS = "methods";
   private static String YAML_KEY_BATCHING = "batching";
   private static String YAML_KEY_THRESHOLDS = "thresholds";
+  private static String YAML_KEY_DESCRIPTOR = "batch_descriptor";
+
   private static String YAML_KEY_BATCHING_ELEMENT_COUNT_THRESHOLD = "element_count_threshold";
   private static String YAML_KEY_BATCHING_DELAY_THRESHOLD_MILLIS = "delay_threshold_millis";
   private static String YAML_KEY_BATCHING_REQUEST_BYTE_THRESHOLD = "request_byte_threshold";
@@ -46,6 +48,10 @@ public class BatchingSettingsConfigParser {
   private static String YAML_KEY_BATCHING_FLOW_CONTROL_LIMIT_EXCEEDED_BEHAVIOR =
       "flow_control_limit_exceeded_behavior";
 
+  private static String YAML_KEY_DESCRIPTOR_BATCHED_FIELD = "batched_field";
+  private static String YAML_KEY_DESCRIPTOR_DISCRIMINATOR_FIELD = "discriminator_fields";
+  private static String YAML_KEY_DESCRIPTOR_SUBRESPONSE_FIELD = "subresponse_field";
+
   public static Optional<List<GapicBatchingSettings>> parse(
       Optional<String> gapicYamlConfigFilePathOpt) {
     return gapicYamlConfigFilePathOpt.isPresent()
@@ -94,6 +100,13 @@ private static Optional<List<GapicBatchingSettings>> parseFromMap(Map<String, Ob
         if (!batchingOuterYamlConfig.containsKey(YAML_KEY_THRESHOLDS)) {
           continue;
         }
+        Preconditions.checkState(
+            batchingOuterYamlConfig.containsKey(YAML_KEY_DESCRIPTOR),
+            String.format(
+                "%s key expected but not found for method %s",
+                YAML_KEY_DESCRIPTOR, (String) methodYamlConfig.get(YAML_KEY_NAME)));
+
+        // Parse the threshold values first.
         Map<String, Object> batchingYamlConfig =
             (Map<String, Object>) batchingOuterYamlConfig.get(YAML_KEY_THRESHOLDS);
         Preconditions.checkState(
@@ -147,6 +160,27 @@ private static Optional<List<GapicBatchingSettings>> parseFromMap(Map<String, Ob
           }
           settingsBuilder.setFlowControlLimitExceededBehavior(behaviorSetting);
         }
+
+        // Parse the descriptor values.
+        Map<String, Object> descriptorYamlConfig =
+            (Map<String, Object>) batchingOuterYamlConfig.get(YAML_KEY_DESCRIPTOR);
+        Preconditions.checkState(
+            descriptorYamlConfig.containsKey(YAML_KEY_DESCRIPTOR_BATCHED_FIELD)
+                && descriptorYamlConfig.containsKey(YAML_KEY_DESCRIPTOR_DISCRIMINATOR_FIELD),
+            String.format(
+                "Batching descriptor YAML config is missing one of %s or %s fields",
+                YAML_KEY_DESCRIPTOR_BATCHED_FIELD, YAML_KEY_DESCRIPTOR_DISCRIMINATOR_FIELD));
+
+        settingsBuilder.setBatchedFieldName(
+            (String) descriptorYamlConfig.get(YAML_KEY_DESCRIPTOR_BATCHED_FIELD));
+        settingsBuilder.setDiscriminatorFieldNames(
+            (List<String>) descriptorYamlConfig.get(YAML_KEY_DESCRIPTOR_DISCRIMINATOR_FIELD));
+
+        if (descriptorYamlConfig.containsKey(YAML_KEY_DESCRIPTOR_SUBRESPONSE_FIELD)) {
+          settingsBuilder.setSubresponseFieldName(
+              (String) descriptorYamlConfig.get(YAML_KEY_DESCRIPTOR_SUBRESPONSE_FIELD));
+        }
+
         settings.add(settingsBuilder.build());
       }
     }

diff --git a/src/test/java/com/google/api/generator/gapic/composer/RetrySettingsComposerTest.java b/src/test/java/com/google/api/generator/gapic/composer/RetrySettingsComposerTest.java
@@ -409,6 +409,9 @@ public void batchingSettings_minimalFlowControlSettings() {
             .setElementCountThreshold(100)
             .setRequestByteThreshold(1048576)
             .setDelayThresholdMillis(10)
+            .setBatchedFieldName("messages")
+            .setDiscriminatorFieldNames(Arrays.asList("topic"))
+            .setSubresponseFieldName("message_ids")
             .build();
 
     Expr builderExpr =
@@ -487,6 +490,8 @@ public void batchingSettings_fullFlowControlSettings() {
             .setFlowControlByteLimit(10485760)
             .setFlowControlLimitExceededBehavior(
                 GapicBatchingSettings.FlowControlLimitExceededBehavior.THROW_EXCEPTION)
+            .setBatchedFieldName("entries")
+            .setDiscriminatorFieldNames(Arrays.asList("log_name", "resource", "labels"))
             .build();
 
     Expr builderExpr =

diff --git a/src/test/java/com/google/api/generator/gapic/model/GapicServiceConfigTest.java b/src/test/java/com/google/api/generator/gapic/model/GapicServiceConfigTest.java
@@ -147,6 +147,8 @@ public void serviceConfig_withBatchingSettings() {
             .setElementCountThreshold(1000)
             .setRequestByteThreshold(2000)
             .setDelayThresholdMillis(3000)
+            .setBatchedFieldName("name")
+            .setDiscriminatorFieldNames(Arrays.asList("severity"))
             .build();
     Optional<List<GapicBatchingSettings>> batchingSettingsOpt =
         Optional.of(Arrays.asList(origBatchingSetting));

diff --git a/...est/java/com/google/api/generator/gapic/protoparser/BatchingSettingsConfigParserTest.java b/...est/java/com/google/api/generator/gapic/protoparser/BatchingSettingsConfigParserTest.java
@@ -75,6 +75,10 @@ public void parseBatchingSettings_logging() {
     assertEquals(
         GapicBatchingSettings.FlowControlLimitExceededBehavior.THROW_EXCEPTION,
         setting.flowControlLimitExceededBehavior());
+
+    assertEquals("entries", setting.batchedFieldName());
+    assertThat(setting.discriminatorFieldNames()).containsExactly("log_name", "resource", "labels");
+    assertThat(setting.subresponseFieldName()).isNull();
   }
 
   @Test
@@ -102,5 +106,9 @@ public void parseBatchingSettings_pubsub() {
     assertEquals(
         GapicBatchingSettings.FlowControlLimitExceededBehavior.IGNORE,
         setting.flowControlLimitExceededBehavior());
+
+    assertEquals("messages", setting.batchedFieldName());
+    assertThat(setting.discriminatorFieldNames()).containsExactly("topic");
+    assertEquals("message_ids", setting.subresponseFieldName());
   }
 }