Merge pull request #2263 from CosmWasm/co/raw-range

Implement `RawRange` query

Author: chipshort

x/wasm/keeper/keeper.go

@@ -939,6 +939,45 @@ func (k Keeper) QueryRaw(ctx context.Context, contractAddress sdk.AccAddress, ke
 	return prefixStore.Get(key)
 }
 
+func (k Keeper) QueryRawRange(ctx context.Context, contractAddress sdk.AccAddress, start, end []byte, limit uint16, reverse bool) (results []wasmvmtypes.RawRangeEntry, nextKey []byte) {
+	defer telemetry.MeasureSince(time.Now(), "wasm", "contract", "query-raw-range")
+
+	prefixStoreKey := types.GetContractStorePrefix(contractAddress)
+	prefixStore := prefix.NewStore(runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx)), prefixStoreKey)
+	var iter storetypes.Iterator
+	if reverse {
+		iter = prefixStore.ReverseIterator(start, end)
+	} else {
+		iter = prefixStore.Iterator(start, end)
+	}
+	defer iter.Close()
+
+	// Make sure to set to empty array because the contract doesn't expect a null JSON value
+	results = []wasmvmtypes.RawRangeEntry{}
+
+	var count uint16 = 0
+	for ; iter.Valid(); iter.Next() {
+		// keep track of count to honor the limit
+		if count == limit {
+			break
+		}
+		count++
+
+		// add key-value pair
+		results = append(results, wasmvmtypes.RawRangeEntry{Key: iter.Key(), Value: iter.Value()})
+	}
+
+	if iter.Valid() {
+		// if there are more results, set the next key
+		key := iter.Key()
+		nextKey = key
+	} else {
+		nextKey = nil
+	}
+
+	return results, nextKey
+}
+
 // internal helper function
 func (k Keeper) contractInstance(ctx context.Context, contractAddress sdk.AccAddress) (types.ContractInfo, types.CodeInfo, wasmvm.KVStore, error) {
 	store := k.storeService.OpenKVStore(ctx)

x/wasm/keeper/keeper_test.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	"bytes"
 	_ "embed"
+	"encoding/binary"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -51,6 +52,9 @@ var hackatomWasm []byte
 //go:embed testdata/replier.wasm
 var replierWasm []byte
 
+//go:embed testdata/queue.wasm
+var queueWasm []byte
+
 var AvailableCapabilities = []string{
 	"iterator", "staking", "stargate", "cosmwasm_1_1", "cosmwasm_1_2", "cosmwasm_1_3",
 	"cosmwasm_1_4", "cosmwasm_2_0", "cosmwasm_2_1", "cosmwasm_2_2", "ibc2",
@@ -2985,3 +2989,170 @@ func TestCheckDiscountEligibility(t *testing.T) {
 		})
 	}
 }
