Add option to limit the resultset returned by paginate helper (#4475)

server/forge/bitbucket/bitbucket.go

@@ -143,7 +143,7 @@ func (c *config) Teams(ctx context.Context, u *model.User) ([]*model.Team, error
 			return nil, err
 		}
 		return convertWorkspaceList(resp.Values), nil
-	})
+	}, -1)
 }
 
 // Repo returns the named Bitbucket repository.
@@ -190,7 +190,7 @@ func (c *config) Repos(ctx context.Context, u *model.User) ([]*model.Repo, error
 			return nil, err
 		}
 		return resp.Values, nil
-	})
+	}, -1)
 	if err != nil {
 		return nil, err
 	}
@@ -331,7 +331,7 @@ func (c *config) Deactivate(ctx context.Context, u *model.User, r *model.Repo, l
 			return nil, err
 		}
 		return hooks.Values, nil
-	})
+	}, -1)
 	if err != nil {
 		return err
 	}

server/forge/bitbucket/internal/client.go

@@ -115,7 +115,7 @@ func (c *Client) ListReposAll(workspace string) ([]*Repo, error) {
 			return nil, err
 		}
 		return resp.Values, nil
-	})
+	}, -1)
 }
 
 func (c *Client) FindHook(owner, name, id string) (*Hook, error) {
@@ -183,7 +183,7 @@ func (c *Client) ListPermissionsAll() ([]*RepoPerm, error) {
 			return nil, err
 		}
 		return resp.Values, nil
-	})
+	}, -1)
 }
 
 func (c *Client) ListBranches(owner, name string, opts *ListOpts) ([]*Branch, error) {

server/forge/forgejo/forgejo.go

@@ -201,7 +201,7 @@ func (c *Forgejo) Teams(ctx context.Context, u *model.User) ([]*model.Team, erro
 			teams = append(teams, toTeam(org, c.url))
 		}
 		return teams, err
-	})
+	}, -1)
 }
 
 // TeamPerm is not supported by the Forgejo driver.
@@ -253,7 +253,7 @@ func (c *Forgejo) Repos(ctx context.Context, u *model.User) ([]*model.Repo, erro
 			},
 		)
 		return repos, err
-	})
+	}, -1)
 
 	result := make([]*model.Repo, 0, len(repos))
 	for _, repo := range repos {
@@ -411,7 +411,7 @@ func (c *Forgejo) Deactivate(ctx context.Context, u *model.User, r *model.Repo,
 			},
 		})
 		return hooks, err
-	})
+	}, -1)
 	if err != nil {
 		return err
 	}
@@ -647,7 +647,7 @@ func (c *Forgejo) getChangedFilesForPR(ctx context.Context, repo *model.Repo, in
 			files = append(files, file.Filename)
 		}
 		return files, nil
-	})
+	}, -1)
 }
 
 func (c *Forgejo) getTagCommitSHA(ctx context.Context, repo *model.Repo, tagName string) (string, error) {

server/forge/gitea/gitea.go

@@ -203,7 +203,7 @@ func (c *Gitea) Teams(ctx context.Context, u *model.User) ([]*model.Team, error)
 			teams = append(teams, toTeam(org, c.url))
 		}
 		return teams, err
-	})
+	}, -1)
 }
 
 // TeamPerm is not supported by the Gitea driver.
@@ -255,7 +255,7 @@ func (c *Gitea) Repos(ctx context.Context, u *model.User) ([]*model.Repo, error)
 			},
 		)
 		return repos, err
-	})
+	}, -1)
 
 	result := make([]*model.Repo, 0, len(repos))
 	for _, repo := range repos {
@@ -413,7 +413,7 @@ func (c *Gitea) Deactivate(ctx context.Context, u *model.User, r *model.Repo, li
 			},
 		})
 		return hooks, err
-	})
+	}, -1)
 	if err != nil {
 		return err
 	}
@@ -654,7 +654,7 @@ func (c *Gitea) getChangedFilesForPR(ctx context.Context, repo *model.Repo, inde
 			files = append(files, file.Filename)
 		}
 		return files, nil
