feat: handle auto pagination for BigQuery v2 (#3829)

Fixes # b/360114900 & b/352331075

Author: PhongChuong

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();
 }