+
+func TestQueryRawRange(t *testing.T) {
+	ctx, keepers := CreateTestInput(t, false, AvailableCapabilities)
+	k := keepers.WasmKeeper
+
+	// Create queue contract and instantiate
+	creator := RandomAccountAddress(t)
+	codeID, _, err := keepers.ContractKeeper.Create(ctx, creator, queueWasm, nil)
+	require.NoError(t, err)
+	initMsgBz := []byte("{}")
+	contractAddress, _, err := keepers.ContractKeeper.Instantiate(ctx, codeID, creator, nil, initMsgBz, "queue", nil)
+
+	type EnqueueMsg struct {
+		Value int32 `json:"value"`
+	}
+	type QueueExecMsg struct {
+		Enqueue *EnqueueMsg `json:"enqueue"`
+		// ...
+	}
+
+	// fill contract storage with 100 items
+	for i := range 100 {
+		enqueueMsg := QueueExecMsg{
+			Enqueue: &EnqueueMsg{Value: int32(i)},
+		}
+		execMsg, err := json.Marshal(enqueueMsg)
+		require.NoError(t, err)
+		_, err = keepers.ContractKeeper.Execute(ctx, contractAddress, creator, execMsg, nil)
+		require.NoError(t, err)
+	}
+
+	type QueueEntry struct {
+		key uint32
+		val int32
+	}
+
+	optUint32 := func(v uint32) *uint32 {
+		return &v
+	}
+
+	specs := map[string]struct {
+		start      *uint32
+		end        *uint32
+		limit      uint16
+		reverse    bool
+		expEntries []QueueEntry
+		expNext    *uint32
+	}{
+		"non-existent range": {
+			start:      optUint32(100),
+			end:        optUint32(200),
+			limit:      10,
+			expEntries: []QueueEntry{},
+			expNext:    nil,
+		},
+		"limited middle range": {
+			start: optUint32(10),
+			end:   optUint32(50),
+			limit: 5,
+			expEntries: []QueueEntry{
+				{key: 10, val: 10},
+				{key: 11, val: 11},
+				{key: 12, val: 12},
+				{key: 13, val: 13},
+				{key: 14, val: 14},
+			},
+			expNext: optUint32(15),
+		},
+		"limited range with no end": {
+			start: optUint32(10),
+			end:   nil,
+			limit: 2,
+			expEntries: []QueueEntry{
+				{key: 10, val: 10},
+				{key: 11, val: 11},
+			},
+			expNext: optUint32(12),
+		},
+		"limited range with no start": {
+			start: nil,
+			end:   optUint32(50),
+			limit: 2,
+			expEntries: []QueueEntry{
+				{key: 0, val: 0},
+				{key: 1, val: 1},
+			},
+			expNext: optUint32(2),
+		},
+		"unbounded range": {
+			start: nil,
+			end:   nil,
+			limit: 1,
+			expEntries: []QueueEntry{
+				{key: 0, val: 0},
+			},
+			expNext: optUint32(1),
+		},
+		"unbounded reversed range": {
+			start:   nil,
+			end:     nil,
+			limit:   1,
+			reverse: true,
+			expEntries: []QueueEntry{
+				{key: 99, val: 99},
+			},
+			expNext: optUint32(98),
+		},
+		"full bounded reversed range": {
+			start:   optUint32(0),
+			end:     optUint32(2),
+			limit:   100,
+			reverse: true,
+			expEntries: []QueueEntry{
+				{key: 1, val: 1},
+				{key: 0, val: 0},
+			},
+			expNext: nil, // no next key because range is fully covered
+		},
+		"start > end, reversed": {
+			start:      optUint32(50),
+			end:        optUint32(10),
+			limit:      5,
+			reverse:    true,
+			expEntries: []QueueEntry{},
+			expNext:    nil,
+		},
+	}
+
+	toBytes := func(v *uint32) []byte {
+		if v == nil {
+			return nil
+		}
+		return binary.BigEndian.AppendUint32(nil, *v)
+	}
+
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			// queue contract uses big endian encoded uint32 as key
+			startBytes := toBytes(spec.start)
+			endBytes := toBytes(spec.end)
+
+			entries, next := k.QueryRawRange(ctx, contractAddress, startBytes, endBytes, spec.limit, spec.reverse)
+			// contract cannot handle nil, so we disallow it
+			require.NotNil(t, entries)
+
+			// converting the entries we get back instead of the entries we put in the spec because
+			// it makes for easier to read test outputs (actual integers instead of byte arrays)
+			convertedEntries := make([]QueueEntry, len(entries))
+			for i, entry := range entries {
+				// values are json-encoded as `{"value":<value>}`
+				// so we need to unmarshal it and extract the value
+				var value EnqueueMsg
+				err := json.Unmarshal(entry.Value, &value)
+				require.NoError(t, err)
+
+				convertedEntries[i] = QueueEntry{
+					key: binary.BigEndian.Uint32(entry.Key),
+					val: value.Value,
+				}
+			}
+
+			expNextBz := toBytes(spec.expNext)
+			assert.Equal(t, spec.expEntries, convertedEntries)
+			assert.Equal(t, expNextBz, next)
+		})
+	}
+}

