feat: add pagination support for listing feature entities

aliyun · Oct 28, 2024 · be32f3d · be32f3d
1 parent 2134bc4
commit be32f3d
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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))