feat: Paginate over prices sheets

erezrokah · erezrokah · commit 3424606ad47d · 2023-04-20T19:06:23.000+03:00
diff --git a/plugins/source/azure/resources/services/consumption/subscription_price_sheets.go b/plugins/source/azure/resources/services/consumption/subscription_price_sheets.go
@@ -2,6 +2,7 @@ package consumption
 
 import (
 	"context"
+	"net/url"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/consumption/armconsumption"
@@ -26,10 +27,37 @@ func fetchSubscriptionPriceSheets(ctx context.Context, meta schema.ClientMeta, p
 	if err != nil {
 		return err
 	}
-	resp, err := svc.Get(ctx, &armconsumption.PriceSheetClientGetOptions{Expand: to.Ptr("properties/meterDetails")})
+	opts := &armconsumption.PriceSheetClientGetOptions{Expand: to.Ptr("properties/meterDetails")}
+	resp, err := svc.Get(ctx, opts)
 	if err != nil {
 		return err
 	}
+
+	// This is a workaround to get all price sheets
+	// Somehow related to https://github.com/Azure/azure-sdk-for-go/issues/4142
+	allPricesSheets := resp.PriceSheetResult.Properties.Pricesheets
+	nextLink := resp.PriceSheetResult.Properties.NextLink
+	for nextLink != nil {
+		parsedNextLink, err := url.Parse(*nextLink)
+		if err != nil {
+			cl.Logger().Warn().Err(err).Msgf("failed to parse next link: %q", *nextLink)
+			break
+		}
+		token := parsedNextLink.Query().Get("skiptoken")
+		if token == "" {
+			cl.Logger().Warn().Msgf("failed to get skiptoken from next link: %q", *nextLink)
+			break
+		}
+		opts.Skiptoken = to.Ptr(token)
+		paginatedResponse, err := svc.Get(ctx, opts)
+		if err != nil {
+			cl.Logger().Warn().Err(err).Msgf("failed to get paginated response for next link: %q", *nextLink)
+			break
+		}
+		allPricesSheets = append(allPricesSheets, paginatedResponse.PriceSheetResult.Properties.Pricesheets...)
+	}
+
+	resp.PriceSheetResult.Properties.Pricesheets = allPricesSheets
 	res <- resp.PriceSheetResult
 	return nil
 }
diff --git a/plugins/source/azure/resources/services/consumption/subscription_price_sheets_test.go b/plugins/source/azure/resources/services/consumption/subscription_price_sheets_test.go
@@ -17,6 +17,7 @@ func createSubscriptionPriceSheets(router *mux.Router) error {
 	if err := faker.FakeObject(&resp); err != nil {
 		return err
 	}
+	resp.PriceSheetResult.Properties.NextLink = nil
 
 	router.HandleFunc("/subscriptions/{subscriptionId}/providers/Microsoft.Consumption/pricesheets/default", func(w http.ResponseWriter, r *http.Request) {
 		b, err := json.Marshal(&resp)

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@ func createSubscriptionPriceSheets(router *mux.Router) error {`
`17`	`17`	`if err := faker.FakeObject(&resp); err != nil {`
`18`	`18`	`return err`
`19`	`19`	`}`
	`20`	`+ resp.PriceSheetResult.Properties.NextLink = nil`
`20`	`21`
`21`	`22`	`router.HandleFunc("/subscriptions/{subscriptionId}/providers/Microsoft.Consumption/pricesheets/default", func(w http.ResponseWriter, r *http.Request) {`
`22`	`23`	`b, err := json.Marshal(&resp)`