diff --git a/pkg/puppetdb/common_test.go b/pkg/puppetdb/common_test.go index 77f6cca..ef4447c 100644 --- a/pkg/puppetdb/common_test.go +++ b/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) } diff --git a/pkg/puppetdb/nodes.go b/pkg/puppetdb/nodes.go index ed29b01..37bf90c 100644 --- a/pkg/puppetdb/nodes.go +++ b/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 +} diff --git a/pkg/puppetdb/nodes_test.go b/pkg/puppetdb/nodes_test.go index 84b3151..567da08 100644 --- a/pkg/puppetdb/nodes_test.go +++ b/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") diff --git a/pkg/puppetdb/testdata/nodes-page-1-response.json b/pkg/puppetdb/testdata/nodes-page-1-response.json new file mode 100644 index 0000000..44a0108 --- /dev/null +++ b/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 + } +] diff --git a/pkg/puppetdb/testdata/nodes-page-2-response.json b/pkg/puppetdb/testdata/nodes-page-2-response.json new file mode 100644 index 0000000..872f171 --- /dev/null +++ b/pkg/puppetdb/testdata/nodes-page-2-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": "6.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": "7.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": "8.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": "9.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": "10.delivery.puppetlabs.net", + "catalog_timestamp": null, + "latest_report_job_id": null, + "latest_report_status": null + } +]