From 249c7fe0adcbe92156c4daab2bc01b8b6e2716ca Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Mon, 13 Jan 2020 17:46:12 +0100
Subject: [PATCH] Add 'raw, all' queries

---
 x/wasm/client/cli/query.go        | 50 +++++++++++++++++++++++++++++--
 x/wasm/internal/keeper/querier.go | 20 +++++++++++--
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/x/wasm/client/cli/query.go b/x/wasm/client/cli/query.go
index cf6cf7332f..840ff985d0 100644
--- a/x/wasm/client/cli/query.go
+++ b/x/wasm/client/cli/query.go
@@ -2,6 +2,7 @@ package cli
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"strconv"
@@ -145,10 +146,26 @@ func GetCmdGetContractInfo(cdc *codec.Codec) *cobra.Command {
 
 // GetCmdGetContractState dumps full internal state of a given contract
 func GetCmdGetContractState(cdc *codec.Codec) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:                        "contract-state",
+		Short:                      "Querying commands for the wasm module",
+		DisableFlagParsing:         true,
+		SuggestionsMinimumDistance: 2,
+		RunE:                       client.ValidateCmd,
+	}
+	cmd.AddCommand(client.GetCommands(
+		GetCmdGetContractStateAll(cdc),
+		GetCmdGetContractStateRaw(cdc),
+	)...)
+	return cmd
+
+}
+
+func GetCmdGetContractStateAll(cdc *codec.Codec) *cobra.Command {
 	return &cobra.Command{
-		Use:   "contract-state [bech32_address]",
-		Short: "Prints out internal state of a contract given its address",
-		Long:  "Prints out internal state of a contract given its address",
+		Use:   "all [bech32_address]",
+		Short: "Prints out all internal state of a contract given its address",
+		Long:  "Prints out all internal state of a contract given its address",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			cliCtx := context.NewCLIContext().WithCodec(cdc)
@@ -168,3 +185,30 @@ func GetCmdGetContractState(cdc *codec.Codec) *cobra.Command {
 		},
 	}
 }
+func GetCmdGetContractStateRaw(cdc *codec.Codec) *cobra.Command {
+	return &cobra.Command{
+		Use:   "raw [bech32_address] [key]",
+		Short: "Prints out internal state for key of a contract given its address",
+		Long:  "Prints out internal state for of a contract given its address",
+		Args:  cobra.ExactArgs(2),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			cliCtx := context.NewCLIContext().WithCodec(cdc)
+
+			addr, err := sdk.AccAddressFromBech32(args[0])
+			if err != nil {
+				return err
+			}
+			key := args[1]
+			if key == "" {
+				return errors.New("key must not be empty")
+			}
+			route := fmt.Sprintf("custom/%s/%s/%s/%s", types.QuerierRoute, keeper.QueryGetContractState, addr.String(), key)
+			res, _, err := cliCtx.Query(route)
+			if err != nil {
+				return err
+			}
+			fmt.Println(string(res))
+			return nil
+		},
+	}
+}
diff --git a/x/wasm/internal/keeper/querier.go b/x/wasm/internal/keeper/querier.go
index 8ec4b89515..8a592d6e31 100644
--- a/x/wasm/internal/keeper/querier.go
+++ b/x/wasm/internal/keeper/querier.go
@@ -1,6 +1,7 @@
 package keeper
 
 import (
+	"bytes"
 	"encoding/json"
 	"strconv"
 
@@ -29,7 +30,18 @@ func NewQuerier(keeper Keeper) sdk.Querier {
 		case QueryListContracts:
 			return queryContractList(ctx, req, keeper)
 		case QueryGetContractState:
-			return queryContractState(ctx, path[1], req, keeper)
+			var accept func([]byte) bool
+			if len(path) > 2 { // passed with key
+				binKey := []byte(path[2])
+				accept = func(b []byte) bool {
+					return bytes.Equal(b, binKey)
+				}
+			} else {
+				accept = func(b []byte) bool {
+					return true
+				}
+			}
+			return selectContractState(ctx, path[1], keeper, accept)
 		case QueryGetCode:
 			return queryCode(ctx, path[1], req, keeper)
 		case QueryListCode:
@@ -67,7 +79,7 @@ func queryContractList(ctx sdk.Context, req abci.RequestQuery, keeper Keeper) ([
 	return bz, nil
 }
 
-func queryContractState(ctx sdk.Context, bech string, req abci.RequestQuery, keeper Keeper) ([]byte, sdk.Error) {
+func selectContractState(ctx sdk.Context, bech string, keeper Keeper, accept func([]byte) bool) ([]byte, sdk.Error) {
 	addr, err := sdk.AccAddressFromBech32(bech)
 	if err != nil {
 		return nil, sdk.ErrUnknownRequest(err.Error())
@@ -76,13 +88,15 @@ func queryContractState(ctx sdk.Context, bech string, req abci.RequestQuery, kee
 
 	var state []types.Model
 	for ; iter.Valid(); iter.Next() {
+		if !accept(iter.Key()) {
+			continue
+		}
 		m := types.Model{
 			Key:   string(iter.Key()),
 			Value: string(iter.Value()),
 		}
 		state = append(state, m)
 	}
-
 	bz, err := json.MarshalIndent(state, "", "  ")
 	if err != nil {
 		return nil, sdk.ErrUnknownRequest(err.Error())