x/wasm/keeper/query_plugins.go

@@ -108,6 +108,7 @@ type wasmQueryKeeper interface {
 	contractMetaDataSource
 	GetCodeInfo(ctx context.Context, codeID uint64) *types.CodeInfo
 	QueryRaw(ctx context.Context, contractAddress sdk.AccAddress, key []byte) []byte
+	QueryRawRange(ctx context.Context, contractAddress sdk.AccAddress, start, end []byte, limit uint16, reverse bool) (results []wasmvmtypes.RawRangeEntry, nextKey []byte)
 	QuerySmart(ctx context.Context, contractAddr sdk.AccAddress, req []byte) ([]byte, error)
 	IsPinnedCode(ctx context.Context, codeID uint64) bool
 }
@@ -705,6 +706,29 @@ func WasmQuerier(k wasmQueryKeeper) func(ctx sdk.Context, request *wasmvmtypes.W
 				Checksum: info.CodeHash,
 			}
 			return json.Marshal(res)
+		case request.RawRange != nil:
+			contractAddr := request.RawRange.ContractAddr
+			addr, err := sdk.AccAddressFromBech32(contractAddr)
+			if err != nil {
+				return nil, errorsmod.Wrap(sdkerrors.ErrInvalidAddress, contractAddr)
+			}
+
+			var reverse bool
+			switch request.RawRange.Order {
+			case "ascending":
+				reverse = false
+			case "descending":
+				reverse = true
+			default:
+				return nil, errorsmod.Wrapf(sdkerrors.ErrInvalidRequest, "unknown order %s", request.RawRange.Order)
+			}
+			data, nextKey := k.QueryRawRange(ctx, addr, request.RawRange.Start, request.RawRange.End, request.RawRange.Limit, reverse)
+			res := wasmvmtypes.RawRangeResponse{
+				Data:    data,
+				NextKey: nextKey,
+			}
+			return json.Marshal(res)
+
 		}
 		return nil, wasmvmtypes.UnsupportedRequest{Kind: "unknown WasmQuery variant"}
 	}

x/wasm/keeper/query_plugins_test.go

@@ -587,6 +587,97 @@ func TestCodeInfoWasmQuerier(t *testing.T) {
 	}
 }
 
