Skip to content

Commit ebdd3c5

Browse files
kyleterrykyleterry
authored
pbd client: adds nodes pagination cursor (#82)
1 parent 07d5308 commit ebdd3c5

File tree

5 files changed

+368
-0
lines changed

5 files changed

+368
-0
lines changed

‎pkg/puppetdb/common_test.go‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package puppetdb
22

33
import (
4+
"fmt"
45
"net/http"
56
"net/url"
67
"os"
8+
"path/filepath"
9+
"strconv"
710
"testing"
811
"time"
912

@@ -31,6 +34,56 @@ func setupGetResponder(t *testing.T, url, query, responseFilename string) {
3134
response.Body.Close()
3235
}
3336

37+
type mockPaginatedGetOptions struct {
38+
limit int
39+
total int
40+
pageFilenames []string
41+
}
42+
43+
func setupPaginatedGetResponder(t *testing.T, url, query string, opts mockPaginatedGetOptions) {
44+
var pages [][]byte
45+
46+
for _, pfn := range opts.pageFilenames {
47+
responseBody, err := os.ReadFile(filepath.Join("testdata", pfn))
48+
require.NoError(t, err)
49+
50+
pages = append(pages, responseBody)
51+
}
52+
53+
responder := func(r *http.Request) (*http.Response, error) {
54+
var (
55+
offset int
56+
pageNum int
57+
err error
58+
)
59+
60+
offsetS := r.URL.Query().Get("offset")
61+
if offsetS != "" {
62+
offset, err = strconv.Atoi(offsetS)
63+
if err != nil {
64+
return nil, err
65+
}
66+
}
67+
68+
if offset > 0 {
69+
pageNum = offset / opts.limit
70+
}
71+
72+
responseBody := pages[pageNum]
73+
74+
response := httpmock.NewBytesResponse(http.StatusOK, responseBody)
75+
response.Header.Set("Content-Type", "application/json")
76+
response.Header.Set("X-Records", fmt.Sprintf("%d", opts.total))
77+
78+
defer response.Body.Close()
79+
80+
return response, nil
81+
}
82+
83+
httpmock.Reset()
84+
httpmock.RegisterResponderWithQuery(http.MethodGet, hostURL+url, query, responder)
85+
}
86+
3487
func setupURLErrorResponder(t *testing.T, url string) {
3588
setupURLResponderWithStatusCode(t, url, http.StatusNotFound)
3689
}

‎pkg/puppetdb/nodes.go‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ package puppetdb
22

33
import (
44
"fmt"
5+
"io"
6+
"math"
57
"strings"
68
)
79

@@ -17,6 +19,39 @@ func (c *Client) Nodes(query string, pagination *Pagination, orderBy *OrderBy) (
1719
return payload, err
1820
}
1921

22+
// PaginatedNodes works just like Nodes, but returns a NodesCursor that
23+
// provides methods for iterating over N pages of nodes and calculates page
24+
// information for tracking progress. If pagination is nil, then a default
25+
// configuration with a limit of 100 is used instead.
26+
func (c *Client) PaginatedNodes(query string, pagination *Pagination, orderBy *OrderBy) (*NodesCursor, error) {
27+
if pagination == nil {
28+
pagination = &Pagination{Limit: 100}
29+
}
30+
31+
tempPagination := Pagination{
32+
Limit: 1,
33+
IncludeTotal: true,
34+
}
35+
36+
// make a call to pdb for 1 node to fetch the total number of nodes for
37+
// page calculations in the cursor.
38+
if _, err := c.Nodes(query, &tempPagination, orderBy); err != nil {
39+
return nil, fmt.Errorf("failed to get node total from pdb: %w", err)
40+
}
41+
42+
pagination.Total = tempPagination.Total
43+
pagination.IncludeTotal = true
44+
45+
nc := &NodesCursor{
46+
client: c,
47+
pagination: pagination,
48+
query: query,
49+
orderBy: orderBy,
50+
}
51+
52+
return nc, nil
53+
}
54+
2055
// Node will return a single node by certname
2156
func (c *Client) Node(certname string) (*Node, error) {
2257
payload := &Node{}
@@ -56,3 +91,57 @@ type Node struct {
5691
LatestReportStatus string `json:"latest_report_status"`
5792
Count int `json:"count"`
5893
}
94+
95+
// NodesCursor is a pagination cursor that provides convenience methods for
96+
// stepping through pages of nodes.
97+
type NodesCursor struct {
98+
client *Client
99+
pagination *Pagination
100+
query string
101+
orderBy *OrderBy
102+
currentPage []Node
103+
}
104+
105+
// Next returns a page of nodes and iterates the pagination cursor by the
106+
// offset. If there are no more results left, the error will be io.EOF.
107+
func (nc *NodesCursor) Next() ([]Node, error) {
108+
// this block increases the offset and checks of it's greater than or equal
109+
// to the total only if we have already returned a first page.
110+
if nc.currentPage != nil {
111+
nc.pagination.Offset = nc.pagination.Offset + nc.pagination.Limit
112+
113+
if nc.pagination.Offset >= nc.pagination.Total {
114+
return []Node{}, io.EOF
115+
}
116+
}
117+
118+
var err error
119+
120+
nc.currentPage, err = nc.client.Nodes(nc.query, nc.pagination, nc.orderBy)
121+
if err != nil {
122+
return nil, fmt.Errorf("client call for Nodes returned an error: %w", err)
123+
}
124+
125+
if nc.CurrentPage() == nc.TotalPages() {
126+
err = io.EOF
127+
}
128+
129+
return nc.currentPage, err
130+
}
131+
132+
// TotalPages returns the total number of pages that can returns nodes.
133+
func (nc *NodesCursor) TotalPages() int {
134+
pagesf := float64(nc.pagination.Total) / float64(nc.pagination.Limit)
135+
pages := int(math.Ceil(pagesf))
136+
137+
return pages
138+
}
139+
140+
// CurrentPage returns the current page number the cursor is at.
141+
func (nc *NodesCursor) CurrentPage() int {
142+
if nc.pagination.Offset == 0 {
143+
return 1
144+
}
145+
146+
return nc.pagination.Offset/nc.pagination.Limit + 1
147+
}

‎pkg/puppetdb/nodes_test.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package puppetdb
22

33
import (
4+
"io"
45
"strings"
56
"testing"
67

@@ -22,6 +23,47 @@ func TestNodes(t *testing.T) {
2223
require.Equal(t, expectedNodes, actual)
2324
}
2425

26+
func TestPaginatedNodes(t *testing.T) {
27+
pagination := Pagination{
28+
Limit: 5,
29+
Offset: 0,
30+
IncludeTotal: true,
31+
}
32+
33+
setupPaginatedGetResponder(t, "/pdb/query/v4/nodes", "", mockPaginatedGetOptions{
34+
limit: pagination.Limit,
35+
total: 10,
36+
pageFilenames: []string{
37+
"nodes-page-1-response.json",
38+
"nodes-page-2-response.json",
39+
},
40+
})
41+
42+
cursor, err := pdbClient.PaginatedNodes("", &pagination, nil)
43+
require.NoError(t, err)
44+
require.Equal(t, 2, cursor.TotalPages())
45+
require.Equal(t, 1, cursor.CurrentPage())
46+
47+
actual, err := cursor.Next()
48+
require.NoError(t, err)
49+
require.Len(t, actual, 5)
50+
require.Equal(t, "1.delivery.puppetlabs.net", actual[0].Certname)
51+
52+
{ // page 2 (last page)
53+
actual, err := cursor.Next()
54+
require.ErrorIs(t, err, io.EOF)
55+
require.Equal(t, 2, cursor.CurrentPage())
56+
require.Len(t, actual, 5)
57+
require.Equal(t, "6.delivery.puppetlabs.net", actual[0].Certname)
58+
}
59+
60+
{
61+
actual, err := cursor.Next()
62+
require.Len(t, actual, 0)
63+
require.ErrorIs(t, err, io.EOF)
64+
}
65+
}
66+
2567
func TestNode(t *testing.T) {
2668
nodeFooURL := strings.ReplaceAll(node, "{certname}", "foo")
2769

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
[
2+
{
3+
"deactivated": null,
4+
"latest_report_hash": "7ccb6fb17b3fe11cecffe00b43b44f3776bcb89d",
5+
"facts_environment": "production",
6+
"cached_catalog_status": "not_used",
7+
"report_environment": "production",
8+
"latest_report_corrective_change": false,
9+
"catalog_environment": "production",
10+
"facts_timestamp": "2020-03-20T10:17:30.394Z",
11+
"latest_report_noop": false,
12+
"expired": null,
13+
"latest_report_noop_pending": false,
14+
"report_timestamp": "2020-03-20T10:17:54.470Z",
15+
"certname": "1.delivery.puppetlabs.net",
16+
"catalog_timestamp": "2020-03-20T10:17:33.991Z",
17+
"latest_report_job_id": "1",
18+
"latest_report_status": "changed"
19+
},
20+
{
21+
"deactivated": null,
22+
"latest_report_hash": null,
23+
"facts_environment": "production",
24+
"cached_catalog_status": null,
25+
"report_environment": null,
26+
"latest_report_corrective_change": null,
27+
"catalog_environment": null,
28+
"facts_timestamp": "2020-03-20T10:10:28.949Z",
29+
"latest_report_noop": null,
30+
"expired": null,
31+
"latest_report_noop_pending": null,
32+
"report_timestamp": null,
33+
"certname": "2.delivery.puppetlabs.net",
34+
"catalog_timestamp": null,
35+
"latest_report_job_id": null,
36+
"latest_report_status": null
37+
},
38+
{
39+
"deactivated": null,
40+
"latest_report_hash": "7ccb6fb17b3fe11cecffe00b43b44f3776bcb89d",
41+
"facts_environment": "production",
42+
"cached_catalog_status": "not_used",
43+
"report_environment": "production",
44+
"latest_report_corrective_change": false,
45+
"catalog_environment": "production",
46+
"facts_timestamp": "2020-03-20T10:17:30.394Z",
47+
"latest_report_noop": false,
48+
"expired": null,
49+
"latest_report_noop_pending": false,
50+
"report_timestamp": "2020-03-20T10:17:54.470Z",
51+
"certname": "3.delivery.puppetlabs.net",
52+
"catalog_timestamp": "2020-03-20T10:17:33.991Z",
53+
"latest_report_job_id": "1",
54+
"latest_report_status": "changed"
55+
},
56+
{
57+
"deactivated": null,
58+
"latest_report_hash": null,
59+
"facts_environment": "production",
60+
"cached_catalog_status": null,
61+
"report_environment": null,
62+
"latest_report_corrective_change": null,
63+
"catalog_environment": null,
64+
"facts_timestamp": "2020-03-20T10:10:28.949Z",
65+
"latest_report_noop": null,
66+
"expired": null,
67+
"latest_report_noop_pending": null,
68+
"report_timestamp": null,
69+
"certname": "4.delivery.puppetlabs.net",
70+
"catalog_timestamp": null,
71+
"latest_report_job_id": null,
72+
"latest_report_status": null
73+
},
74+
{
75+
"deactivated": null,
76+
"latest_report_hash": null,
77+
"facts_environment": "production",
78+
"cached_catalog_status": null,
79+
"report_environment": null,
80+
"latest_report_corrective_change": null,
81+
"catalog_environment": null,
82+
"facts_timestamp": "2020-03-20T10:10:28.949Z",
83+
"latest_report_noop": null,
84+
"expired": null,
85+
"latest_report_noop_pending": null,
86+
"report_timestamp": null,
87+
"certname": "5.delivery.puppetlabs.net",
88+
"catalog_timestamp": null,
89+
"latest_report_job_id": null,
90+
"latest_report_status": null
91+
}
92+
]

0 commit comments

Comments
 (0)