reuse getTokensByAccountAndCollection from move-nft

initia-labs · Jun 13, 2024 · 16d729d · 16d729d
1 parent 30be9d8
commit 16d729d
Showing 1 changed file with 25 additions and 13 deletions.
diff --git a/submodules/wasm-nft/grpc_query.go b/submodules/wasm-nft/grpc_query.go
@@ -229,24 +229,29 @@ func (sm WasmNFTSubmodule) 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 {
+			sm.Logger(ctx).Warn("index mismatch found", "account", ownerSdkAddr, "collection", colSdkAddr, "action", "GetTokensByAccountAndCollection")
+			continue
+		}
+		token.CollectionName, _ = sm.getCollectionNameFromPairSubmodule(ctx, token.CollectionName)
+		res = append(res, &token)
+	}
+
 	res = slices.DeleteFunc(res, func(item *nfttypes.IndexedToken) bool {
 		return item == nil
 	})
@@ -290,3 +295,10 @@ func WithCollectionPaginationTriplePrefix[K1, K2, K3 any](prefix K1) func(o *que
 		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
+	}
+}