Skip to content

Commit d9e2bb7

Browse files
kyleterrykyleterry
authored
pdb client: generalizes the pagination cursor (#83)
1 parent ebdd3c5 commit d9e2bb7

File tree

6 files changed

+313
-66
lines changed

6 files changed

+313
-66
lines changed

‎pkg/puppetdb/facts.go‎

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

3+
import (
4+
"errors"
5+
"fmt"
6+
"io"
7+
)
8+
39
const (
410
factNames = "/pdb/query/v4/fact-names"
511
factPaths = "/pdb/query/v4/fact-paths"
@@ -28,6 +34,19 @@ func (c *Client) Facts(query string, pagination *Pagination, orderBy *OrderBy) (
2834
return payload, err
2935
}
3036

37+
func (c *Client) PaginatedFacts(query string, pagination *Pagination, orderBy *OrderBy) (*FactsCursor, error) {
38+
pc, err := newPageCursor(c, facts, query, pagination, orderBy)
39+
if err != nil {
40+
return nil, fmt.Errorf("failed to initialize page cursor: %w", err)
41+
}
42+
43+
cursor := FactsCursor{
44+
pageCursor: pc,
45+
}
46+
47+
return &cursor, nil
48+
}
49+
3150
// FactContents will return all facts matching the given query on the fact-contents endpoint. Facts for deactivated nodes are not included in the response.
3251
// - https://puppet.com/docs/puppetdb/latest/api/query/v4/fact-contents.html
3352
func (c *Client) FactContents(query string, pagination *Pagination, orderBy *OrderBy) ([]Fact, error) {
@@ -60,3 +79,17 @@ type FactPath struct {
6079
Type string `json:"type"`
6180
Count int `json:"count"`
6281
}
82+
83+
type FactsCursor struct {
84+
*pageCursor
85+
}
86+
87+
func (fc FactsCursor) Next() ([]Fact, error) {
88+
payload := []Fact{}
89+
err := fc.next(&payload)
90+
if err != nil && !errors.Is(err, io.EOF) {
91+
return nil, err
92+
}
93+
94+
return payload, err
95+
}

‎pkg/puppetdb/facts_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
"testing"
56

67
"github.com/stretchr/testify/require"
@@ -34,6 +35,47 @@ func TestFacts(t *testing.T) {
3435
require.Equal(t, expectedFacts, actual)
3536
}
3637

38+
func TestPaginatedFacts(t *testing.T) {
39+
pagination := Pagination{
40+
Limit: 10,
41+
Offset: 0,
42+
IncludeTotal: true,
43+
}
44+
45+
setupPaginatedGetResponder(t, facts, "", mockPaginatedGetOptions{
46+
limit: pagination.Limit,
47+
total: 20,
48+
pageFilenames: []string{
49+
"facts-page-1-response.json",
50+
"facts-page-2-response.json",
51+
},
52+
})
53+
54+
cursor, err := pdbClient.PaginatedFacts("", &pagination, nil)
55+
require.NoError(t, err)
56+
require.Equal(t, 2, cursor.TotalPages())
57+
require.Equal(t, 1, cursor.CurrentPage())
58+
59+
actual, err := cursor.Next()
60+
require.NoError(t, err)
61+
require.Len(t, actual, 10)
62+
require.Equal(t, "1.delivery.puppetlabs.net", actual[0].Certname)
63+
64+
{ // page 2 (last page)
65+
actual, err := cursor.Next()
66+
require.ErrorIs(t, err, io.EOF)
67+
require.Equal(t, 2, cursor.CurrentPage())
68+
require.Len(t, actual, 10)
69+
require.Equal(t, "6.delivery.puppetlabs.net", actual[0].Certname)
70+
}
71+
72+
{
73+
actual, err := cursor.Next()
74+
require.Len(t, actual, 0)
75+
require.ErrorIs(t, err, io.EOF)
76+
}
77+
}
78+
3779
// TestFactContents performs a test on the FactContents method, and verified the expected response is returned,
3880
func TestFactContents(t *testing.T) {
3981
// Test with query

‎pkg/puppetdb/nodes.go‎

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

33
import (
4+
"errors"
45
"fmt"
56
"io"
6-
"math"
77
"strings"
88
)
99

@@ -24,32 +24,16 @@ func (c *Client) Nodes(query string, pagination *Pagination, orderBy *OrderBy) (
2424
// information for tracking progress. If pagination is nil, then a default
2525
// configuration with a limit of 100 is used instead.
2626
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)
27+
pc, err := newPageCursor(c, nodes, query, pagination, orderBy)
28+
if err != nil {
29+
return nil, fmt.Errorf("failed to initialize page cursor: %w", err)
4030
}
4131

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,
32+
cursor := NodesCursor{
33+
pageCursor: pc,
5034
}
5135

52-
return nc, nil
36+
return &cursor, nil
5337
}
5438

5539
// Node will return a single node by certname
@@ -95,53 +79,17 @@ type Node struct {
9579
// NodesCursor is a pagination cursor that provides convenience methods for
9680
// stepping through pages of nodes.
9781
type NodesCursor struct {
98-
client *Client
99-
pagination *Pagination
100-
query string
101-
orderBy *OrderBy
102-
currentPage []Node
82+
*pageCursor
10383
}
10484

10585
// Next returns a page of nodes and iterates the pagination cursor by the
10686
// offset. If there are no more results left, the error will be io.EOF.
10787
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
88+
payload := []Node{}
89+
err := nc.next(&payload)
90+
if err != nil && !errors.Is(err, io.EOF) {
91+
return nil, err
14492
}
14593

146-
return nc.pagination.Offset/nc.pagination.Limit + 1
94+
return payload, err
14795
}

‎pkg/puppetdb/pagination.go‎

Lines changed: 101 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
package puppetdb
22

3-
import "strconv"
3+
import (
4+
"fmt"
5+
"io"
6+
"math"
7+
"strconv"
8+
)
49

510
// toParams will take the Pagination struct and convert into a form Client SetQueryParam accepts
611
func (p Pagination) toParams() map[string]string {
@@ -25,3 +30,98 @@ type Pagination struct {
2530
IncludeTotal bool
2631
Total int
2732
}
33+
34+
func NewDefaultPagination() *Pagination {
35+
return &Pagination{
36+
Limit: 100,
37+
IncludeTotal: true,
38+
}
39+
}
40+
41+
// pageCursor is a pagination cursor that provides convenience methods for
42+
// stepping through pbd response pages.
43+
type pageCursor struct {
44+
client *Client
45+
path string
46+
query string
47+
pagination *Pagination
48+
orderBy *OrderBy
49+
currentPage any
50+
}
51+
52+
// next makes a request for a page of results and iterates the pagination
53+
// cursor by the offset. If there are no more results left, the error will be
54+
// io.EOF.
55+
func (pc *pageCursor) next(response any) error {
56+
// this block increases the offset and checks of it's greater than or equal
57+
// to the total only if we have already returned a first page.
58+
if pc.currentPage != nil {
59+
pc.pagination.Offset = pc.pagination.Offset + pc.pagination.Limit
60+
61+
if pc.pagination.Offset >= pc.pagination.Total {
62+
return io.EOF
63+
}
64+
}
65+
66+
var err error
67+
68+
err = getRequest(pc.client, pc.path, pc.query, pc.pagination, pc.orderBy, response)
69+
if err != nil {
70+
return fmt.Errorf("page cursor: client call returned an error: %w", err)
71+
}
72+
73+
pc.currentPage = response
74+
75+
if pc.CurrentPage() == pc.TotalPages() {
76+
err = io.EOF
77+
}
78+
79+
return err
80+
}
81+
82+
// TotalPages returns the total number of pages that can returns nodes.
83+
func (pc *pageCursor) TotalPages() int {
84+
pagesf := float64(pc.pagination.Total) / float64(pc.pagination.Limit)
85+
pages := int(math.Ceil(pagesf))
86+
87+
return pages
88+
}
89+
90+
// CurrentPage returns the current page number the cursor is at.
91+
func (pc *pageCursor) CurrentPage() int {
92+
if pc.pagination.Offset == 0 {
93+
return 1
94+
}
95+
96+
return pc.pagination.Offset/pc.pagination.Limit + 1
97+
}
98+
99+
func newPageCursor(c *Client, path, query string, p *Pagination, orderBy *OrderBy) (*pageCursor, error) {
100+
if p == nil {
101+
p = NewDefaultPagination()
102+
}
103+
104+
tempPagination := Pagination{
105+
Limit: 1,
106+
IncludeTotal: true,
107+
}
108+
109+
// make a call to pdb for 1 object to fetch the total number of results for
110+
// page calculations in the cursor.
111+
if err := getRequest(c, path, query, &tempPagination, orderBy, &[]any{}); err != nil {
112+
return nil, fmt.Errorf("failed to get result total from pdb: %w", err)
113+
}
114+
115+
p.Total = tempPagination.Total
116+
p.IncludeTotal = true
117+
118+
pc := &pageCursor{
119+
client: c,
120+
path: path,
121+
query: query,
122+
pagination: p,
123+
orderBy: orderBy,
124+
}
125+
126+
return pc, nil
127+
}

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
[
2+
{
3+
"name": "os",
4+
"certname": "1.delivery.puppetlabs.net",
5+
"value": "Debian",
6+
"environment": "production"
7+
},
8+
{
9+
"name": "uptime_days",
10+
"certname": "1.delivery.puppetlabs.net",
11+
"value": "1000000 days",
12+
"environment": "production"
13+
},
14+
{
15+
"name": "os",
16+
"certname": "2.delivery.puppetlabs.net",
17+
"value": "Debian",
18+
"environment": "production"
19+
},
20+
{
21+
"name": "uptime_days",
22+
"certname": "2.delivery.puppetlabs.net",
23+
"value": "1000000 days",
24+
"environment": "production"
25+
},
26+
{
27+
"name": "os",
28+
"certname": "3.delivery.puppetlabs.net",
29+
"value": "Debian",
30+
"environment": "production"
31+
},
32+
{
33+
"name": "uptime_days",
34+
"certname": "3.delivery.puppetlabs.net",
35+
"value": "1000000 days",
36+
"environment": "production"
37+
},
38+
{
39+
"name": "os",
40+
"certname": "4.delivery.puppetlabs.net",
41+
"value": "Debian",
42+
"environment": "production"
43+
},
44+
{
45+
"name": "uptime_days",
46+
"certname": "4.delivery.puppetlabs.net",
47+
"value": "1000000 days",
48+
"environment": "production"
49+
},
50+
{
51+
"name": "os",
52+
"certname": "5.delivery.puppetlabs.net",
53+
"value": "Debian",
54+
"environment": "production"
55+
},
56+
{
57+
"name": "uptime_days",
58+
"certname": "5.delivery.puppetlabs.net",
59+
"value": "1000000 days",
60+
"environment": "production"
61+
}
62+
]

0 commit comments

Comments
 (0)