feat: adds augmented pagination to account for BQ family of APIs (#2372)

Co-authored-by: Anthonios Partheniou <[email protected]>

Author: chalmerlowe

gapic/schema/wrappers.py

@@ -1830,13 +1830,34 @@ def ident(self) -> metadata.Address:
         """Return the identifier data to be used in templates."""
         return self.meta.address
 
+    def _validate_paged_field_size_type(self, page_field_size) -> bool:
+        """Validates allowed paged_field_size type(s).
+
+        Confirms whether the paged_field_size.type is an allowed wrapper type:
+        The norm is for type to be int, but an additional check is included to
+        account for BigQuery legacy APIs which allowed UInt32Value and
+        Int32Value.
+        """
+
+        pb_type = page_field_size.type
+
+        return pb_type == int or (
+            isinstance(pb_type, MessageType)
+            and pb_type.message_pb.name in {"UInt32Value", "Int32Value"}
+        )
+
     @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 +1866,18 @@ 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
+
+        # Confirm whether the paged_field_size is an allowed type.
+        if not self._validate_paged_field_size_type(page_field_size=page_field_size):
             return None
 
         # Return the first repeated field.

tests/fragments/test_pagination_max_results_and_wrapper.proto

@@ -0,0 +1,36 @@
+// Copyright (C) 2025  Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package google.fragment;
+
+import "google/api/client.proto";
+import "google/protobuf/wrappers.proto";
+
+service MaxResultsDatasetService {
+  option (google.api.default_host) = "my.example.com";
+  rpc ListMaxResultsDataset(ListMaxResultsDatasetRequest) returns (ListMaxResultsDatasetResponse) {
+  }
+}
+
+message ListMaxResultsDatasetRequest {
+  google.protobuf.UInt32Value max_results = 2;
+  string page_token = 3;
+}
+
+message ListMaxResultsDatasetResponse {
+  string next_page_token = 3;
+  repeated string datasets = 4;
+}

tests/unit/schema/wrappers/test_method.py

@@ -23,6 +23,7 @@
 from google.api import routing_pb2
 from google.cloud import extended_operations_pb2 as ex_ops_pb2
 from google.protobuf import descriptor_pb2
+from google.protobuf import wrappers_pb2
 
 from gapic.schema import metadata
 from gapic.schema import wrappers
@@ -189,6 +190,63 @@ def test_method_paged_result_field_no_page_field():
     assert method.paged_result_field is None
 
 
+def test_method_paged_result_field_invalid_wrapper_type():
+    """Validate paged_result_field() returns None if page_size/max_results wrappertypes
+    are not allowed types.
+    """
+
+    # page_size is not allowed wrappertype
+    parent = make_field(name="parent", type="TYPE_STRING")
+    page_size = make_field(name="page_size", type="TYPE_DOUBLE")  # not an allowed type
+    page_token = make_field(name="page_token", type="TYPE_STRING")
+    foos = make_field(name="foos", message=make_message("Foo"), repeated=True)
+    next_page_token = make_field(name="next_page_token", type="TYPE_STRING")
+
+    input_msg = make_message(
+        name="ListFoosRequest",
+        fields=(
+            parent,
+            page_size,
+            page_token,
+        ),
+    )
+    output_msg = make_message(
+        name="ListFoosResponse",
+        fields=(
+            foos,
+            next_page_token,
+        ),
+    )
+    method = make_method(
+        "ListFoos",
+        input_message=input_msg,
+        output_message=output_msg,
+    )
+    assert method.paged_result_field is None
+
+    # max_results is not allowed wrappertype
+    max_results = make_field(
+        name="max_results", type="TYPE_STRING"
+    )  # not an allowed type
+
+    input_msg = make_message(
+        name="ListFoosRequest",
+        fields=(
+            parent,
+            max_results,
+            page_token,
+        ),
+    )
+
+    method = make_method(
+        "ListFoos",
+        input_message=input_msg,
+        output_message=output_msg,
+    )
+
+    assert method.paged_result_field is None
+
+
 def test_method_paged_result_ref_types():
     input_msg = make_message(
         name="ListSquidsRequest",
@@ -999,3 +1057,63 @@ def test_mixin_rule():
         "city": {},
     }
     assert e == m.sample_request
+
+
+@pytest.mark.parametrize(
+    "field_type, pb_type, expected",
+    [
+        # valid paged_result_field candidates
+        (int, "TYPE_INT32", True),
+        (wrappers_pb2.UInt32Value, "TYPE_MESSAGE", True),
+        (wrappers_pb2.Int32Value, "TYPE_MESSAGE", True),
+        # invalid paged_result_field candidates
+        (float, "TYPE_DOUBLE", False),
+        (wrappers_pb2.UInt32Value, "TYPE_DOUBLE", False),
+        (wrappers_pb2.Int32Value, "TYPE_DOUBLE", False),
+    ],
+)
+def test__validate_paged_field_size_type(field_type, pb_type, expected):
+    """Test _validate_paged_field_size_type with wrapper types and type indicators."""
+
+    # Setup
+    if pb_type in {"TYPE_INT32", "TYPE_DOUBLE"}:
+        page_size = make_field(name="page_size", type=pb_type)
+    else:
+        # expecting TYPE_MESSAGE which in this context is associated with
+        # *Int32Value in legacy APIs such as BigQuery.
+        # See: https://github.com/googleapis/gapic-generator-python/blob/c8b7229ba2865d6a2f5966aa151be121de81f92d/gapic/schema/wrappers.py#L378C1-L411C10
+        page_size = make_field(
+            name="max_results",
+            type=pb_type,
+            message=make_message(name=field_type.DESCRIPTOR.name),
+        )
+
+    parent = make_field(name="parent", type="TYPE_STRING")
+    page_token = make_field(name="page_token", type="TYPE_STRING")
+    next_page_token = make_field(name="next_page_token", type="TYPE_STRING")
+
+    input_msg = make_message(
+        name="ListFoosRequest",
+        fields=(
+            parent,
+            page_size,
+            page_token,
+        ),
+    )
+
+    output_msg = make_message(
+        name="ListFoosResponse",
+        fields=(
+            make_field(name="foos", message=make_message("Foo"), repeated=True),
+            next_page_token,
+        ),
+    )
+
+    method = make_method(
+        "ListFoos",
+        input_message=input_msg,
+        output_message=output_msg,
+    )
+
+    actual = method._validate_paged_field_size_type(page_field_size=page_size)
+    assert actual == expected