add and use RESTWithNext

nate smith · nate smith · commit a3fba669bd23 · 2022-03-29T10:15:33.000-05:00
diff --git a/api/client.go b/api/client.go
@@ -316,41 +316,55 @@ func graphQLClient(h *http.Client, hostname string) *graphql.Client {
 
 // REST performs a REST request and parses the response.
 func (c Client) REST(hostname string, method string, p string, body io.Reader, data interface{}) error {
+	_, err := c.RESTWithNext(hostname, method, p, body, data)
+	return err
+}
+
+func (c Client) RESTWithNext(hostname string, method string, p string, body io.Reader, data interface{}) (string, error) {
 	req, err := http.NewRequest(method, restURL(hostname, p), body)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	req.Header.Set("Content-Type", "application/json; charset=utf-8")
 
 	resp, err := c.http.Do(req)
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer resp.Body.Close()
 
 	success := resp.StatusCode >= 200 && resp.StatusCode < 300
 	if !success {
-		return HandleHTTPError(resp)
+		return "", HandleHTTPError(resp)
 	}
 
 	if resp.StatusCode == http.StatusNoContent {
-		return nil
+		return "", nil
 	}
 
 	b, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	err = json.Unmarshal(b, &data)
 	if err != nil {
-		return err
+		return "", err
 	}
 
-	return nil
+	var next string
+	for _, m := range linkRE.FindAllStringSubmatch(resp.Header.Get("Link"), -1) {
+		if len(m) > 2 && m[2] == "next" {
+			next = m[1]
+		}
+	}
+
+	return next, nil
 }
 
+var linkRE = regexp.MustCompile(`<([^>]+)>;\s*rel="([^"]+)"`)
+
 func restURL(hostname string, pathOrURL string) string {
 	if strings.HasPrefix(pathOrURL, "https://") || strings.HasPrefix(pathOrURL, "http://") {
 		return pathOrURL
diff --git a/pkg/cmd/status/status.go b/pkg/cmd/status/status.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"regexp"
 	"sort"
 	"strings"
 	"time"
@@ -217,6 +218,17 @@ func (s *StatusGetter) ActualMention(n Notification) (string, error) {
 // These are split up by endpoint since it is along that boundary we parallelize
 // work
 
+var linkRE = regexp.MustCompile(`<([^>]+)>;\s*rel="([^"]+)"`)
+
+func findNextPage(resp *http.Response) (string, bool) {
+	for _, m := range linkRE.FindAllStringSubmatch(resp.Header.Get("Link"), -1) {
+		if len(m) > 2 && m[2] == "next" {
+			return m[1], true
+		}
+	}
+	return "", false
+}
+
 // Populate .Mentions
 func (s *StatusGetter) LoadNotifications() error {
 	perPage := 100
@@ -233,21 +245,24 @@ func (s *StatusGetter) LoadNotifications() error {
 	// not work with PATs right now.
 	var ns []Notification
 	var resp []Notification
-	pages := 3
-	for page := 1; page <= pages; page++ {
-		query.Add("page", fmt.Sprintf("%d", page))
-		p := fmt.Sprintf("notifications?%s", query.Encode())
-		err := c.REST(ghinstance.Default(), "GET", p, nil, &resp)
+	pages := 0
+	p := fmt.Sprintf("notifications?%s", query.Encode())
+	for pages < 3 {
+		next, err := c.RESTWithNext(ghinstance.Default(), "GET", p, nil, &resp)
 		if err != nil {
 			var httpErr api.HTTPError
 			if !errors.As(err, &httpErr) || httpErr.StatusCode != 404 {
 				return fmt.Errorf("could not get notifications: %w", err)
 			}
 		}
 		ns = append(ns, resp...)
-		if len(resp) == 0 || len(resp) < perPage {
+
+		if next == "" || len(resp) < perPage {
 			break
 		}
+
+		pages++
+		p = next
 	}
 
 	s.Mentions = []StatusItem{}
@@ -430,21 +445,23 @@ func (s *StatusGetter) LoadEvents() error {
 
 	var events []Event
 	var resp []Event
-	pages := 2
-	for page := 1; page <= pages; page++ {
-		query.Add("page", fmt.Sprintf("%d", page))
-		p := fmt.Sprintf("users/%s/received_events?%s", currentUsername, query.Encode())
-		err := c.REST(ghinstance.Default(), "GET", p, nil, &resp)
+	pages := 0
+	p := fmt.Sprintf("users/%s/received_events?%s", currentUsername, query.Encode())
+	for pages < 2 {
+		next, err := c.RESTWithNext(ghinstance.Default(), "GET", p, nil, &resp)
 		if err != nil {
 			var httpErr api.HTTPError
 			if !errors.As(err, &httpErr) || httpErr.StatusCode != 404 {
 				return fmt.Errorf("could not get events: %w", err)
 			}
 		}
 		events = append(events, resp...)
-		if len(resp) == 0 || len(resp) < perPage {
+		if next == "" || len(resp) < perPage {
 			break
 		}
+
+		pages++
+		p = next
 	}
 
 	s.RepoActivity = []StatusItem{}
