From be32f3d3171b7f03604328e8d39b2b6ee0aea14d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=83=E6=9E=AB?= <qianfeng.syf@alibaba-inc.com>
Date: Mon, 28 Oct 2024 17:12:14 +0800
Subject: [PATCH] feat: add pagination support for listing feature entities

---
 api/api_feature_entity.go            |  6 ++++--
 api/response_feature_entity.go       |  1 +
 featurestore/feature_store_client.go | 31 ++++++++++++++++++----------
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/api/api_feature_entity.go b/api/api_feature_entity.go
index fb206de..9befa24 100644
--- a/api/api_feature_entity.go
+++ b/api/api_feature_entity.go
@@ -14,19 +14,21 @@ FeatureEntityApiService List FeatureEntities
 
 @return InlineResponse20079
 */
-func (a *FeatureEntityApiService) ListFeatureEntities(projectId string) (ListFeatureEntitiesResponse, error) {
+func (a *FeatureEntityApiService) ListFeatureEntities(pagesize, pagenumber int32, projectId string) (ListFeatureEntitiesResponse, error) {
 	var (
 		localVarReturnValue ListFeatureEntitiesResponse
 	)
 	request := paifeaturestore.ListFeatureEntitiesRequest{}
 	request.SetProjectId(projectId)
-	request.SetPageSize(100)
+	request.SetPageSize(pagesize)
+	request.SetPageNumber(pagenumber)
 
 	response, err := a.client.ListFeatureEntities(&a.client.instanceId, &request)
 	if err != nil {
 		return localVarReturnValue, err
 	}
 
+	localVarReturnValue.TotalCount = int(*response.Body.TotalCount)
 	var featureEntities []*FeatureEntity
 
 	for _, entity := range response.Body.FeatureEntities {
diff --git a/api/response_feature_entity.go b/api/response_feature_entity.go
index b40e5f9..bb04d7b 100644
--- a/api/response_feature_entity.go
+++ b/api/response_feature_entity.go
@@ -1,6 +1,7 @@
 package api
 
 type ListFeatureEntitiesResponse struct {
+	TotalCount      int `json:"total_count"`
 	FeatureEntities []*FeatureEntity
 }
 
diff --git a/featurestore/feature_store_client.go b/featurestore/feature_store_client.go
index 40fd10b..8152d75 100644
--- a/featurestore/feature_store_client.go
+++ b/featurestore/feature_store_client.go
@@ -249,23 +249,32 @@ func (c *FeatureStoreClient) LoadProjectData() error {
 		project := domain.NewProject(p, c.datasourceInitClient)
 		projectData[project.ProjectName] = project
 
+		var (
+			pagesize   = 100
+			pagenumber = 1
+		)
+
 		// get feature entities
-		listFeatureEntitiesResponse, err := c.client.FeatureEntityApi.ListFeatureEntities(strconv.Itoa(p.ProjectId))
-		if err != nil {
-			c.logError(fmt.Errorf("list feature entities error, err=%v", err))
-			return err
-		}
+		for {
+			listFeatureEntitiesResponse, err := c.client.FeatureEntityApi.ListFeatureEntities(int32(pagesize), int32(pagenumber), strconv.Itoa(p.ProjectId))
+			if err != nil {
+				c.logError(fmt.Errorf("list feature entities error, err=%v", err))
+				return err
+			}
 
-		for _, entity := range listFeatureEntitiesResponse.FeatureEntities {
-			if entity.ProjectId == project.ProjectId {
+			for _, entity := range listFeatureEntitiesResponse.FeatureEntities {
 				project.FeatureEntityMap[entity.FeatureEntityName] = domain.NewFeatureEntity(entity)
 			}
+
+			if len(listFeatureEntitiesResponse.FeatureEntities) == 0 || pagesize*pagenumber > listFeatureEntitiesResponse.TotalCount {
+				break
+			}
+
+			pagenumber++
+
 		}
 
-		var (
-			pagesize   = 100
-			pagenumber = 1
-		)
+		pagenumber = 1
 		// get feature views
 		for {
 			listFeatureViews, err := c.client.FeatureViewApi.ListFeatureViews(int32(pagesize), int32(pagenumber), strconv.Itoa(p.ProjectId))