feat: handle auto pagination for BigQuery v2

gapic-generator-java/src/main/java/com/google/api/generator/engine/ast/TypeNode.java

@@ -74,6 +74,17 @@ public enum TypeKind {
       withReference(
           VaporReference.builder().setName("Value").setPakkage("com.google.protobuf").build());
 
+  public static final TypeNode UINT32VALUE =
+      withReference(
+          VaporReference.builder()
+              .setName("UInt32Value")
+              .setPakkage("com.google.protobuf")
+              .build());
+
+  public static final TypeNode INT32VALUE =
+      withReference(
+          VaporReference.builder().setName("Int32Value").setPakkage("com.google.protobuf").build());
+
   private static final Map<TypeNode, TypeNode> BOXED_TYPE_MAP = createBoxedTypeMap();
 
   public static final TypeNode VOID = builder().setTypeKind(TypeKind.VOID).build();

gapic-generator-java/src/main/java/com/google/api/generator/gapic/composer/common/AbstractServiceStubSettingsClassComposer.java

@@ -149,6 +149,20 @@ public abstract class AbstractServiceStubSettingsClassComposer implements ClassC
 
   private final TransportContext transportContext;
 
+  // Maps of BigQuery methods to pagination types.
+  private static final ImmutableMap<String, TypeNode> BIGQUERY_PAGINATE_MAX_RESULT_TYPES =
+      ImmutableMap.of(
+          "com.google.cloud.bigquery.v2.ListDatasets",
+          TypeNode.UINT32VALUE,
+          "com.google.cloud.bigquery.v2.ListJobs",
+          TypeNode.INT32VALUE,
+          "com.google.cloud.bigquery.v2.ListModels",
+          TypeNode.UINT32VALUE,
+          "com.google.cloud.bigquery.v2.ListRoutines",
+          TypeNode.UINT32VALUE,
+          "com.google.cloud.bigquery.v2.ListTables",
+          TypeNode.UINT32VALUE);
+
   protected static final VariableExpr DEFAULT_SERVICE_SCOPES_VAR_EXPR =
       createDefaultServiceScopesVarExpr();
 
@@ -727,9 +741,12 @@ private static Expr createPagedListDescriptorAssignExpr(
             .build());
 
     // Create injectPageSize method.
+    String methodFullName =
+        String.format("%s.%s", method.inputType().reference().pakkage(), method.name());
     VariableExpr pageSizeVarExpr =
         VariableExpr.withVariable(
             Variable.builder().setType(TypeNode.INT).setName("pageSize").build());
+
     // Re-declare for clarity and easier readability.
     returnType = method.inputType();
     returnExpr =
@@ -738,6 +755,20 @@ private static Expr createPagedListDescriptorAssignExpr(
             .setMethodName("set" + JavaStyle.toUpperCamelCase(method.pageSizeFieldName()))
             .setArguments(pageSizeVarExpr)
             .build();
+    if (BIGQUERY_PAGINATE_MAX_RESULT_TYPES.containsKey(methodFullName)) {
+      returnExpr =
+          MethodInvocationExpr.builder()
+              .setExprReferenceExpr(newBuilderExpr)
+              .setMethodName("set" + JavaStyle.toUpperCamelCase(method.pageSizeFieldName()))
+              .setArguments(
+                  MethodInvocationExpr.builder()
+                      .setStaticReferenceType(
+                          BIGQUERY_PAGINATE_MAX_RESULT_TYPES.get(methodFullName))
+                      .setMethodName("of")
+                      .setArguments(pageSizeVarExpr)
+                      .build())
+              .build();
+    }
     returnExpr =
         MethodInvocationExpr.builder()
             .setExprReferenceExpr(returnExpr)
@@ -758,17 +789,32 @@ private static Expr createPagedListDescriptorAssignExpr(
     // TODO(miraleung): Test the edge cases where these proto fields aren't present.
     // Create extractPageSize method.
     returnType = TypeNode.INT_OBJECT;
+    returnExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(payloadVarExpr)
+            .setMethodName("get" + JavaStyle.toUpperCamelCase(method.pageSizeFieldName()))
+            .setReturnType(returnType)
+            .build();
+    if (BIGQUERY_PAGINATE_MAX_RESULT_TYPES.containsKey(methodFullName)) {
+      // Return type is UINT32VALUE or INT32VALUE so use getValue to unwrap.
+      returnExpr =
+          MethodInvocationExpr.builder()
+              .setExprReferenceExpr(payloadVarExpr)
+              .setMethodName("get" + JavaStyle.toUpperCamelCase(method.pageSizeFieldName()))
+              .build();
+      returnExpr =
+          MethodInvocationExpr.builder()
+              .setExprReferenceExpr(returnExpr)
+              .setMethodName("getValue")
+              .setReturnType(returnType)
+              .build();
+    }
     anonClassMethods.add(
         methodStarterBuilder
             .setReturnType(returnType)
             .setName("extractPageSize")
             .setArguments(payloadVarExpr.toBuilder().setIsDecl(true).build())
-            .setReturnExpr(
-                MethodInvocationExpr.builder()
-                    .setExprReferenceExpr(payloadVarExpr)
-                    .setMethodName("get" + JavaStyle.toUpperCamelCase(method.pageSizeFieldName()))
-                    .setReturnType(returnType)
-                    .build())
+            .setReturnExpr(returnExpr)
             .build());
 
     // Create extractNextToken method.

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

@@ -121,6 +121,16 @@ enum SelectiveGapicType {
   private static final Set<String> MIXIN_JAVA_PACKAGE_ALLOWLIST =
       ImmutableSet.of("com.google.iam.v1", "com.google.longrunning", "com.google.cloud.location");
 
+  // List of BigQuery methods that can use max_result field as an alternative to page_size for
+  // pagination.
+  private static final ImmutableSet<String> BIGQUERY_LEGACY_PAGINATION_ALLOWLIST =
+      ImmutableSet.of(
+          "google.cloud.bigquery.v2.JobService.ListJobs",
+          "google.cloud.bigquery.v2.RoutineService.ListRoutines",
+          "google.cloud.bigquery.v2.DatasetService.ListDatasets",
+          "google.cloud.bigquery.v2.ModelService.ListModels",
+          "google.cloud.bigquery.v2.TableService.ListTables");
+
   // Allow other parsers to access this.
   protected static final SourceCodeInfoParser SOURCE_CODE_INFO_PARSER = new SourceCodeInfoParser();
 
@@ -1029,7 +1039,8 @@ static String parsePageSizeFieldName(
       // page_size gets priority over max_results if both are present
       List<String> fieldNames = new ArrayList<>();
       fieldNames.add("page_size");
-      if (transport == Transport.REST) {
+      if ((transport == Transport.REST)
+          || (BIGQUERY_LEGACY_PAGINATION_ALLOWLIST.contains(methodDescriptor.getFullName()))) {
         fieldNames.add("max_results");
       }
       for (String fieldName : fieldNames) {

gapic-generator-java/src/test/java/com/google/api/generator/gapic/composer/ClientLibraryReflectConfigComposerTest.java

@@ -61,6 +61,8 @@ void composeReflectConfigs_showcase() {
             "com.google.showcase.v1beta1.Foobar$Builder",
             "com.google.showcase.v1beta1.Object",
             "com.google.showcase.v1beta1.Object$Builder",
+            "com.google.showcase.v1beta1.PagedExpandLegacyRequest",
+            "com.google.showcase.v1beta1.PagedExpandLegacyRequest$Builder",
             "com.google.showcase.v1beta1.PagedExpandRequest",
             "com.google.showcase.v1beta1.PagedExpandRequest$Builder",
             "com.google.showcase.v1beta1.PagedExpandResponse",

gapic-generator-java/src/test/java/com/google/api/generator/gapic/composer/grpc/GrpcServiceStubClassComposerTest.java

@@ -101,4 +101,15 @@ void generateGrpcServiceStubClass_callableNameType() {
     Assert.assertGoldenClass(this.getClass(), clazz, "GrpcCallableNameTypeStub.golden");
     Assert.assertEmptySamples(clazz.samples());
   }
+
+  @Test
+  void generateGrpcServiceStubClass_bigQuery() {
+    GapicContext context = GrpcTestProtoLoader.instance().parseBigqueryService();
+
+    Service bigqueryJobService = context.services().get(0);
+    GapicClass clazz =
+        GrpcServiceStubClassComposer.instance().generate(context, bigqueryJobService);
+    Assert.assertGoldenClass(this.getClass(), clazz, "GrpcBigQueryJobServiceStub.golden");
+    Assert.assertEmptySamples(clazz.samples());
+  }
 }

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

@@ -62,7 +62,13 @@ static Stream<Arguments> data() {
             GrpcTestProtoLoader.instance().parseSelectiveGenerationTesting(),
             "localhost:7469",
             "v1beta1",
-            1));
+            1),
+        Arguments.of(
+            "JobServiceStubSettings",
+            GrpcTestProtoLoader.instance().parseBigqueryService(),
+            "bigquery",
+            "v2",
+            0));
   }
 
   @ParameterizedTest

gapic-generator-java/src/test/java/com/google/api/generator/gapic/composer/grpc/goldens/EchoClient.golden

@@ -206,6 +206,20 @@ import javax.annotation.Generated;
  *      </ul>
  *       </td>
  *    </tr>
+ *    <tr>
+ *      <td><p> PagedExpandLegacy</td>
+ *      <td><p> </td>
+ *      <td>
+ *      <p>Request object method variants only take one parameter, a request object, which must be constructed before the call.</p>
+ *      <ul>
+ *           <li><p> pagedExpandLegacy(PagedExpandLegacyRequest request)
+ *      </ul>
+ *      <p>Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.</p>
+ *      <ul>
+ *           <li><p> pagedExpandLegacyCallable()
+ *      </ul>
+ *       </td>
+ *    </tr>
  *  </table>
  *
  * <p>See the individual methods for example code.
@@ -1124,6 +1138,63 @@ public class EchoClient implements BackgroundResource {
     return stub.collideNameCallable();
   }
 
+  // AUTO-GENERATED DOCUMENTATION AND METHOD.
+  /**
+   * Sample code:
+   *
+   * <pre>{@code
+   * // This snippet has been automatically generated and should be regarded as a code template only.
+   * // It will require modifications to work:
+   * // - It may require correct/in-range values for request initialization.
+   * // - It may require specifying regional endpoints when creating the service client as shown in
+   * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+   * try (EchoClient echoClient = EchoClient.create()) {
+   *   PagedExpandLegacyRequest request =
+   *       PagedExpandLegacyRequest.newBuilder()
+   *           .setContent("content951530617")
+   *           .setMaxResults(1128457243)
+   *           .setPageToken("pageToken873572522")
+   *           .build();
+   *   PagedExpandResponse response = echoClient.pagedExpandLegacy(request);
+   * }
+   * }</pre>
+   *
+   * @param request The request object containing all of the parameters for the API call.
+   * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+   */
+  public final PagedExpandResponse pagedExpandLegacy(PagedExpandLegacyRequest request) {
+    return pagedExpandLegacyCallable().call(request);
+  }
+
+  // AUTO-GENERATED DOCUMENTATION AND METHOD.
+  /**
+   * Sample code:
+   *
+   * <pre>{@code
+   * // This snippet has been automatically generated and should be regarded as a code template only.
+   * // It will require modifications to work:
+   * // - It may require correct/in-range values for request initialization.
+   * // - It may require specifying regional endpoints when creating the service client as shown in
+   * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+   * try (EchoClient echoClient = EchoClient.create()) {
+   *   PagedExpandLegacyRequest request =
+   *       PagedExpandLegacyRequest.newBuilder()
+   *           .setContent("content951530617")
+   *           .setMaxResults(1128457243)
+   *           .setPageToken("pageToken873572522")
+   *           .build();
+   *   ApiFuture<PagedExpandResponse> future =
+   *       echoClient.pagedExpandLegacyCallable().futureCall(request);
+   *   // Do something.
+   *   PagedExpandResponse response = future.get();
+   * }
+   * }</pre>
+   */
+  public final UnaryCallable<PagedExpandLegacyRequest, PagedExpandResponse>
+      pagedExpandLegacyCallable() {
+    return stub.pagedExpandLegacyCallable();
+  }
+
   @Override
   public final void close() {
     stub.close();

gapic-generator-java/src/test/java/com/google/api/generator/gapic/composer/grpc/goldens/EchoClientTest.golden

@@ -26,6 +26,7 @@ import com.google.protobuf.Timestamp;
 import com.google.rpc.Status;
 import io.grpc.StatusRuntimeException;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.UUID;
@@ -890,4 +891,55 @@ public class EchoClientTest {
       // Expected exception.
     }
   }
+
+  @Test
+  public void pagedExpandLegacyTest() throws Exception {
+    PagedExpandResponse expectedResponse =
+        PagedExpandResponse.newBuilder()
+            .addAllResponses(new ArrayList<EchoResponse>())
+            .setNextPageToken("nextPageToken-1386094857")
+            .build();
+    mockEcho.addResponse(expectedResponse);
+
+    PagedExpandLegacyRequest request =
+        PagedExpandLegacyRequest.newBuilder()
+            .setContent("content951530617")
+            .setMaxResults(1128457243)
+            .setPageToken("pageToken873572522")
+            .build();
+
+    PagedExpandResponse actualResponse = client.pagedExpandLegacy(request);
+    Assert.assertEquals(expectedResponse, actualResponse);
+
+    List<AbstractMessage> actualRequests = mockEcho.getRequests();
+    Assert.assertEquals(1, actualRequests.size());
+    PagedExpandLegacyRequest actualRequest = ((PagedExpandLegacyRequest) actualRequests.get(0));
+
+    Assert.assertEquals(request.getContent(), actualRequest.getContent());
+    Assert.assertEquals(request.getMaxResults(), actualRequest.getMaxResults());
+    Assert.assertEquals(request.getPageToken(), actualRequest.getPageToken());
+    Assert.assertTrue(
+        channelProvider.isHeaderSent(
+            ApiClientHeaderProvider.getDefaultApiClientHeaderKey(),
+            GaxGrpcProperties.getDefaultApiClientHeaderPattern()));
+  }
+
+  @Test
+  public void pagedExpandLegacyExceptionTest() throws Exception {
+    StatusRuntimeException exception = new StatusRuntimeException(io.grpc.Status.INVALID_ARGUMENT);
+    mockEcho.addException(exception);
+
+    try {
+      PagedExpandLegacyRequest request =
+          PagedExpandLegacyRequest.newBuilder()
+              .setContent("content951530617")
+              .setMaxResults(1128457243)
+              .setPageToken("pageToken873572522")
+              .build();
+      client.pagedExpandLegacy(request);
+      Assert.fail("No exception raised");
+    } catch (InvalidArgumentException e) {
+      // Expected exception.
+    }
+  }
 }

gapic-generator-java/src/test/java/com/google/api/generator/gapic/composer/grpc/goldens/EchoSettings.golden

@@ -159,6 +159,12 @@ public class EchoSettings extends ClientSettings<EchoSettings> {
     return ((EchoStubSettings) getStubSettings()).collideNameSettings();
   }
 
+  /** Returns the object with the settings used for calls to pagedExpandLegacy. */
+  public UnaryCallSettings<PagedExpandLegacyRequest, PagedExpandResponse>
+      pagedExpandLegacySettings() {
+    return ((EchoStubSettings) getStubSettings()).pagedExpandLegacySettings();
+  }
+
   public static final EchoSettings create(EchoStubSettings stub) throws IOException {
     return new EchoSettings.Builder(stub.toBuilder()).build();
   }
@@ -314,6 +320,12 @@ public class EchoSettings extends ClientSettings<EchoSettings> {
       return getStubSettingsBuilder().collideNameSettings();
     }
 
+    /** Returns the builder for the settings used for calls to pagedExpandLegacy. */
+    public UnaryCallSettings.Builder<PagedExpandLegacyRequest, PagedExpandResponse>
+        pagedExpandLegacySettings() {
+      return getStubSettingsBuilder().pagedExpandLegacySettings();
+    }
+
     @Override
     public EchoSettings build() throws IOException {
       return new EchoSettings(this);

gapic-generator-java/src/test/java/com/google/api/generator/gapic/composer/grpc/goldens/EchoStub.golden

@@ -18,6 +18,7 @@ import com.google.showcase.v1beta1.EchoRequest;
 import com.google.showcase.v1beta1.EchoResponse;
 import com.google.showcase.v1beta1.ExpandRequest;
 import com.google.showcase.v1beta1.Object;
+import com.google.showcase.v1beta1.PagedExpandLegacyRequest;
 import com.google.showcase.v1beta1.PagedExpandRequest;
 import com.google.showcase.v1beta1.PagedExpandResponse;
 import com.google.showcase.v1beta1.WaitMetadata;
@@ -92,6 +93,10 @@ public abstract class EchoStub implements BackgroundResource {
     throw new UnsupportedOperationException("Not implemented: collideNameCallable()");
   }
 
+  public UnaryCallable<PagedExpandLegacyRequest, PagedExpandResponse> pagedExpandLegacyCallable() {
+    throw new UnsupportedOperationException("Not implemented: pagedExpandLegacyCallable()");
+  }
+
   @Override
   public abstract void close();
 }