refactor(pagination): encode keys on refactored queries

x/deployment/keeper/grpc_query.go

@@ -15,6 +15,7 @@
 	"pkg.akt.dev/go/node/deployment/v1"
 	types "pkg.akt.dev/go/node/deployment/v1beta4"
 
+	"pkg.akt.dev/node/util/query"
 	"pkg.akt.dev/node/x/deployment/keeper/keys"
 )
 
@@ -33,50 +34,74 @@
 
 	if req.Pagination == nil {
 		req.Pagination = &sdkquery.PageRequest{}
+	} else if req.Pagination.Offset > 0 && req.Filters.State == "" {
+		return nil, status.Error(codes.InvalidArgument, "invalid request parameters. if offset is set, filter.state must be provided")
 	}
 
 	if req.Pagination.Limit == 0 {
 		req.Pagination.Limit = sdkquery.DefaultLimit
 	}
 
-	if len(req.Pagination.Key) > 0 {
-		return nil, status.Error(codes.InvalidArgument, "key-based pagination is not supported")
-	}
-
 	ctx := sdk.UnwrapSDKContext(c)
 
-	// Determine which states to iterate
-	states := []v1.Deployment_State{v1.DeploymentActive, v1.DeploymentClosed}
-	if req.Filters.State != "" {
+	states := make([]byte, 0, 2)
+	var resumePK *keys.DeploymentPrimaryKey
+
+	// nolint: gocritic
+	if len(req.Pagination.Key) > 0 {
+		var pkBytes []byte
+		var err error
+		states, _, pkBytes, _, err = query.DecodePaginationKey(req.Pagination.Key)
+		if err != nil {
+			return nil, status.Error(codes.InvalidArgument, err.Error())
+		}
+
+		_, pk, err := k.deployments.KeyCodec().Decode(pkBytes)
+		if err != nil {
+			return nil, status.Error(codes.Internal, err.Error())
+		}
+		resumePK = &pk
+	} else if req.Filters.State != "" {
 		stateVal := v1.Deployment_State(v1.Deployment_State_value[req.Filters.State])
+
 		if stateVal == v1.DeploymentStateInvalid {
 			return nil, status.Error(codes.InvalidArgument, "invalid state value")
 		}
-		states = []v1.Deployment_State{stateVal}
+
+		states = append(states, byte(stateVal))
+	} else {
+		states = append(states, byte(v1.DeploymentActive), byte(v1.DeploymentClosed))
 	}
 
-	if req.Pagination.Reverse {
+	if len(req.Pagination.Key) == 0 && req.Pagination.Reverse {
 		for i, j := 0, len(states)-1; i < j; i, j = i+1, j-1 {
 			states[i], states[j] = states[j], states[i]
 		}
 	}
 
 	var deployments types.DeploymentResponses
-	limit := req.Pagination.Limit
+	var nextKey []byte
+	total := uint64(0)
 	offset := req.Pagination.Offset
-	skipped := uint64(0)
-	countTotal := req.Pagination.CountTotal
-	var total uint64
-	var acctErr error
+	var scanErr error
 
-	for _, state := range states {
-		if limit == 0 && !countTotal {
+	for idx := range states {
+		if req.Pagination.Limit == 0 && len(nextKey) > 0 {
 			break
 		}
 
+		state := v1.Deployment_State(states[idx])
+
 		var iter indexes.MultiIterator[int32, keys.DeploymentPrimaryKey]
 		var err error
-		if req.Pagination.Reverse {
+
+		if idx == 0 && resumePK != nil {
+			r := collections.NewPrefixedPairRange[int32, keys.DeploymentPrimaryKey](int32(state)).StartInclusive(*resumePK)
+			if req.Pagination.Reverse {
+				r = collections.NewPrefixedPairRange[int32, keys.DeploymentPrimaryKey](int32(state)).EndInclusive(*resumePK).Descending()
+			}
+			iter, err = k.deployments.Indexes.State.Iterate(ctx, r)
+		} else if req.Pagination.Reverse {
 			iter, err = k.deployments.Indexes.State.Iterate(ctx,
 				collections.NewPrefixedPairRange[int32, keys.DeploymentPrimaryKey](int32(state)).Descending())
 		} else {
@@ -86,33 +111,43 @@
 			return nil, status.Error(codes.Internal, err.Error())
 		}
 
+		count := uint64(0)
+
 		err = indexes.ScanValues(ctx, k.deployments, iter, func(deployment v1.Deployment) bool {
 			if !req.Filters.Accept(deployment, state) {
 				return false
 			}
 
-			if countTotal {
-				total++
-			}
-
-			if limit == 0 {
-				return !countTotal
+			if offset > 0 {
+				offset--
+				return false
 			}
 
-			if skipped < offset {
-				skipped++
-				return false
+			if req.Pagination.Limit == 0 {
+				// Page is full — encode this item's PK as NextKey
+				pk := keys.DeploymentIDToKey(deployment.ID)
+				pkBuf := make([]byte, k.deployments.KeyCodec().Size(pk))
+				if _, encErr := k.deployments.KeyCodec().Encode(pkBuf, pk); encErr != nil {
+					scanErr = encErr
+					return true
+				}
+				var encErr error
+				nextKey, encErr = query.EncodePaginationKey(states[idx:], []byte{states[idx]}, pkBuf, nil)
+				if encErr != nil {
+					scanErr = encErr
+				}
+				return true
 			}
 
-			account, acctE := k.ekeeper.GetAccount(ctx, deployment.ID.ToEscrowAccountID())
-			if acctE != nil {
-				acctErr = fmt.Errorf("%w: fetching escrow account for DeploymentID=%s", acctE, deployment.ID)
+			account, acctErr := k.ekeeper.GetAccount(ctx, deployment.ID.ToEscrowAccountID())
+			if acctErr != nil {
+				scanErr = fmt.Errorf("%w: fetching escrow account for DeploymentID=%s", acctErr, deployment.ID)
 				return true
 			}
 
 			groups, grpErr := k.GetGroups(ctx, deployment.ID)
 			if grpErr != nil {
-				acctErr = fmt.Errorf("%w: fetching groups for DeploymentID=%s", grpErr, deployment.ID)
+				scanErr = fmt.Errorf("%w: fetching groups for DeploymentID=%s", grpErr, deployment.ID)
 				return true
 			}
 
@@ -121,28 +156,28 @@
 				Groups:        groups,
 				EscrowAccount: account,
 			})
-			limit--
+			req.Pagination.Limit--
+			count++
 
 			return false
 		})
 		if err != nil {
 			return nil, status.Error(codes.Internal, err.Error())
 		}
-		if acctErr != nil {
-			return nil, status.Error(codes.Internal, acctErr.Error())
+		if scanErr != nil {
+			return nil, status.Error(codes.Internal, scanErr.Error())
 		}
-	}
 
-	resp := &types.QueryDeploymentsResponse{
-		Deployments: deployments,
-		Pagination:  &sdkquery.PageResponse{},
+		total += count
 	}
 
-	if countTotal {
-		resp.Pagination.Total = total
-	}
-
-	return resp, nil
+	return &types.QueryDeploymentsResponse{
+		Deployments: deployments,
+		Pagination: &sdkquery.PageResponse{
+			Total:   total,
+			NextKey: nextKey,
+		},
+	}, nil
 }
 
 // Deployment returns deployment details based on DeploymentID

x/deployment/keeper/grpc_query_test.go

@@ -270,6 +270,93 @@ func TestGRPCQueryDeployments(t *testing.T) {
 			page1.Deployments[0].Deployment.ID,
 			"offset pagination must return different deployments")
 	})
+
+	// Validate offset without state filter is rejected
+	t.Run("offset without state filter returns error", func(t *testing.T) {
+		_, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Offset: 1, Limit: 1},
+		})
+		require.Error(t, err)
+	})
+
+	// Validate NextKey is set when there are more results
+	t.Run("NextKey set when more results exist", func(t *testing.T) {
+		res, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Limit: 1},
+		})
+		require.NoError(t, err)
+		require.Len(t, res.Deployments, 1)
+		require.NotNil(t, res.Pagination.NextKey, "NextKey must be set when more results exist")
+		assert.Equal(t, uint64(1), res.Pagination.Total, "Total should equal count of returned items")
+	})
+
+	// Validate NextKey is nil when all results fit
+	t.Run("NextKey nil when all results returned", func(t *testing.T) {
+		res, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Limit: 100},
+		})
+		require.NoError(t, err)
+		require.Len(t, res.Deployments, 3)
+		require.Nil(t, res.Pagination.NextKey, "NextKey must be nil when all results returned")
+		assert.Equal(t, uint64(3), res.Pagination.Total)
+	})
+
+	// Validate key-based pagination returns correct next page
+	t.Run("key-based pagination returns next page", func(t *testing.T) {
+		// Get first page
+		page1, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Limit: 1},
+		})
+		require.NoError(t, err)
+		require.Len(t, page1.Deployments, 1)
+		require.NotNil(t, page1.Pagination.NextKey)
+
+		// Get second page using NextKey
+		page2, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Key: page1.Pagination.NextKey, Limit: 1},
+		})
+		require.NoError(t, err)
+		require.Len(t, page2.Deployments, 1)
+		require.NotNil(t, page2.Pagination.NextKey, "second page should have NextKey (3 items total)")
+
+		// Get third page using NextKey
+		page3, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Key: page2.Pagination.NextKey, Limit: 1},
+		})
+		require.NoError(t, err)
+		require.Len(t, page3.Deployments, 1)
+		require.Nil(t, page3.Pagination.NextKey, "last page should not have NextKey")
+
+		// All three pages should have different deployments
+		ids := map[string]bool{
+			page1.Deployments[0].Deployment.ID.String(): true,
+			page2.Deployments[0].Deployment.ID.String(): true,
+			page3.Deployments[0].Deployment.ID.String(): true,
+		}
+		assert.Len(t, ids, 3, "all pages should return distinct deployments")
+	})
+
+	// Validate key-based pagination with state filter
+	t.Run("key-based pagination with state filter", func(t *testing.T) {
+		page1, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Filters:    v1beta4.DeploymentFilters{State: v1.DeploymentActive.String()},
+			Pagination: &sdkquery.PageRequest{Limit: 1},
+		})
+		require.NoError(t, err)
+		require.Len(t, page1.Deployments, 1)
+		require.NotNil(t, page1.Pagination.NextKey, "should have next key for active deployments")
+
+		page2, err := suite.queryClient.Deployments(suite.ctx, &v1beta4.QueryDeploymentsRequest{
+			Pagination: &sdkquery.PageRequest{Key: page1.Pagination.NextKey, Limit: 10},
+		})
+		require.NoError(t, err)
+		require.Len(t, page2.Deployments, 1, "should return remaining active deployment")
+		require.Nil(t, page2.Pagination.NextKey, "no more active deployments")
+
+		require.NotEqual(t, page1.Deployments[0].Deployment.ID,
+			page2.Deployments[0].Deployment.ID,
+			"pages must return different deployments")
+	})
 }
 
 type deploymentFilterModifier struct {