pbd client: adds nodes pagination cursor

pkg/puppetdb/common_test.go

@@ -1,9 +1,12 @@
 package puppetdb
 
 import (
+	"fmt"
 	"net/http"
 	"net/url"
 	"os"
+	"path/filepath"
+	"strconv"
 	"testing"
 	"time"
 
@@ -31,6 +34,56 @@ func setupGetResponder(t *testing.T, url, query, responseFilename string) {
 	response.Body.Close()
 }
 
+type mockPaginatedGetOptions struct {
+	limit         int
+	total         int
+	pageFilenames []string
+}
+
+func setupPaginatedGetResponder(t *testing.T, url, query string, opts mockPaginatedGetOptions) {
+	var pages [][]byte
+
+	for _, pfn := range opts.pageFilenames {
+		responseBody, err := os.ReadFile(filepath.Join("testdata", pfn))
+		require.NoError(t, err)
+
+		pages = append(pages, responseBody)
+	}
+
+	responder := func(r *http.Request) (*http.Response, error) {
+		var (
+			offset  int
+			pageNum int
+			err     error
+		)
+
+		offsetS := r.URL.Query().Get("offset")
+		if offsetS != "" {
+			offset, err = strconv.Atoi(offsetS)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		if offset > 0 {
+			pageNum = offset / opts.limit
+		}
+
+		responseBody := pages[pageNum]
+
+		response := httpmock.NewBytesResponse(http.StatusOK, responseBody)
+		response.Header.Set("Content-Type", "application/json")
+		response.Header.Set("X-Records", fmt.Sprintf("%d", opts.total))
+
+		defer response.Body.Close()
+
+		return response, nil
+	}
+
+	httpmock.Reset()
+	httpmock.RegisterResponderWithQuery(http.MethodGet, hostURL+url, query, responder)
+}
+
 func setupURLErrorResponder(t *testing.T, url string) {
 	setupURLResponderWithStatusCode(t, url, http.StatusNotFound)
 }

pkg/puppetdb/nodes.go

@@ -2,6 +2,8 @@ package puppetdb
 
 import (
 	"fmt"
+	"io"
+	"math"
 	"strings"
 )
 
@@ -17,6 +19,39 @@ func (c *Client) Nodes(query string, pagination *Pagination, orderBy *OrderBy) (
 	return payload, err
 }
 
+// PaginatedNodes works just like Nodes, but returns a NodesCursor that
+// provides methods for iterating over N pages of nodes and calculates page
+// information for tracking progress. If pagination is nil, then a default
+// configuration with a limit of 100 is used instead.
+func (c *Client) PaginatedNodes(query string, pagination *Pagination, orderBy *OrderBy) (*NodesCursor, error) {
+	if pagination == nil {
+		pagination = &Pagination{Limit: 100}
+	}
+
+	tempPagination := Pagination{
+		Limit:        1,
+		IncludeTotal: true,
+	}
+
+	// make a call to pdb for 1 node to fetch the total number of nodes for
+	// page calculations in the cursor.
+	if _, err := c.Nodes(query, &tempPagination, orderBy); err != nil {
+		return nil, fmt.Errorf("failed to get node total from pdb: %w", err)
+	}
+
+	pagination.Total = tempPagination.Total
+	pagination.IncludeTotal = true
+
+	nc := &NodesCursor{
+		client:     c,
+		pagination: pagination,
+		query:      query,
+		orderBy:    orderBy,
+	}
+
+	return nc, nil
+}
+
 // Node will return a single node by certname
 func (c *Client) Node(certname string) (*Node, error) {
 	payload := &Node{}
@@ -56,3 +91,57 @@ type Node struct {
 	LatestReportStatus           string      `json:"latest_report_status"`
 	Count                        int         `json:"count"`
 }
+
+// NodesCursor is a pagination cursor that provides convenience methods for
+// stepping through pages of nodes.
+type NodesCursor struct {
+	client      *Client
+	pagination  *Pagination
+	query       string
+	orderBy     *OrderBy
+	currentPage []Node
+}
+
+// Next returns a page of nodes and iterates the pagination cursor by the
+// offset. If there are no more results left, the error will be io.EOF.
+func (nc *NodesCursor) Next() ([]Node, error) {
+	// this block increases the offset and checks of it's greater than or equal
+	// to the total only if we have already returned a first page.
+	if nc.currentPage != nil {
+		nc.pagination.Offset = nc.pagination.Offset + nc.pagination.Limit
+
+		if nc.pagination.Offset >= nc.pagination.Total {
+			return []Node{}, io.EOF
+		}
+	}
+
+	var err error
+
+	nc.currentPage, err = nc.client.Nodes(nc.query, nc.pagination, nc.orderBy)
+	if err != nil {
+		return nil, fmt.Errorf("client call for Nodes returned an error: %w", err)
+	}
+
+	if nc.CurrentPage() == nc.TotalPages() {
+		err = io.EOF
+	}
+
+	return nc.currentPage, err
+}
+
+// TotalPages returns the total number of pages that can returns nodes.
+func (nc *NodesCursor) TotalPages() int {
+	pagesf := float64(nc.pagination.Total) / float64(nc.pagination.Limit)
+	pages := int(math.Ceil(pagesf))
+
+	return pages
+}
+
+// CurrentPage returns the current page number the cursor is at.
+func (nc *NodesCursor) CurrentPage() int {
+	if nc.pagination.Offset == 0 {
+		return 1
+	}
+
+	return nc.pagination.Offset/nc.pagination.Limit + 1
+}

pkg/puppetdb/nodes_test.go

@@ -1,6 +1,7 @@
 package puppetdb
 
 import (
+	"io"
 	"strings"
 	"testing"
 
@@ -22,6 +23,47 @@ func TestNodes(t *testing.T) {
 	require.Equal(t, expectedNodes, actual)
 }
 
+func TestPaginatedNodes(t *testing.T) {
+	pagination := Pagination{
+		Limit:        5,
+		Offset:       0,
+		IncludeTotal: true,
+	}
+
+	setupPaginatedGetResponder(t, "/pdb/query/v4/nodes", "", mockPaginatedGetOptions{
+		limit: pagination.Limit,
+		total: 10,
+		pageFilenames: []string{
+			"nodes-page-1-response.json",
+			"nodes-page-2-response.json",
+		},
+	})
+
+	cursor, err := pdbClient.PaginatedNodes("", &pagination, nil)
+	require.NoError(t, err)
+	require.Equal(t, 2, cursor.TotalPages())
+	require.Equal(t, 1, cursor.CurrentPage())
+
+	actual, err := cursor.Next()
+	require.NoError(t, err)
+	require.Len(t, actual, 5)
+	require.Equal(t, "1.delivery.puppetlabs.net", actual[0].Certname)
+
+	{ // page 2 (last page)
+		actual, err := cursor.Next()
+		require.ErrorIs(t, err, io.EOF)
+		require.Equal(t, 2, cursor.CurrentPage())
+		require.Len(t, actual, 5)
+		require.Equal(t, "6.delivery.puppetlabs.net", actual[0].Certname)
+	}
+
+	{
+		actual, err := cursor.Next()
+		require.Len(t, actual, 0)
+		require.ErrorIs(t, err, io.EOF)
+	}
+}
+
 func TestNode(t *testing.T) {
 	nodeFooURL := strings.ReplaceAll(node, "{certname}", "foo")
 

pkg/puppetdb/testdata/nodes-page-1-response.json

@@ -0,0 +1,92 @@
+[
+  {
+    "deactivated": null,
+    "latest_report_hash": "7ccb6fb17b3fe11cecffe00b43b44f3776bcb89d",
+    "facts_environment": "production",
+    "cached_catalog_status": "not_used",
+    "report_environment": "production",
+    "latest_report_corrective_change": false,
+    "catalog_environment": "production",
+    "facts_timestamp": "2020-03-20T10:17:30.394Z",
+    "latest_report_noop": false,
+    "expired": null,
+    "latest_report_noop_pending": false,
+    "report_timestamp": "2020-03-20T10:17:54.470Z",
+    "certname": "1.delivery.puppetlabs.net",
+    "catalog_timestamp": "2020-03-20T10:17:33.991Z",
+    "latest_report_job_id": "1",
+    "latest_report_status": "changed"
+  },
+  {
+    "deactivated": null,
+    "latest_report_hash": null,
+    "facts_environment": "production",
+    "cached_catalog_status": null,
+    "report_environment": null,
+    "latest_report_corrective_change": null,
+    "catalog_environment": null,
+    "facts_timestamp": "2020-03-20T10:10:28.949Z",
+    "latest_report_noop": null,
+    "expired": null,
+    "latest_report_noop_pending": null,
+    "report_timestamp": null,
+    "certname": "2.delivery.puppetlabs.net",
+    "catalog_timestamp": null,
+    "latest_report_job_id": null,
+    "latest_report_status": null
+  },
+  {
+    "deactivated": null,
+    "latest_report_hash": "7ccb6fb17b3fe11cecffe00b43b44f3776bcb89d",
+    "facts_environment": "production",
+    "cached_catalog_status": "not_used",
+    "report_environment": "production",
+    "latest_report_corrective_change": false,
+    "catalog_environment": "production",
+    "facts_timestamp": "2020-03-20T10:17:30.394Z",
+    "latest_report_noop": false,
+    "expired": null,
+    "latest_report_noop_pending": false,
+    "report_timestamp": "2020-03-20T10:17:54.470Z",
+    "certname": "3.delivery.puppetlabs.net",
+    "catalog_timestamp": "2020-03-20T10:17:33.991Z",
+    "latest_report_job_id": "1",
+    "latest_report_status": "changed"
+  },
+  {
+    "deactivated": null,
+    "latest_report_hash": null,
+    "facts_environment": "production",
+    "cached_catalog_status": null,
+    "report_environment": null,
+    "latest_report_corrective_change": null,
+    "catalog_environment": null,
+    "facts_timestamp": "2020-03-20T10:10:28.949Z",
+    "latest_report_noop": null,
+    "expired": null,
+    "latest_report_noop_pending": null,
+    "report_timestamp": null,
+    "certname": "4.delivery.puppetlabs.net",
+    "catalog_timestamp": null,
+    "latest_report_job_id": null,
+    "latest_report_status": null
+  },
+  {
+    "deactivated": null,
+    "latest_report_hash": null,
+    "facts_environment": "production",
+    "cached_catalog_status": null,
+    "report_environment": null,
+    "latest_report_corrective_change": null,
+    "catalog_environment": null,
+    "facts_timestamp": "2020-03-20T10:10:28.949Z",
+    "latest_report_noop": null,
+    "expired": null,
+    "latest_report_noop_pending": null,
+    "report_timestamp": null,
+    "certname": "5.delivery.puppetlabs.net",
+    "catalog_timestamp": null,
+    "latest_report_job_id": null,
+    "latest_report_status": null
+  }
+]