+func TestRawRangeWasmQuerier(t *testing.T) {
+	myValidContractAddr := keeper.RandomBech32AccountAddress(t)
+	validResponse := wasmvmtypes.RawRangeResponse{
+		Data: []wasmvmtypes.RawRangeEntry{
+			{
+				Key:   []byte("key1"),
+				Value: []byte("value"),
+			},
+		},
+		NextKey: nil,
+	}
+	var ctx sdk.Context
+	specs := map[string]struct {
+		req    *wasmvmtypes.WasmQuery
+		mock   mockWasmQueryKeeper
+		expRes wasmvmtypes.RawRangeResponse
+		expErr bool
+	}{
+		"all good": {
+			req: &wasmvmtypes.WasmQuery{
+				RawRange: &wasmvmtypes.RawRangeQuery{
+					ContractAddr: myValidContractAddr,
+					Start:        []byte("key0"),
+					End:          []byte("key2"),
+					Limit:        10,
+					Order:        "ascending",
+				},
+			},
+			mock: mockWasmQueryKeeper{
+				QueryRawRangeFn: func(ctx context.Context, contractAddress sdk.AccAddress, start, end []byte, limit uint16, reverse bool) (results []wasmvmtypes.RawRangeEntry, nextKey []byte) {
+					return validResponse.Data, validResponse.NextKey
+				},
+			},
+			expRes: validResponse,
+		},
+		"all good - descending": {
+			req: &wasmvmtypes.WasmQuery{
+				RawRange: &wasmvmtypes.RawRangeQuery{
+					ContractAddr: myValidContractAddr,
+					Start:        []byte("start"),
+					End:          []byte("end"),
+					Limit:        10,
+					Order:        "descending",
+				},
+			},
+			mock: mockWasmQueryKeeper{
+				QueryRawRangeFn: func(ctx context.Context, contractAddress sdk.AccAddress, start, end []byte, limit uint16, reverse bool) (results []wasmvmtypes.RawRangeEntry, nextKey []byte) {
+					return []wasmvmtypes.RawRangeEntry{}, nil
+				},
+			},
+			expRes: wasmvmtypes.RawRangeResponse{
+				Data:    []wasmvmtypes.RawRangeEntry{},
+				NextKey: nil,
+			},
+		},
+		"invalid addr": {
+			req: &wasmvmtypes.WasmQuery{
+				RawRange: &wasmvmtypes.RawRangeQuery{
+					ContractAddr: "not a valid addr",
+					Order:        "ascending",
+				},
+			},
+			expErr: true,
+		},
+		"invalid order": {
+			req: &wasmvmtypes.WasmQuery{
+				RawRange: &wasmvmtypes.RawRangeQuery{
+					ContractAddr: myValidContractAddr,
+					Order:        "not a valid order",
+				},
+			},
+			expErr: true,
+		},
+	}
+
+	for name, spec := range specs {
+		t.Run(name, func(t *testing.T) {
+			q := keeper.WasmQuerier(spec.mock)
+			gotBz, gotErr := q(ctx, spec.req)
+			if spec.expErr {
+				require.Error(t, gotErr)
+				return
+			}
+			require.NoError(t, gotErr)
+			var gotRes wasmvmtypes.RawRangeResponse
+			require.NoError(t, json.Unmarshal(gotBz, &gotRes))
+			assert.Equal(t, spec.expRes, gotRes)
+		})
+	}
+}
+
 func TestQueryErrors(t *testing.T) {
 	specs := map[string]struct {
 		src    error
@@ -628,6 +719,7 @@ type mockWasmQueryKeeper struct {
 	GetContractInfoFn func(ctx context.Context, contractAddress sdk.AccAddress) *types.ContractInfo
 	QueryRawFn        func(ctx context.Context, contractAddress sdk.AccAddress, key []byte) []byte
 	QuerySmartFn      func(ctx context.Context, contractAddr sdk.AccAddress, req types.RawContractMessage) ([]byte, error)
+	QueryRawRangeFn   func(ctx context.Context, contractAddress sdk.AccAddress, start, end []byte, limit uint16, reverse bool) (results []wasmvmtypes.RawRangeEntry, nextKey []byte)
 	IsPinnedCodeFn    func(ctx context.Context, codeID uint64) bool
 	GetCodeInfoFn     func(ctx context.Context, codeID uint64) *types.CodeInfo
 }
@@ -646,6 +738,13 @@ func (m mockWasmQueryKeeper) QueryRaw(ctx context.Context, contractAddress sdk.A
 	return m.QueryRawFn(ctx, contractAddress, key)
 }
 
+func (m mockWasmQueryKeeper) QueryRawRange(ctx context.Context, contractAddress sdk.AccAddress, start, end []byte, limit uint16, reverse bool) (results []wasmvmtypes.RawRangeEntry, nextKey []byte) {
+	if m.QueryRawRangeFn == nil {
+		panic("not expected to be called")
+	}
+	return m.QueryRawRangeFn(ctx, contractAddress, start, end, limit, reverse)
+}
+
 func (m mockWasmQueryKeeper) QuerySmart(ctx context.Context, contractAddr sdk.AccAddress, req []byte) ([]byte, error) {
 	if m.QuerySmartFn == nil {
 		panic("not expected to be called")