-	})
+	}, -1)
 }
 
 func (c *Gitea) getTagCommitSHA(ctx context.Context, repo *model.Repo, tagName string) (string, error) {

server/forge/github/github.go

@@ -660,7 +660,7 @@ func (c *client) loadChangedFilesFromPullRequest(ctx context.Context, pull *gith
 			opts.Page = resp.NextPage
 		}
 		return utils.DeduplicateStrings(fileList), nil
-	})
+	}, -1)
 
 	return pipeline, err
 }

shared/utils/paginate.go

@@ -15,19 +15,37 @@
 package utils
 
 // Paginate iterates over a func call until it does not return new items and return it as list.
-func Paginate[T any](get func(page int) ([]T, error)) ([]T, error) {
+func Paginate[T any](get func(page int) ([]T, error), limit int) ([]T, error) {
 	items := make([]T, 0, 10)
 	page := 1
 	lenFirstBatch := -1
 
 	for {
+		// limit < 1 means get all results
+		remaining := 0
+		if limit > 0 {
+			remaining = limit - len(items)
+			if remaining <= 0 {
+				break
+			}
+		}
+
 		batch, err := get(page)
 		if err != nil {
 			return nil, err
 		}
+
+		// Take only what we need from this batch if limit > 0
+		if limit > 0 && len(batch) > remaining {
+			batch = batch[:remaining]
+		}
+
 		items = append(items, batch...)
 
 		if page == 1 {
+			if len(batch) == 0 {
+				return items, nil
+			}
 			lenFirstBatch = len(batch)
 		} else if len(batch) < lenFirstBatch || len(batch) == 0 {
 			break

shared/utils/paginate_test.go

@@ -21,27 +21,72 @@ import (
 )
 
 func TestPaginate(t *testing.T) {
-	apiExec := 0
+	// Generic mock generator that can handle all cases
-	apiMock := func(page int) []int {
+	createMock := func(pages [][]int) func(page int) []int {
-		apiExec++
+		return func(page int) []int {
-		switch page {
+			if page <= 0 {
-		case 0, 1:
+				page = 0
-			return []int{11, 12, 13}
+			} else {
-		case 2:
+				page--
-			return []int{21, 22, 23}
+			}
-		case 3:
+
-			return []int{31, 32}
+			if page >= len(pages) {
-		default:
+				return []int{}
-			return []int{}
+			}
+
+			return pages[page]
 		}
 	}
 
-	result, _ := Paginate(func(page int) ([]int, error) {
+	tests := []struct {
-		return apiMock(page), nil
+		name     string
-	})
+		limit    int
+		pages    [][]int
+		expected []int
+		apiCalls int
+	}{
+		{
+			name:     "multiple pages",
+			limit:    -1,
+			pages:    [][]int{{11, 12, 13}, {21, 22, 23}, {31, 32}},
+			expected: []int{11, 12, 13, 21, 22, 23, 31, 32},
+			apiCalls: 3,
+		},
+		{
+			name:     "zero limit",
+			limit:    0,
+			pages:    [][]int{{1, 2, 3}, {1, 2, 3}, {1, 2, 3}},
+			expected: []int{1, 2, 3, 1, 2, 3, 1, 2, 3},
+			apiCalls: 4,
+		},
+		{
+			name:     "empty result",
+			limit:    5,
+			pages:    [][]int{{}},
+			expected: []int{},
+			apiCalls: 1,
+		},
+		{
+			name:     "limit less than batch",
+			limit:    2,
+			pages:    [][]int{{1, 2, 3, 4, 5}},
+			expected: []int{1, 2},
+			apiCalls: 1,
+		},
+	}
 
-	assert.EqualValues(t, 3, apiExec)
+	for _, tt := range tests {
-	if assert.Len(t, result, 8) {
+		t.Run(tt.name, func(t *testing.T) {
-		assert.EqualValues(t, []int{11, 12, 13, 21, 22, 23, 31, 32}, result)
+			apiExec := 0
+			mock := createMock(tt.pages)
+
+			result, _ := Paginate(func(page int) ([]int, error) {
+				apiExec++
+				return mock(page), nil
+			}, tt.limit)
+
+			assert.EqualValues(t, tt.apiCalls, apiExec)
+			assert.EqualValues(t, tt.expected, result)
+		})
 	}
 }