Merge pull request #38 from initia-labs/fix/submodules

fix token query with owner/coll
initia-labs · Apr 25, 2024 · 8d4f87e · 8d4f87e
2 parents e48d0bf + 4a0041c
commit 8d4f87e
Showing 1 changed file with 25 additions and 19 deletions.
diff --git a/submodules/move-nft/grpc_query.go b/submodules/move-nft/grpc_query.go
@@ -186,10 +186,7 @@ func (sm MoveNftSubmodule) getTokensByAccount(ctx context.Context, req *nfttypes
 	ownerSdkAddr := getCosmosAddress(ownerAddr)
 	identifiers := []collections.Pair[sdk.AccAddress, string]{}
 
-	_, pageRes, err := query.CollectionFilteredPaginate(ctx, sm.tokenOwnerMap, req.Pagination,
-		func(k collections.Triple[sdk.AccAddress, sdk.AccAddress, string], _ bool) (bool, error) {
-			return true, nil
-		},
+	_, pageRes, err := query.CollectionPaginate(ctx, sm.tokenOwnerMap, req.Pagination,
 		func(k collections.Triple[sdk.AccAddress, sdk.AccAddress, string], v bool) (bool, error) {
 			identifiers = append(identifiers, collections.Join(k.K2(), k.K3()))
 			return v, nil
@@ -228,24 +225,28 @@ func (sm MoveNftSubmodule) getTokensByAccountAndCollection(ctx context.Context,
 		return nil, status.Error(codes.InvalidArgument, err.Error())
 	}
 	ownerSdkAddr := getCosmosAddress(ownerAddr)
-	ownerAddrStr := ownerSdkAddr.String()
 
-	res, pageRes, err := query.CollectionFilteredPaginate(ctx, sm.tokenMap, req.Pagination,
-		func(k collections.Pair[sdk.AccAddress, string], v nfttypes.IndexedToken) (bool, error) {
-			if slices.Equal(k.K1(), colSdkAddr) && (v.OwnerAddr == ownerAddrStr) {
-				return true, nil
-			}
-			return false, nil
-		},
-		func(k collections.Pair[sdk.AccAddress, string], v nfttypes.IndexedToken) (*nfttypes.IndexedToken, error) {
-			v.CollectionName, _ = sm.getCollectionNameFromPairSubmodule(ctx, v.CollectionName)
-			return &v, nil
+	identifiers := []collections.Pair[sdk.AccAddress, string]{}
+	_, pageRes, err := query.CollectionPaginate(ctx, sm.tokenOwnerMap, req.Pagination,
+		func(k collections.Triple[sdk.AccAddress, sdk.AccAddress, string], v bool) (bool, error) {
+			identifiers = append(identifiers, collections.Join(k.K2(), k.K3()))
+			return v, nil
 		},
+		WithCollectionPaginationTriplePrefix2[sdk.AccAddress, sdk.AccAddress, string](ownerSdkAddr, colSdkAddr),
 	)
-
 	if err != nil {
-		return nil, status.Error(codes.Internal, err.Error())
+		return nil, handleCollectionErr(err)
 	}
+	res := []*nfttypes.IndexedToken{}
+	for _, identifier := range identifiers {
+		token, err := sm.tokenMap.Get(ctx, identifier)
+		if err != nil {
+			return nil, handleCollectionErr(err)
+		}
+		token.CollectionName, _ = sm.getCollectionNameFromPairSubmodule(ctx, token.CollectionName)
+		res = append(res, &token)
+	}
+
 	res = slices.DeleteFunc(res, func(item *nfttypes.IndexedToken) bool {
 		return item == nil
 	})
@@ -281,11 +282,16 @@ func (sm MoveNftSubmodule) getTokensByAccountCollectionAndTokenId(ctx context.Co
 	}, nil
 }
 
-// WithCollectionPaginationTriplePrefix applies a prefix to a collection, whose key is a collection.Triple,
-// being paginated that needs prefixing.
 func WithCollectionPaginationTriplePrefix[K1, K2, K3 any](prefix K1) func(o *query.CollectionsPaginateOptions[collections.Triple[K1, K2, K3]]) {
 	return func(o *query.CollectionsPaginateOptions[collections.Triple[K1, K2, K3]]) {
 		prefix := collections.TriplePrefix[K1, K2, K3](prefix)
 		o.Prefix = &prefix
 	}
 }
+
+func WithCollectionPaginationTriplePrefix2[K1, K2, K3 any](prefix K1, prefix2 K2) func(o *query.CollectionsPaginateOptions[collections.Triple[K1, K2, K3]]) {
+	return func(o *query.CollectionsPaginateOptions[collections.Triple[K1, K2, K3]]) {
+		prefix := collections.TripleSuperPrefix[K1, K2, K3](prefix, prefix2)
+		o.Prefix = &prefix
+	}
+}