feat: adds augmented pagination to account for BQ family of APIs

gapic/schema/wrappers.py

@@ -1495,6 +1495,8 @@ class Method:
         default_factory=metadata.Metadata,
     )
 
+    PACKAGES_WITH_ALLOWED_WRAPPERS = {"google.cloud.bigquery.v2"}
+
     def __getattr__(self, name):
         return getattr(self.method_pb, name)
 
@@ -1832,11 +1834,16 @@ def ident(self) -> metadata.Address:
 
     @utils.cached_property
     def paged_result_field(self) -> Optional[Field]:
-        """Return the response pagination field if the method is paginated."""
-        # If the request field lacks any of the expected pagination fields,
-        # then the method is not paginated.
+        """Return the response pagination field if the method is paginated.
+
+        The request field must have a page_token field and a page_size field (or
+        for legacy APIs, a max_results field) and the response field
+        must have a next_token_field and a repeated field.
+
+        For the purposes of supporting legacy APIs, additional wrapper types are
+        allowed.
+        """
 
-        # The request must have page_token and next_page_token as they keep track of pages
         for source, source_type, name in (
             (self.input, str, "page_token"),
             (self.output, str, "next_page_token"),
@@ -1845,13 +1852,32 @@ def paged_result_field(self) -> Optional[Field]:
             if not field or field.type != source_type:
                 return None
 
-        # The request must have max_results or page_size
+        # The request must have page_size (or max_results if legacy API)
         page_fields = (
             self.input.fields.get("max_results", None),
             self.input.fields.get("page_size", None),
         )
         page_field_size = next((field for field in page_fields if field), None)
-        if not page_field_size or page_field_size.type != int:
+        if not page_field_size:
+            return None
+
+        # If the a legacy API uses the UInt32Value and Int32Value wrappers,
+        # the service package must be in the allowlist.
+        package_name = self.input.meta.address.proto_package
+
+        if page_field_size.type == int or (
+            # The following additional checks are related to several members of the BQ family of
+            # APIs, which use the legacy wrapper types: "UInt32Value" and "Int32Value"}
+            # for max_results.
+            # NOTE:
+            #   bigquery_v2 should be paginated
+            #   but bigquery_connection_v1beta1 should NOT be paginated
+            isinstance(page_field_size.type, MessageType)
+            and page_field_size.type.message_pb.name in {"UInt32Value", "Int32Value"}
+            and package_name in self.PACKAGES_WITH_ALLOWED_WRAPPERS
+        ):
+            pass
+        else:
             return None
 
         # Return the first repeated field.