From 9622e515a3948c3955298dd6bb56be058cd30114 Mon Sep 17 00:00:00 2001
From: Anna Khmelnitsky <akhmelnitsky@vmware.com>
Date: Mon, 8 Mar 2021 19:46:42 +0000
Subject: [PATCH] Support pagination in policy group data source

---
 nsxt/data_source_nsxt_policy_group.go | 32 +++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/nsxt/data_source_nsxt_policy_group.go b/nsxt/data_source_nsxt_policy_group.go
index d978d9570..65b67fda9 100644
--- a/nsxt/data_source_nsxt_policy_group.go
+++ b/nsxt/data_source_nsxt_policy_group.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/vmware/vsphere-automation-sdk-go/runtime/protocol/client"
 	"github.com/vmware/vsphere-automation-sdk-go/services/nsxt/infra/domains"
 	"github.com/vmware/vsphere-automation-sdk-go/services/nsxt/model"
 )
@@ -26,6 +27,32 @@ func dataSourceNsxtPolicyGroup() *schema.Resource {
 	}
 }
 
+func listPolicyGroups(domain string, connector *client.RestConnector) ([]model.Group, error) {
+	// Local Manager only
+	client := domains.NewDefaultGroupsClient(connector)
+
+	var results []model.Group
+	var cursor *string
+	total := 0
+
+	for {
+		groups, err := client.List(domain, cursor, nil, nil, nil, nil, nil, nil)
+		if err != nil {
+			return results, err
+		}
+		results = append(results, groups.Results...)
+		if total == 0 && groups.ResultCount != nil {
+			// first response
+			total = int(*groups.ResultCount)
+		}
+
+		cursor = groups.Cursor
+		if len(results) >= total {
+			return results, nil
+		}
+	}
+}
+
 func dataSourceNsxtPolicyGroupRead(d *schema.ResourceData, m interface{}) error {
 	if isPolicyGlobalManager(m) {
 		domain := d.Get("domain").(string)
@@ -56,14 +83,15 @@ func dataSourceNsxtPolicyGroupRead(d *schema.ResourceData, m interface{}) error
 		return fmt.Errorf("Error obtaining Group ID or name during read")
 	} else {
 		// Get by full name/prefix
-		objList, err := client.List(domain, nil, nil, nil, nil, nil, nil, nil)
+		objList, err := listPolicyGroups(domain, connector)
 		if err != nil {
 			return handleListError("Group", err)
 		}
+
 		// go over the list to find the correct one (prefer a perfect match. If not - prefix match)
 		var perfectMatch []model.Group
 		var prefixMatch []model.Group
-		for _, objInList := range objList.Results {
+		for _, objInList := range objList {
 			if strings.HasPrefix(*objInList.DisplayName, objName) {
 				prefixMatch = append(prefixMatch, objInList)
 			}