From 754ba7891dc2f83c5fc2e4b60b6bfc3c8beb2c31 Mon Sep 17 00:00:00 2001 From: Kyle Terry Date: Wed, 5 Apr 2023 14:33:09 -0700 Subject: [PATCH 1/3] pbd client: adds nodes pagination cursor To make iterating over pages of nodes easier and consistent, I added a pagination cursor that allows a user of the nodes client to call NodesCursor.Next() to keep fetching pages until an io.EOF error is returned. This keeps the client use minimal of implementing packages (such as scarp) and centralizes and hides the pagination state/logic so it doesn't need to be implemented over and over again. I think we should generalize this a bit more and add it to all the pdb calls we export in this package as clients. I'm not sure if we can generalize it for the entire go-pe-client tree because other pagination schemes across PE might be different than pdb's API. --- pkg/puppetdb/common_test.go | 53 +++++++++++ pkg/puppetdb/nodes.go | 81 ++++++++++++++++ pkg/puppetdb/nodes_test.go | 42 +++++++++ .../testdata/nodes-page-1-response.json | 92 +++++++++++++++++++ .../testdata/nodes-page-2-response.json | 92 +++++++++++++++++++ 5 files changed, 360 insertions(+) create mode 100644 pkg/puppetdb/testdata/nodes-page-1-response.json create mode 100644 pkg/puppetdb/testdata/nodes-page-2-response.json diff --git a/pkg/puppetdb/common_test.go b/pkg/puppetdb/common_test.go index 5a04b0c..f21db60 100644 --- a/pkg/puppetdb/common_test.go +++ b/pkg/puppetdb/common_test.go @@ -1,9 +1,12 @@ package puppetdb import ( + "fmt" "io/ioutil" "net/http" "net/url" + "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 := ioutil.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..d4e409f 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,31 @@ func (c *Client) Nodes(query string, pagination *Pagination, orderBy *OrderBy) ( return payload, err } +func (c *Client) PaginatedNodes(query string, pagination *Pagination, orderBy *OrderBy) (*NodesCursor, error) { + tp := 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, &tp, orderBy); err != nil { + return nil, fmt.Errorf("failed to get node total from pdb: %w", err) + } + + pagination.Total = tp.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 +83,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 + } +] From 7ea9258fb22856b9d0bf62bc81d4fa597f13101e Mon Sep 17 00:00:00 2001 From: Kyle Terry Date: Wed, 5 Apr 2023 14:54:49 -0700 Subject: [PATCH 2/3] pdb client: adds default pagination if param is nil If the pagination parameter passed to client.PaginatedNodes is nil, then we create a default one with a limit of 100 set. --- pkg/puppetdb/nodes.go | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/pkg/puppetdb/nodes.go b/pkg/puppetdb/nodes.go index d4e409f..37bf90c 100644 --- a/pkg/puppetdb/nodes.go +++ b/pkg/puppetdb/nodes.go @@ -19,19 +19,27 @@ 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) { - tp := Pagination{ + 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, &tp, orderBy); err != nil { + if _, err := c.Nodes(query, &tempPagination, orderBy); err != nil { return nil, fmt.Errorf("failed to get node total from pdb: %w", err) } - pagination.Total = tp.Total + pagination.Total = tempPagination.Total pagination.IncludeTotal = true nc := &NodesCursor{ From 40e62bb6b40c80777736697078264052d4f602be Mon Sep 17 00:00:00 2001 From: Kyle Terry Date: Wed, 5 Apr 2023 15:42:18 -0700 Subject: [PATCH 3/3] pdb client: removes call to deprecated ioutil package --- pkg/puppetdb/common_test.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pkg/puppetdb/common_test.go b/pkg/puppetdb/common_test.go index 9b73aea..ef4447c 100644 --- a/pkg/puppetdb/common_test.go +++ b/pkg/puppetdb/common_test.go @@ -2,7 +2,6 @@ package puppetdb import ( "fmt" - "io/ioutil" "net/http" "net/url" "os" @@ -45,7 +44,7 @@ func setupPaginatedGetResponder(t *testing.T, url, query string, opts mockPagina var pages [][]byte for _, pfn := range opts.pageFilenames { - responseBody, err := ioutil.ReadFile(filepath.Join("testdata", pfn)) + responseBody, err := os.ReadFile(filepath.Join("testdata", pfn)) require.NoError(t, err) pages = append(pages, responseBody)