Skip to content

Commit 3d2c51b

Browse files
ayanamistayanamist
authored
add NextToken support to #260 (#290)
using NextToken to paginate
1 parent 3d3cd3c commit 3d2c51b

File tree

2 files changed

+158
-29
lines changed

2 files changed

+158
-29
lines changed

‎openapi/pager.go‎

Lines changed: 61 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -39,22 +39,27 @@ var PagerFlag = &cli.Flag{Category: "caller",
3939
{Key: "PageNumber", DefaultValue: "PageNumber", Short: i18n.T(" PageNumber", "指定PageNumber的属性")},
4040
{Key: "PageSize", DefaultValue: "PageSize", Short: i18n.T("PageSize", "")},
4141
{Key: "TotalCount", DefaultValue: "TotalCount", Short: i18n.T("TotalCount", "")},
42+
{Key: "NextToken", DefaultValue: "NextToken", Short: i18n.T("NextToken", "")},
4243
},
4344
ExcludeWith: []string{WaiterFlag.Name},
4445
}
4546

4647
type Pager struct {
4748
PageNumberFlag string
4849
PageSizeFlag string
50+
NextTokenFlag string
4951

5052
PageNumberExpr string
5153
PageSizeExpr string
5254
TotalCountExpr string
55+
NextTokenExpr string
5356

5457
PageSize int
5558

5659
totalCount int
5760
currentPageNumber int
61+
nextTokenMode bool
62+
nextToken string
5863
collectionPath string
5964

6065
results []interface{}
@@ -75,6 +80,11 @@ func GetPager() *Pager {
7580
pager.PageSizeExpr, _ = PagerFlag.GetFieldValue("PageSize")
7681
pager.TotalCountExpr, _ = PagerFlag.GetFieldValue("TotalCount")
7782

83+
nextTokenFlagTemp, _ := PagerFlag.GetFieldValue("NextToken")
84+
tempStr = strings.Split(nextTokenFlagTemp, ".")
85+
pager.NextTokenFlag = tempStr[len(tempStr)-1]
86+
pager.NextTokenExpr = nextTokenFlagTemp
87+
7888
pager.collectionPath, _ = PagerFlag.GetFieldValue("path")
7989
return pager
8090
}
@@ -101,6 +111,9 @@ func (a *Pager) CallWith(invoker Invoker) (string, error) {
101111
}
102112

103113
func (a *Pager) HasMore() bool {
114+
if a.nextTokenMode {
115+
return a.nextToken != ""
116+
}
104117
pages := int(math.Ceil(float64(a.totalCount) / float64(a.PageSize)))
105118
return a.currentPageNumber < pages
106119
}
@@ -135,40 +148,56 @@ func (a *Pager) FeedResponse(body string) error {
135148
return fmt.Errorf("unmarshal %s", err.Error())
136149
}
137150

138-
if total, err := jmespath.Search(a.TotalCountExpr, j); err == nil {
139-
var totalCount float64
140-
if strCount, ok := total.(string); ok {
141-
totalCount, _ = strconv.ParseFloat(strCount, 64)
151+
if a.nextTokenMode {
152+
a.nextToken = ""
153+
}
154+
if a.NextTokenExpr != "" {
155+
// allow to ignore NextToken mode
156+
if val, err := jmespath.Search(a.NextTokenExpr, j); err == nil {
157+
if nextToken, ok := val.(string); ok {
158+
a.nextToken = nextToken
159+
a.nextTokenMode = true
160+
}
142161
} else {
143-
totalCount = total.(float64)
162+
return fmt.Errorf("jmespath: '%s' failed %s", a.NextTokenExpr, err)
144163
}
145-
a.totalCount = int(totalCount)
146-
} else {
147-
return fmt.Errorf("jmespath: '%s' failed %s", a.TotalCountExpr, err)
148164
}
165+
if !a.nextTokenMode {
166+
if total, err := jmespath.Search(a.TotalCountExpr, j); err == nil {
167+
var totalCount float64
168+
if strCount, ok := total.(string); ok {
169+
totalCount, _ = strconv.ParseFloat(strCount, 64)
170+
} else {
171+
totalCount = total.(float64)
172+
}
173+
a.totalCount = int(totalCount)
174+
} else {
175+
return fmt.Errorf("jmespath: '%s' failed %s", a.TotalCountExpr, err)
176+
}
149177

150-
if pageNumber, err := jmespath.Search(a.PageNumberExpr, j); err == nil {
151-
var currentPageNumber float64
152-
if strpageNumber, ok := pageNumber.(string); ok {
153-
currentPageNumber, _ = strconv.ParseFloat(strpageNumber, 64)
178+
if pageNumber, err := jmespath.Search(a.PageNumberExpr, j); err == nil {
179+
var currentPageNumber float64
180+
if strpageNumber, ok := pageNumber.(string); ok {
181+
currentPageNumber, _ = strconv.ParseFloat(strpageNumber, 64)
182+
} else {
183+
currentPageNumber = pageNumber.(float64)
184+
}
185+
a.currentPageNumber = int(currentPageNumber)
154186
} else {
155-
currentPageNumber = pageNumber.(float64)
187+
return fmt.Errorf("jmespath: '%s' failed %s", a.PageNumberExpr, err)
156188
}
157-
a.currentPageNumber = int(currentPageNumber)
158-
} else {
159-
return fmt.Errorf("jmespath: '%s' failed %s", a.PageNumberExpr, err)
160-
}
161189

162-
if pageSize, err := jmespath.Search(a.PageSizeExpr, j); err == nil {
163-
var PageSize float64
164-
if strpageSize, ok := pageSize.(string); ok {
165-
PageSize, _ = strconv.ParseFloat(strpageSize, 64)
190+
if pageSize, err := jmespath.Search(a.PageSizeExpr, j); err == nil {
191+
var PageSize float64
192+
if strpageSize, ok := pageSize.(string); ok {
193+
PageSize, _ = strconv.ParseFloat(strpageSize, 64)
194+
} else {
195+
PageSize = pageSize.(float64)
196+
}
197+
a.PageSize = int(PageSize)
166198
} else {
167-
PageSize = pageSize.(float64)
199+
return fmt.Errorf("jmespath: '%s' failed %s", a.PageSizeExpr, err)
168200
}
169-
a.PageSize = int(PageSize)
170-
} else {
171-
return fmt.Errorf("jmespath: '%s' failed %s", a.PageSizeExpr, err)
172201
}
173202

174203
if a.collectionPath == "" {
@@ -183,10 +212,13 @@ func (a *Pager) FeedResponse(body string) error {
183212
}
184213

185214
func (a *Pager) MoveNextPage(request *requests.CommonRequest) {
186-
a.currentPageNumber = a.currentPageNumber + 1
187-
// cli.Printf("Move to page %d", a.currentPageNumber)
188-
189-
request.QueryParams[a.PageNumberFlag] = strconv.Itoa(a.currentPageNumber)
215+
if a.nextTokenMode {
216+
request.QueryParams[a.NextTokenFlag] = a.nextToken
217+
} else {
218+
a.currentPageNumber = a.currentPageNumber + 1
219+
// cli.Printf("Move to page %d", a.currentPageNumber)
220+
request.QueryParams[a.PageNumberFlag] = strconv.Itoa(a.currentPageNumber)
221+
}
190222
}
191223

192224
func (a *Pager) mergeCollections(body interface{}) error {

‎openapi/pager_test.go‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,30 @@ func TestPager_CallWith(t *testing.T) {
124124
assert.Contains(t, err.Error(), "please check your accessKey with secret, and read the user guide")
125125
}
126126

127+
func TestPager_NextToken(t *testing.T) {
128+
pager := &Pager{
129+
PageNumberFlag: "PageNumber",
130+
PageSizeFlag: "PageSize",
131+
NextTokenFlag: "NextToken",
132+
PageNumberExpr: "PageNumber",
133+
PageSizeExpr: "PageSize",
134+
TotalCountExpr: "TotalCount",
135+
NextTokenExpr: "NextToken",
136+
}
137+
assert.NotNil(t, pager)
138+
assert.Nil(t, pager.FeedResponse(pagerTestJsonNextToken1))
139+
assert.True(t, pager.HasMore())
140+
request := requests.NewCommonRequest()
141+
pager.MoveNextPage(request)
142+
assert.Equal(t, "7a758b45874db76a0147eb118b6f597f2df87e458f7fbf48e0b5e8707e68181f", request.QueryParams["NextToken"])
143+
assert.Nil(t, pager.FeedResponse(pagerTestJsonNextToken2))
144+
assert.False(t, pager.HasMore())
145+
responseStr := pager.GetResponseCollection()
146+
var j map[string]interface{}
147+
assert.Nil(t, json.Unmarshal([]byte(responseStr), &j))
148+
assert.Equal(t, 7, len(j["TagResources"].(map[string]interface{})["TagResource"].([]interface{})))
149+
}
150+
127151
func TestPager_HasMore(t *testing.T) {
128152
pager := Pager{
129153
PageSize: 5,
@@ -135,6 +159,10 @@ func TestPager_HasMore(t *testing.T) {
135159
pager.currentPageNumber = 3
136160
istrue = pager.HasMore()
137161
assert.False(t, istrue)
162+
163+
pager.nextTokenMode = true
164+
pager.nextToken = "aaaaaa"
165+
assert.True(t, pager.HasMore())
138166
}
139167

140168
func TestPager_GetResponseCollection(t *testing.T) {
@@ -176,6 +204,10 @@ func TestPager_FeedResponse(t *testing.T) {
176204
pager.PageSizeExpr = "PageSize"
177205
err = pager.FeedResponse(body)
178206
assert.Nil(t, err)
207+
208+
pager.NextTokenExpr = "NextToken"
209+
err = pager.FeedResponse(body)
210+
assert.Nil(t, err)
179211
}
180212

181213
func TestPager_MoveNextPage(t *testing.T) {
@@ -186,6 +218,12 @@ func TestPager_MoveNextPage(t *testing.T) {
186218
}
187219
pager.MoveNextPage(request)
188220
assert.Equal(t, 1, pager.currentPageNumber)
221+
222+
pager.nextTokenMode = true
223+
pager.nextToken = "aaaaa"
224+
pager.NextTokenFlag = "NextToken"
225+
pager.MoveNextPage(request)
226+
assert.Equal(t, "aaaaa", request.QueryParams["NextToken"])
189227
}
190228

191229
func Test_detectArrayPath(t *testing.T) {
@@ -533,3 +571,62 @@ var pagerTestJson2 = []byte(`{
533571
]
534572
}
535573
}`)
574+
575+
const pagerTestJsonNextToken1 = `{
576+
"NextToken": "7a758b45874db76a0147eb118b6f597f2df87e458f7fbf48e0b5e8707e68181f",
577+
"RequestId": "EAB4BE0E-8CB2-4445-B671-40C8648308D0",
578+
"TagResources": {
579+
"TagResource": [
580+
{
581+
"ResourceId": "d-8vbi818ykkw7pmbvb2q7",
582+
"ResourceType": "disk",
583+
"TagKey": "AppGroup",
584+
"TagValue": "daily-test-ecs"
585+
},
586+
{
587+
"ResourceId": "d-8vbi818ykkw7pmbvb2q7",
588+
"ResourceType": "disk",
589+
"TagKey": "Provider",
590+
"TagValue": "sigma"
591+
},
592+
{
593+
"ResourceId": "d-8vb0bdoq6qxw1s11flho",
594+
"ResourceType": "disk",
595+
"TagKey": "AppGroup",
596+
"TagValue": "daily-test-ecs"
597+
},
598+
{
599+
"ResourceId": "d-8vb0bdoq6qxw1s11flho",
600+
"ResourceType": "disk",
601+
"TagKey": "Provider",
602+
"TagValue": "sigma"
603+
},
604+
{
605+
"ResourceId": "d-8vb4qthb8rk0uswg3apn",
606+
"ResourceType": "disk",
607+
"TagKey": "AppGroup",
608+
"TagValue": "daily-test-ecs"
609+
},
610+
{
611+
"ResourceId": "d-8vb4qthb8rk0uswg3apn",
612+
"ResourceType": "disk",
613+
"TagKey": "Provider",
614+
"TagValue": "sigma"
615+
}
616+
]
617+
}
618+
}`
619+
620+
const pagerTestJsonNextToken2 = `{
621+
"RequestId": "EADUSE0E-8CB2-4445-B671-40C8648308D0",
622+
"TagResources": {
623+
"TagResource": [
624+
{
625+
"ResourceId": "d-8vbi818ykkw7pdsadadf",
626+
"ResourceType": "disk",
627+
"TagKey": "AppGroup",
628+
"TagValue": "daily-test-ecs"
629+
}
630+
]
631+
}
632+
}`

0 commit comments

Comments
 (0)