Skip to content

Commit 0e5c33b

Browse files
committed
refactor(pagination): encode keys on refactored queries
Signed-off-by: Artur Troian <troian@users.noreply.github.com>
1 parent 5203a52 commit 0e5c33b

3 files changed

Lines changed: 472 additions & 201 deletions

File tree

x/deployment/keeper/grpc_query.go

Lines changed: 77 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"pkg.akt.dev/go/node/deployment/v1"
1616
types "pkg.akt.dev/go/node/deployment/v1beta4"
1717

18+
"pkg.akt.dev/node/util/query"
1819
"pkg.akt.dev/node/x/deployment/keeper/keys"
1920
)
2021

@@ -33,50 +34,74 @@ func (k Querier) Deployments(c context.Context, req *types.QueryDeploymentsReque
3334

3435
if req.Pagination == nil {
3536
req.Pagination = &sdkquery.PageRequest{}
37+
} else if req.Pagination.Offset > 0 && req.Filters.State == "" {
38+
return nil, status.Error(codes.InvalidArgument, "invalid request parameters. if offset is set, filter.state must be provided")
3639
}
3740

3841
if req.Pagination.Limit == 0 {
3942
req.Pagination.Limit = sdkquery.DefaultLimit
4043
}
4144

42-
if len(req.Pagination.Key) > 0 {
43-
return nil, status.Error(codes.InvalidArgument, "key-based pagination is not supported")
44-
}
45-
4645
ctx := sdk.UnwrapSDKContext(c)
4746

48-
// Determine which states to iterate
49-
states := []v1.Deployment_State{v1.DeploymentActive, v1.DeploymentClosed}
50-
if req.Filters.State != "" {
47+
states := make([]byte, 0, 2)
48+
var resumePK *keys.DeploymentPrimaryKey
49+
50+
// nolint: gocritic
51+
if len(req.Pagination.Key) > 0 {
52+
var pkBytes []byte
53+
var err error
54+
states, _, pkBytes, _, err = query.DecodePaginationKey(req.Pagination.Key)
55+
if err != nil {
56+
return nil, status.Error(codes.InvalidArgument, err.Error())
57+
}
58+
59+
_, pk, err := k.deployments.KeyCodec().Decode(pkBytes)
60+
if err != nil {
61+
return nil, status.Error(codes.Internal, err.Error())
62+
}
63+
resumePK = &pk
64+
} else if req.Filters.State != "" {
5165
stateVal := v1.Deployment_State(v1.Deployment_State_value[req.Filters.State])
66+
5267
if stateVal == v1.DeploymentStateInvalid {
5368
return nil, status.Error(codes.InvalidArgument, "invalid state value")
5469
}
55-
states = []v1.Deployment_State{stateVal}
70+
71+
states = append(states, byte(stateVal))
72+
} else {
73+
states = append(states, byte(v1.DeploymentActive), byte(v1.DeploymentClosed))
5674
}
5775

58-
if req.Pagination.Reverse {
76+
if len(req.Pagination.Key) == 0 && req.Pagination.Reverse {
5977
for i, j := 0, len(states)-1; i < j; i, j = i+1, j-1 {
6078
states[i], states[j] = states[j], states[i]
6179
}
6280
}
6381

6482
var deployments types.DeploymentResponses
65-
limit := req.Pagination.Limit
83+
var nextKey []byte
84+
total := uint64(0)
6685
offset := req.Pagination.Offset
67-
skipped := uint64(0)
68-
countTotal := req.Pagination.CountTotal
69-
var total uint64
70-
var acctErr error
86+
var scanErr error
7187

72-
for _, state := range states {
73-
if limit == 0 && !countTotal {
88+
for idx := range states {
89+
if req.Pagination.Limit == 0 && len(nextKey) > 0 {
7490
break
7591
}
7692

93+
state := v1.Deployment_State(states[idx])
94+
7795
var iter indexes.MultiIterator[int32, keys.DeploymentPrimaryKey]
7896
var err error
79-
if req.Pagination.Reverse {
97+
98+
if idx == 0 && resumePK != nil {
99+
r := collections.NewPrefixedPairRange[int32, keys.DeploymentPrimaryKey](int32(state)).StartInclusive(*resumePK)
100+
if req.Pagination.Reverse {
101+
r = collections.NewPrefixedPairRange[int32, keys.DeploymentPrimaryKey](int32(state)).EndInclusive(*resumePK).Descending()
102+
}
103+
iter, err = k.deployments.Indexes.State.Iterate(ctx, r)
104+
} else if req.Pagination.Reverse {
80105
iter, err = k.deployments.Indexes.State.Iterate(ctx,
81106
collections.NewPrefixedPairRange[int32, keys.DeploymentPrimaryKey](int32(state)).Descending())
82107
} else {
@@ -86,33 +111,43 @@ func (k Querier) Deployments(c context.Context, req *types.QueryDeploymentsReque
86111
return nil, status.Error(codes.Internal, err.Error())
87112
}
88113

114+
count := uint64(0)
115+
89116
err = indexes.ScanValues(ctx, k.deployments, iter, func(deployment v1.Deployment) bool {
90117
if !req.Filters.Accept(deployment, state) {
91118
return false
92119
}
93120

94-
if countTotal {
95-
total++
96-
}
97-
98-
if limit == 0 {
99-
return !countTotal
121+
if offset > 0 {
122+
offset--
123+
return false
100124
}
101125

102-
if skipped < offset {
103-
skipped++
104-
return false
126+
if req.Pagination.Limit == 0 {
127+
// Page is full — encode this item's PK as NextKey
128+
pk := keys.DeploymentIDToKey(deployment.ID)
129+
pkBuf := make([]byte, k.deployments.KeyCodec().Size(pk))
130+
if _, encErr := k.deployments.KeyCodec().Encode(pkBuf, pk); encErr != nil {
131+
scanErr = encErr
132+
return true
133+
}
134+
var encErr error
135+
nextKey, encErr = query.EncodePaginationKey(states[idx:], []byte{states[idx]}, pkBuf, nil)
136+
if encErr != nil {
137+
scanErr = encErr
138+
}
139+
return true
105140
}
106141

107-
account, acctE := k.ekeeper.GetAccount(ctx, deployment.ID.ToEscrowAccountID())
108-
if acctE != nil {
109-
acctErr = fmt.Errorf("%w: fetching escrow account for DeploymentID=%s", acctE, deployment.ID)
142+
account, acctErr := k.ekeeper.GetAccount(ctx, deployment.ID.ToEscrowAccountID())
143+
if acctErr != nil {
144+
scanErr = fmt.Errorf("%w: fetching escrow account for DeploymentID=%s", acctErr, deployment.ID)
110145
return true
111146
}
112147

113148
groups, grpErr := k.GetGroups(ctx, deployment.ID)
114149
if grpErr != nil {
115-
acctErr = fmt.Errorf("%w: fetching groups for DeploymentID=%s", grpErr, deployment.ID)
150+
scanErr = fmt.Errorf("%w: fetching groups for DeploymentID=%s", grpErr, deployment.ID)
116151
return true
117152
}
118153

@@ -121,28 +156,28 @@ func (k Querier) Deployments(c context.Context, req *types.QueryDeploymentsReque
121156
Groups: groups,
122157
EscrowAccount: account,
123158
})
124-
limit--
159+
req.Pagination.Limit--
160+
count++
125161

126162
return false
127163
})
128164
if err != nil {
129165
return nil, status.Error(codes.Internal, err.Error())
130166
}
131-
if acctErr != nil {
132-
return nil, status.Error(codes.Internal, acctErr.Error())
167+
if scanErr != nil {
168+
return nil, status.Error(codes.Internal, scanErr.Error())
133169
}
134-
}
135170

136-
resp := &types.QueryDeploymentsResponse{
137-
Deployments: deployments,
138-
Pagination: &sdkquery.PageResponse{},
171+
total += count
139172
}
140173

141-
if countTotal {
142-
resp.Pagination.Total = total
143-
}
144-
145-
return resp, nil
174+
return &types.QueryDeploymentsResponse{
175+
Deployments: deployments,
176+
Pagination: &sdkquery.PageResponse{
177+
Total: total,
178+
NextKey: nextKey,
179+
},
180+
}, nil
146181
}
147182

148183
// Deployment returns deployment details based on DeploymentID

x/deployment/keeper/grpc_query_test.go

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,93 @@ func TestGRPCQueryDeployments(t *testing.T) {
270270
page1.Deployments[0].Deployment.ID,
271271
"offset pagination must return different deployments")
272272
})
273+
274+
// Validate offset without state filter is rejected
275+
t.Run("offset without state filter returns error", func(t *testing.T) {
276+
_, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
277+
Pagination: &sdkquery.PageRequest{Offset: 1, Limit: 1},
278+
})
279+
require.Error(t, err)
280+
})
281+
282+
// Validate NextKey is set when there are more results
283+
t.Run("NextKey set when more results exist", func(t *testing.T) {
284+
res, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
285+
Pagination: &sdkquery.PageRequest{Limit: 1},
286+
})
287+
require.NoError(t, err)
288+
require.Len(t, res.Deployments, 1)
289+
require.NotNil(t, res.Pagination.NextKey, "NextKey must be set when more results exist")
290+
assert.Equal(t, uint64(1), res.Pagination.Total, "Total should equal count of returned items")
291+
})
292+
293+
// Validate NextKey is nil when all results fit
294+
t.Run("NextKey nil when all results returned", func(t *testing.T) {
295+
res, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
296+
Pagination: &sdkquery.PageRequest{Limit: 100},
297+
})
298+
require.NoError(t, err)
299+
require.Len(t, res.Deployments, 3)
300+
require.Nil(t, res.Pagination.NextKey, "NextKey must be nil when all results returned")
301+
assert.Equal(t, uint64(3), res.Pagination.Total)
302+
})
303+
304+
// Validate key-based pagination returns correct next page
305+
t.Run("key-based pagination returns next page", func(t *testing.T) {
306+
// Get first page
307+
page1, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
308+
Pagination: &sdkquery.PageRequest{Limit: 1},
309+
})
310+
require.NoError(t, err)
311+
require.Len(t, page1.Deployments, 1)
312+
require.NotNil(t, page1.Pagination.NextKey)
313+
314+
// Get second page using NextKey
315+
page2, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
316+
Pagination: &sdkquery.PageRequest{Key: page1.Pagination.NextKey, Limit: 1},
317+
})
318+
require.NoError(t, err)
319+
require.Len(t, page2.Deployments, 1)
320+
require.NotNil(t, page2.Pagination.NextKey, "second page should have NextKey (3 items total)")
321+
322+
// Get third page using NextKey
323+
page3, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
324+
Pagination: &sdkquery.PageRequest{Key: page2.Pagination.NextKey, Limit: 1},
325+
})
326+
require.NoError(t, err)
327+
require.Len(t, page3.Deployments, 1)
328+
require.Nil(t, page3.Pagination.NextKey, "last page should not have NextKey")
329+
330+
// All three pages should have different deployments
331+
ids := map[string]bool{
332+
page1.Deployments[0].Deployment.ID.String(): true,
333+
page2.Deployments[0].Deployment.ID.String(): true,
334+
page3.Deployments[0].Deployment.ID.String(): true,
335+
}
336+
assert.Len(t, ids, 3, "all pages should return distinct deployments")
337+
})
338+
339+
// Validate key-based pagination with state filter
340+
t.Run("key-based pagination with state filter", func(t *testing.T) {
341+
page1, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
342+
Filters: v1beta4.DeploymentFilters{State: v1.DeploymentActive.String()},
343+
Pagination: &sdkquery.PageRequest{Limit: 1},
344+
})
345+
require.NoError(t, err)
346+
require.Len(t, page1.Deployments, 1)
347+
require.NotNil(t, page1.Pagination.NextKey, "should have next key for active deployments")
348+
349+
page2, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
350+
Pagination: &sdkquery.PageRequest{Key: page1.Pagination.NextKey, Limit: 10},
351+
})
352+
require.NoError(t, err)
353+
require.Len(t, page2.Deployments, 1, "should return remaining active deployment")
354+
require.Nil(t, page2.Pagination.NextKey, "no more active deployments")
355+
356+
require.NotEqual(t, page1.Deployments[0].Deployment.ID,
357+
page2.Deployments[0].Deployment.ID,
358+
"pages must return different deployments")
359+
})
273360
}
274361

275362
type deploymentFilterModifier struct {

0 commit comments

Comments
 (0)