feat(sidekick): support aggregated pagination (#2432)

Author: coryan

internal/sidekick/internal/parser/disco_test.go

@@ -112,6 +112,30 @@ func TestDisco_ParsePagination(t *testing.T) {
 	}
 }
 
+func TestDisco_ParsePaginationAggregate(t *testing.T) {
+	model, err := ParseDisco(discoSourceFile, "", map[string]string{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	updateMethodPagination(nil, model)
+	wantID := "..machineTypes.aggregatedList"
+	got, ok := model.State.MethodByID[wantID]
+	if !ok {
+		t.Fatalf("expected method %s in the API model", wantID)
+	}
+	wantPagination := &api.Field{
+		Name:     "pageToken",
+		JSONName: "pageToken",
+		ID:       "..machineTypes.aggregatedListRequest.pageToken",
+		Typez:    api.STRING_TYPE,
+		TypezID:  "string",
+		Optional: true,
+	}
+	if diff := cmp.Diff(wantPagination, got.Pagination, cmpopts.IgnoreFields(api.Field{}, "Documentation")); diff != "" {
+		t.Errorf("mismatch (-want, +got):\n%s", diff)
+	}
+}
+
 func TestDisco_ParseDeprecatedEnum(t *testing.T) {
 	model, err := ParseDisco(discoSourceFile, "", map[string]string{})
 	if err != nil {

internal/sidekick/internal/parser/pagination.go

@@ -122,12 +122,16 @@ func paginationResponseItem(overrides []config.PaginationOverride, methodID stri
 		return response.Fields[fieldIdx]
 	}
 
+	var mapItems *api.Field
 	for _, field := range response.Fields {
+		if field.Map && mapItems == nil {
+			mapItems = field
+		}
 		if field.Repeated && field.Typez == api.MESSAGE_TYPE {
 			return field
 		}
 	}
-	return nil
+	return mapItems
 }
 
 func paginationResponseNextPageToken(response *api.Message) *api.Field {

internal/sidekick/internal/parser/pagination_test.go

@@ -536,7 +536,97 @@ func TestPaginationResponseErrors(t *testing.T) {
 	}
 }
 
-func TestPaginationResponseItem(t *testing.T) {
+func TestPaginationResponseItemMatching(t *testing.T) {
+	for _, test := range []struct {
+		Repeated bool
+		Map      bool
+		Typez    api.Typez
+		Name     string
+	}{
+		{false, true, api.MESSAGE_TYPE, "items"},
+		{true, false, api.MESSAGE_TYPE, "items"},
+	} {
+		response := &api.Message{
+			Name: "Response",
+			ID:   ".package.Response",
+			Fields: []*api.Field{
+				{
+					Name:     test.Name,
+					JSONName: test.Name,
+					Typez:    test.Typez,
+					Repeated: test.Repeated,
+					Map:      test.Map,
+				},
+			},
+		}
+		got := paginationResponseItem(nil, "package.Service.List", response)
+		if diff := cmp.Diff(response.Fields[0], got); diff != "" {
+			t.Errorf("mismatch (-want, +got):\n%s", diff)
+		}
+	}
+}
+
+func TestPaginationResponseItemMatchingMany(t *testing.T) {
+	for _, test := range []struct {
+		Repeated bool
+		Map      bool
+	}{
+		{true, false},
+		{false, true},
+	} {
+		response := &api.Message{
+			Name: "Response",
+			ID:   ".package.Response",
+			Fields: []*api.Field{
+				{
+					Name:     "first",
+					JSONName: "first",
+					Typez:    api.MESSAGE_TYPE,
+					Repeated: test.Repeated,
+					Map:      test.Map,
+				},
+				{
+					Name:     "second",
+					JSONName: "second",
+					Typez:    api.MESSAGE_TYPE,
+					Repeated: test.Repeated,
+					Map:      test.Map,
+				},
+			},
+		}
+		got := paginationResponseItem(nil, "package.Service.List", response)
+		if diff := cmp.Diff(response.Fields[0], got); diff != "" {
+			t.Errorf("mismatch (-want, +got):\n%s", diff)
+		}
+	}
+}
+
+func TestPaginationResponseItemMatchingPreferRepeatedOverMap(t *testing.T) {
+	response := &api.Message{
+		Name: "Response",
+		ID:   ".package.Response",
+		Fields: []*api.Field{
+			{
+				Name:     "map",
+				JSONName: "map",
+				Typez:    api.MESSAGE_TYPE,
+				Map:      true,
+			},
+			{
+				Name:     "repeated",
+				JSONName: "repeated",
+				Typez:    api.MESSAGE_TYPE,
+				Repeated: true,
+			},
+		},
+	}
+	got := paginationResponseItem(nil, "package.Service.List", response)
+	if diff := cmp.Diff(response.Fields[1], got); diff != "" {
+		t.Errorf("mismatch (-want, +got):\n%s", diff)
+	}
+}
+
+func TestPaginationResponseItemNotMatching(t *testing.T) {
 	overrides := []config.PaginationOverride{
 		{ID: ".package.Service.List", ItemField: "--invalid--"},
 	}

internal/sidekick/internal/rust/templates/common/message.mustache

@@ -225,11 +225,21 @@ impl wkt::message::Message for {{Codec.Name}} {
 #[doc(hidden)]
 impl gax::paginator::internal::PageableResponse for {{Codec.Name}} {
     {{#PageableItem}}
+    {{^Map}}
     type PageItem = {{{Codec.PrimitiveFieldType}}};
 
     fn items(self) -> std::vec::Vec<Self::PageItem> {
         self.{{{Codec.FieldName}}}
     }
+    {{/Map}}
+    {{#Map}}
+    type PageItem = ({{Codec.KeyType}}, {{Codec.ValueType}});
+
+    fn items(self) -> std::vec::Vec<Self::PageItem> {
+        use std::iter::{IntoIterator, Iterator};
+        self.{{{Codec.FieldName}}}.into_iter().collect()
+    }
+    {{/Map}}
 
     {{/PageableItem}}
     {{#NextPageToken}}