diff --git a/cli/main.go b/cli/main.go index 7f46777..9273ee9 100644 --- a/cli/main.go +++ b/cli/main.go @@ -19,6 +19,8 @@ import ( plog "github.com/pingcap/log" ) +const Json_separator = "Xuan123#@!xuan" + var ( pdAddr = flag.String("pd", "localhost:2379", "PD addr") clientLog = flag.String("log-file", "/dev/null", "TiKV client log file") @@ -37,11 +39,13 @@ var RegisteredCmds = []tcli.Cmd{ kvcmds.ScanPrefixCmd{}, kvcmds.HeadCmd{}, kvcmds.PutCmd{}, + kvcmds.PutJsonCmd{}, kvcmds.BackupCmd{}, kvcmds.NewBenchCmd( kvcmds.NewYcsbBench(*pdAddr), ), kvcmds.GetCmd{}, + kvcmds.GetJsonCmd{}, kvcmds.LoadCsvCmd{}, kvcmds.DeleteCmd{}, kvcmds.DeletePrefixCmd{}, @@ -126,31 +130,62 @@ func main() { shell.AutoHelp(false) // register shell commands - for _, cmd := range RegisteredCmds { + for i, cmd := range RegisteredCmds { handler := cmd.Handler() - //completer := cmd.Completer() longhelp := cmd.LongHelp() - shell.AddCmd(&ishell.Cmd{ - Name: cmd.Name(), - Help: cmd.Help(), - LongHelp: cmd.LongHelp(), - Aliases: cmd.Alias(), - Func: func(c *ishell.Context) { - ctx := context.WithValue(context.TODO(), "ishell", c) - if strings.ToLower(*clientLogLevel) == "debug" { - fmt.Fprintln(os.Stderr, color.YellowString("Input:"), c.RawArgs) - for _, arg := range c.Args { - fmt.Fprintln(os.Stderr, color.YellowString("Arg:"), arg) - } - fmt.Fprintf(os.Stderr, "\033[33mOutput:\033[0m\n") - } - if len(c.Args) > 0 && c.Args[0] == "--help" { - c.Println(longhelp) - return - } - handler(ctx) - }, - }) + // putjson is when i == 4. We need a multi-line input + if i == 4 { + shell.AddCmd(&ishell.Cmd{ + Name: cmd.Name(), + Help: cmd.Help(), + LongHelp: cmd.LongHelp(), + Aliases: cmd.Alias(), + Func: func(c *ishell.Context) { + c.ShowPrompt(false) + defer c.ShowPrompt(true) + /* first type in a key, then enter. + screen print "Please enter a json value end with 'EOF'." + now input the multi-line json value. It has to end with "EOF".*/ + c.Println("Please enter a json value end with 'EOF'.") + lines := c.ReadMultiLines("EOF") + /* connect the key and value with Json_separator to one long string. + the value of Json_separator and "EOF" should be strings that will not exist in the customer input. + separate the combined string later on by Json_separator and get the key and multi-line jsonvalue + */ + if len(c.Args) == 1 { + c.Args[0] = c.Args[0] + Json_separator + lines + c.RawArgs[1] = c.RawArgs[1] + " " + lines + } + ctx := context.WithValue(context.TODO(), "ishell", c) + fmt.Println(color.WhiteString("Input:"), strings.Join(c.RawArgs, " ")) + fmt.Print("\n") + handler(ctx) + }, + }) + } else { + shell.AddCmd(&ishell.Cmd{ + Name: cmd.Name(), + Help: cmd.Help(), + LongHelp: cmd.LongHelp(), + Aliases: cmd.Alias(), + Func: func(c *ishell.Context) { + ctx := context.WithValue(context.TODO(), "ishell", c) + if strings.ToLower(*clientLogLevel) == "debug" { + fmt.Fprintln(os.Stderr, color.YellowString("Input:"), c.RawArgs) + for _, arg := range c.Args { + fmt.Fprintln(os.Stderr, color.YellowString("Arg:"), arg) + } + fmt.Fprintf(os.Stderr, "\033[33mOutput:\033[0m\n") + } + if len(c.Args) > 0 && c.Args[0] == "--help" { + c.Println(longhelp) + return + } + handler(ctx) + }, + }) + + } } shell.Run() shell.Close() diff --git a/client/client.go b/client/client.go index 8fd1bd6..ee561fa 100644 --- a/client/client.go +++ b/client/client.go @@ -108,6 +108,7 @@ type Client interface { GetPDClient() pd.Client Put(ctx context.Context, kv KV) error + PutJson(ctx context.Context, kv KV) error BatchPut(ctx context.Context, kv []KV) error Get(ctx context.Context, k Key) (KV, error) diff --git a/client/rawkv_client.go b/client/rawkv_client.go index f1b0842..351f26e 100644 --- a/client/rawkv_client.go +++ b/client/rawkv_client.go @@ -62,6 +62,10 @@ func (c *rawkvClient) Put(ctx context.Context, kv KV) error { return c.rawClient.Put(context.TODO(), kv.K, kv.V) } +func (c *rawkvClient) PutJson(ctx context.Context, kv KV) error { + return c.rawClient.Put(context.TODO(), kv.K, kv.V) +} + func (c *rawkvClient) BatchPut(ctx context.Context, kvs []KV) error { for _, kv := range kvs { err := c.rawClient.Put(context.TODO(), kv.K[:], kv.V[:]) @@ -83,6 +87,17 @@ func (c *rawkvClient) Get(ctx context.Context, k Key) (KV, error) { return KV{k, v}, nil } +func (c *rawkvClient) GetJson(ctx context.Context, k Key) (KV, error) { + v, err := c.rawClient.Get(context.TODO(), k) + if err != nil { + return KV{}, err + } + if v == nil { + return KV{}, errors.New("key not found") + } + return KV{k, v}, nil +} + func (c *rawkvClient) Scan(ctx context.Context, prefix []byte) (KVS, int, error) { scanOpts := utils.PropFromContext(ctx) diff --git a/client/txnkv_client.go b/client/txnkv_client.go index 43aa089..441aa4a 100644 --- a/client/txnkv_client.go +++ b/client/txnkv_client.go @@ -91,6 +91,24 @@ func (c *txnkvClient) Put(ctx context.Context, kv KV) error { return nil } +func (c *txnkvClient) PutJson(ctx context.Context, kv KV) error { + tx, err := c.txnClient.Begin() + if err != nil { + return err + } + + tx.Set(kv.K, kv.V) + + err = tx.Commit(context.TODO()) + if err != nil { + err = tx.Rollback() + if err != nil { + return err + } + } + return nil +} + func (c *txnkvClient) Scan(ctx context.Context, startKey []byte) (KVS, int, error) { scanOpts := utils.PropFromContext(ctx) tx, err := c.txnClient.Begin() @@ -164,6 +182,18 @@ func (c *txnkvClient) Get(ctx context.Context, k Key) (KV, error) { return KV{K: k, V: v}, nil } +func (c *txnkvClient) GetJson(ctx context.Context, k Key) (KV, error) { + tx, err := c.txnClient.Begin() + if err != nil { + return KV{}, err + } + v, err := tx.Get(context.TODO(), k) + if err != nil { + return KV{}, err + } + return KV{K: k, V: v}, nil +} + func (c *txnkvClient) Delete(ctx context.Context, k Key) error { tx, err := c.txnClient.Begin() if err != nil { diff --git a/go.mod b/go.mod index cc2acc4..86d75aa 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( github.com/magiconair/properties v1.8.0 github.com/manifoldco/promptui v0.8.0 github.com/olekukonko/tablewriter v0.0.5 + github.com/oliveagle/jsonpath v0.0.0-20180606110733-2e52cf6e6852 // indirect github.com/pingcap/go-ycsb v0.0.0-20210727125954-0c816a248fc3 github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4 github.com/pkg/errors v0.9.1 diff --git a/go.sum b/go.sum index 70c3147..dfb8b9e 100644 --- a/go.sum +++ b/go.sum @@ -347,6 +347,8 @@ github.com/olekukonko/tablewriter v0.0.0-20170122224234-a0225b3f23b5/go.mod h1:v github.com/olekukonko/tablewriter v0.0.4/go.mod h1:zq6QwlOf5SlnkVbMSr5EoBv3636FWnp+qbPhuoO21uA= github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec= github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= +github.com/oliveagle/jsonpath v0.0.0-20180606110733-2e52cf6e6852 h1:Yl0tPBa8QPjGmesFh1D0rDy+q1Twx6FyU7VWHi8wZbI= +github.com/oliveagle/jsonpath v0.0.0-20180606110733-2e52cf6e6852/go.mod h1:eqOVx5Vwu4gd2mmMZvVZsgIqNSaW3xxRThUJ0k/TPk4= github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= diff --git a/kvcmds/cmd_get_json.go b/kvcmds/cmd_get_json.go new file mode 100755 index 0000000..8bc045e --- /dev/null +++ b/kvcmds/cmd_get_json.go @@ -0,0 +1,66 @@ +package kvcmds + +import ( + "context" + "github.com/c4pt0r/tcli/utils" + "github.com/c4pt0r/tcli/client" + "fmt" + "encoding/json" + "github.com/oliveagle/jsonpath" +) + +type GetJsonCmd struct{} + +func (c GetJsonCmd) Name() string { return "getjson" } +func (c GetJsonCmd) Alias() []string { return []string{"getj", } } +func (c GetJsonCmd) Help() string { + return `getjson [key] [jsonpath]` +} + +func (c GetJsonCmd) LongHelp() string { + return c.Help() +} + +func (c GetJsonCmd) Handler() func(ctx context.Context) { + return func(ctx context.Context) { + utils.OutputWithElapse(func() error { + ic := utils.ExtractIshellContext(ctx) + if len(ic.Args) != 2 { + utils.Print(c.Help()) + return nil + } + + s := ic.RawArgs[1] + // it's a hex string literal + k, err := utils.GetStringLit(s) + if err != nil { + return err + } + + kv, err := client.GetTiKVClient().Get(context.TODO(), client.Key(k)) + if err != nil { + return err + } + //json_rawdata stores the jsonvalue of the key + json_rawdata := string(kv.V) + + //json path starts with "$.", but there is no "$." from the customer input. Add on + edited_jsonpath := "$." + ic.RawArgs[2] + + var json_data interface{} + + //change the type of json_rawdata(string) to the type available for jsonpath, and store it to json_data + json.Unmarshal([]byte(json_rawdata), &json_data) + + //find the jsonpath in the json_data and print + ans, err := jsonpath.JsonPathLookup(json_data, edited_jsonpath) + if err == nil { + fmt.Println(ans) + } else { + fmt.Println(err) + } + + return nil + }) + } +} diff --git a/kvcmds/cmd_put_json.go b/kvcmds/cmd_put_json.go new file mode 100644 index 0000000..d6cc3ce --- /dev/null +++ b/kvcmds/cmd_put_json.go @@ -0,0 +1,69 @@ +package kvcmds + +import ( + "context" + "fmt" + "github.com/c4pt0r/tcli/utils" + "github.com/c4pt0r/tcli/client" + "encoding/json" + "strings" +) + +const Json_separator = "Xuan123#@!xuan" + +type PutJsonCmd struct{} + +func (c PutJsonCmd) Name() string { return "putjson" } +func (c PutJsonCmd) Alias() []string { return []string{"putjson", "putj"} } +func (c PutJsonCmd) Help() string { + return `putjson [key] [value]` +} +func (c PutJsonCmd) LongHelp() string { + return c.Help() +} + +func (c PutJsonCmd) Handler() func(ctx context.Context) { + return func(ctx context.Context) { + utils.OutputWithElapse(func() error { + ic := utils.ExtractIshellContext(ctx) + // combine key and value together with the separator "Json_separator", so the length should be 1 + if len(ic.Args) != 1 { + fmt.Println(c.LongHelp()) + return nil + } + + //use split to separate k and v by Json_separator + //now s[0] should be the key, and s[1] should be the value，with extra 'EOF' at the end + s := strings.Split(ic.Args[0], Json_separator) + + //value now allowed to be empty, which is just "EOF" + if s[1] == "EOF" { + fmt.Println(c.LongHelp()) + return nil + } + + //take away the "EOF" from value + k := s[0] + v_and_EOF := strings.Split(s[1], "EOF") + v := v_and_EOF[0] + + //checkJson if the value is valid jsonvalue + checkJson := json.Valid([]byte(v)) + + if checkJson == true { + fmt.Println("This is a valud json value. Successfully stored") + err := client.GetTiKVClient().Put(context.TODO(), client.KV{K: []byte(k), V: []byte(v)}) + if err != nil { + fmt.Println("we have error") + return err + } + return nil + } else { + fmt.Println("This is an invalud json value. Please try again") + return nil + } + + return nil + }) + } +}