From 9ad848d04b1653061543f1f645c0e583b981c7cd Mon Sep 17 00:00:00 2001 From: Vladimir Mihailenco Date: Mon, 20 Aug 2012 13:42:33 +0300 Subject: [PATCH] Add support for scripting commands. --- README.md | 5 ++- commands.go | 51 +++++++++++++++++++++++++++--- parser.go | 72 +++++++++++++++++++++++++++++------------- redis.go | 2 ++ redis_test.go | 87 +++++++++++++++++++++++++++++++++++++++++++++++++-- req.go | 47 +++++++++++++++++++++++++--- 6 files changed, 231 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index b49bceb8..f655d144 100644 --- a/README.md +++ b/README.md @@ -242,7 +242,10 @@ Some corner cases: client.ZRangeByScoreWithScores("zset", "-inf", "+inf", 0, 2) ZINTERSTORE out 2 zset1 zset2 WEIGHTS 2 3 AGGREGATE SUM - client.ZInterStore("out", 2, redis.ZStore{Weights: []int64{2, 3}}, "zset1", "zset2") + client.ZInterStore("out", redis.ZStore{Weights: []int64{2, 3}}, "zset1", "zset2") + + EVAL "return {KEYS[1],ARGV[1]}" 1 "key" "hello" + client.Eval("return {KEYS[1],ARGV[1]}", []string{"key"}, []string{"hello"}) Contributing ------------ diff --git a/commands.go b/commands.go index ebb4239d..20ad871f 100644 --- a/commands.go +++ b/commands.go @@ -742,11 +742,10 @@ func (c *Client) ZIncrBy(key string, increment float64, member string) *FloatReq func (c *Client) ZInterStore( destination string, - numkeys int64, store ZStore, keys ...string, ) *IntReq { - args := []string{"ZINTERSTORE", destination, strconv.FormatInt(numkeys, 10)} + args := []string{"ZINTERSTORE", destination, strconv.FormatInt(int64(len(keys)), 10)} args = append(args, keys...) if len(store.Weights) > 0 { args = append(args, "WEIGHTS") @@ -904,11 +903,10 @@ func (c *Client) ZScore(key, member string) *FloatReq { func (c *Client) ZUnionStore( destination string, - numkeys int64, store ZStore, keys ...string, ) *IntReq { - args := []string{"ZUNIONSTORE", destination, strconv.FormatInt(numkeys, 10)} + args := []string{"ZUNIONSTORE", destination, strconv.FormatInt(int64(len(keys)), 10)} args = append(args, keys...) if len(store.Weights) > 0 { args = append(args, "WEIGHTS") @@ -1031,3 +1029,48 @@ func (c *Client) Time() *StringSliceReq { c.Process(req) return req } + +//------------------------------------------------------------------------------ + +func (c *Client) Eval(script string, keys []string, args []string) *IfaceReq { + reqArgs := []string{"EVAL", script, strconv.FormatInt(int64(len(keys)), 10)} + reqArgs = append(reqArgs, keys...) + reqArgs = append(reqArgs, args...) + req := NewIfaceReq(reqArgs...) + c.Process(req) + return req +} + +func (c *Client) EvalSha(sha1 string, keys []string, args []string) *IfaceReq { + reqArgs := []string{"EVALSHA", sha1, strconv.FormatInt(int64(len(keys)), 10)} + reqArgs = append(reqArgs, keys...) + reqArgs = append(reqArgs, args...) + req := NewIfaceReq(reqArgs...) + c.Process(req) + return req +} + +func (c *Client) ScriptExists(scripts ...string) *BoolSliceReq { + args := append([]string{"SCRIPT", "EXISTS"}, scripts...) + req := NewBoolSliceReq(args...) + c.Process(req) + return req +} + +func (c *Client) ScriptFlush() *StatusReq { + req := NewStatusReq("SCRIPT", "FLUSH") + c.Process(req) + return req +} + +func (c *Client) ScriptKill() *StatusReq { + req := NewStatusReq("SCRIPT", "KILL") + c.Process(req) + return req +} + +func (c *Client) ScriptLoad(script string) *StringReq { + req := NewStringReq("SCRIPT", "LOAD", script) + c.Process(req) + return req +} diff --git a/parser.go b/parser.go index 1073eb88..7523fb63 100644 --- a/parser.go +++ b/parser.go @@ -57,15 +57,25 @@ func readLine(rd reader) ([]byte, error) { //------------------------------------------------------------------------------ +const ( + ifaceSlice = iota + stringSlice + boolSlice +) + func parseReply(rd reader) (interface{}, error) { - return _parseReply(rd, false) + return _parseReply(rd, ifaceSlice) } -func parseIfaceSliceReply(rd reader) (interface{}, error) { - return _parseReply(rd, true) +func parseStringSliceReply(rd reader) (interface{}, error) { + return _parseReply(rd, stringSlice) } -func _parseReply(rd reader, useIfaceSlice bool) (interface{}, error) { +func parseBoolSliceReply(rd reader) (interface{}, error) { + return _parseReply(rd, boolSlice) +} + +func _parseReply(rd reader, multiBulkType int) (interface{}, error) { line, err := readLine(rd) if err != nil { return 0, err @@ -113,7 +123,42 @@ func _parseReply(rd reader, useIfaceSlice bool) (interface{}, error) { return nil, Nil } - if useIfaceSlice { + switch multiBulkType { + case stringSlice: + numReplies, err := strconv.ParseInt(string(line[1:]), 10, 64) + if err != nil { + return nil, err + } + + vals := make([]string, 0, numReplies) + for i := int64(0); i < numReplies; i++ { + v, err := parseReply(rd) + if err != nil { + return nil, err + } else { + vals = append(vals, v.(string)) + } + } + + return vals, nil + case boolSlice: + numReplies, err := strconv.ParseInt(string(line[1:]), 10, 64) + if err != nil { + return nil, err + } + + vals := make([]bool, 0, numReplies) + for i := int64(0); i < numReplies; i++ { + v, err := parseReply(rd) + if err != nil { + return nil, err + } else { + vals = append(vals, v.(int64) == 1) + } + } + + return vals, nil + default: numReplies, err := strconv.ParseInt(string(line[1:]), 10, 64) if err != nil { return nil, err @@ -131,23 +176,6 @@ func _parseReply(rd reader, useIfaceSlice bool) (interface{}, error) { } } - return vals, nil - } else { - numReplies, err := strconv.ParseInt(string(line[1:]), 10, 64) - if err != nil { - return nil, err - } - - vals := make([]string, 0, numReplies) - for i := int64(0); i < numReplies; i++ { - v, err := parseReply(rd) - if err != nil { - return nil, err - } else { - vals = append(vals, v.(string)) - } - } - return vals, nil } default: diff --git a/redis.go b/redis.go index fca6c7b0..8046b911 100644 --- a/redis.go +++ b/redis.go @@ -137,6 +137,7 @@ func (c *BaseClient) Close() error { type Client struct { *BaseClient + TryEvalShaMinLen int } func NewClient(openConn OpenConnFunc, closeConn CloseConnFunc, initConn InitConnFunc) *Client { @@ -145,6 +146,7 @@ func NewClient(openConn OpenConnFunc, closeConn CloseConnFunc, initConn InitConn ConnPool: NewMultiConnPool(openConn, closeConn, 10), InitConn: initConn, }, + TryEvalShaMinLen: 80, } } diff --git a/redis_test.go b/redis_test.go index e0894223..f1bb2040 100644 --- a/redis_test.go +++ b/redis_test.go @@ -2,6 +2,7 @@ package redis_test import ( "bytes" + "fmt" "io" "net" "runtime" @@ -1818,7 +1819,7 @@ func (t *RedisTest) TestZInterStore(c *C) { c.Assert(zAdd.Err(), IsNil) zInterStore := t.client.ZInterStore( - "out", 2, redis.ZStore{Weights: []int64{2, 3}}, "zset1", "zset2") + "out", redis.ZStore{Weights: []int64{2, 3}}, "zset1", "zset2") c.Assert(zInterStore.Err(), IsNil) c.Assert(zInterStore.Val(), Equals, int64(2)) @@ -2023,7 +2024,7 @@ func (t *RedisTest) TestZUnionStore(c *C) { c.Assert(zAdd.Err(), IsNil) zUnionStore := t.client.ZUnionStore( - "out", 2, redis.ZStore{Weights: []int64{2, 3}}, "zset1", "zset2") + "out", redis.ZStore{Weights: []int64{2, 3}}, "zset1", "zset2") c.Assert(zUnionStore.Err(), IsNil) c.Assert(zUnionStore.Val(), Equals, int64(3)) @@ -2599,6 +2600,88 @@ func (t *RedisTest) TestTime(c *C) { //------------------------------------------------------------------------------ +func (t *RedisTest) TestScriptingEval(c *C) { + eval := t.client.Eval( + "return {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}", + []string{"key1", "key2"}, + []string{"first", "second"}, + ) + c.Assert(eval.Err(), IsNil) + c.Assert(eval.Val(), DeepEquals, []interface{}{"key1", "key2", "first", "second"}) + + eval = t.client.Eval( + "return redis.call('set',KEYS[1],'bar')", + []string{"foo"}, + []string{}, + ) + c.Assert(eval.Err(), IsNil) + c.Assert(eval.Val(), Equals, "OK") + + eval = t.client.Eval("return 10", []string{}, []string{}) + c.Assert(eval.Err(), IsNil) + c.Assert(eval.Val(), Equals, int64(10)) + + eval = t.client.Eval("return {1,2,{3,'Hello World!'}}", []string{}, []string{}) + c.Assert(eval.Err(), IsNil) + // DeepEquals can't compare nested slices. + c.Assert( + fmt.Sprintf("%#v", eval.Val()), + Equals, + `[]interface {}{1, 2, []interface {}{3, "Hello World!"}}`, + ) +} + +func (t *RedisTest) TestScriptingEvalSha(c *C) { + set := t.client.Set("foo", "bar") + c.Assert(set.Err(), IsNil) + c.Assert(set.Val(), Equals, "OK") + + eval := t.client.Eval("return redis.call('get','foo')", []string{}, []string{}) + c.Assert(eval.Err(), IsNil) + c.Assert(eval.Val(), Equals, "bar") + + evalSha := t.client.EvalSha("6b1bf486c81ceb7edf3c093f4c48582e38c0e791", []string{}, []string{}) + c.Assert(evalSha.Err(), IsNil) + c.Assert(evalSha.Val(), Equals, "bar") + + evalSha = t.client.EvalSha("ffffffffffffffffffffffffffffffffffffffff", []string{}, []string{}) + c.Assert(evalSha.Err(), ErrorMatches, "NOSCRIPT No matching script. Please use EVAL.") + c.Assert(evalSha.Val(), Equals, nil) +} + +func (t *RedisTest) TestScriptingScriptExists(c *C) { + scriptLoad := t.client.ScriptLoad("return 1") + c.Assert(scriptLoad.Err(), IsNil) + c.Assert(scriptLoad.Val(), Equals, "e0e1f9fabfc9d4800c877a703b823ac0578ff8db") + + scriptExists := t.client.ScriptExists( + "e0e1f9fabfc9d4800c877a703b823ac0578ff8db", + "ffffffffffffffffffffffffffffffffffffffff", + ) + c.Assert(scriptExists.Err(), IsNil) + c.Assert(scriptExists.Val(), DeepEquals, []bool{true, false}) +} + +func (t *RedisTest) TestScriptingScriptFlush(c *C) { + scriptFlush := t.client.ScriptFlush() + c.Assert(scriptFlush.Err(), IsNil) + c.Assert(scriptFlush.Val(), Equals, "OK") +} + +func (t *RedisTest) TestScriptingScriptKill(c *C) { + scriptKill := t.client.ScriptKill() + c.Assert(scriptKill.Err(), ErrorMatches, "ERR No scripts in execution right now.") + c.Assert(scriptKill.Val(), Equals, "") +} + +func (t *RedisTest) TestScriptingScriptLoad(c *C) { + scriptLoad := t.client.ScriptLoad("return redis.call('get','foo')") + c.Assert(scriptLoad.Err(), IsNil) + c.Assert(scriptLoad.Val(), Equals, "6b1bf486c81ceb7edf3c093f4c48582e38c0e791") +} + +//------------------------------------------------------------------------------ + func (t *RedisTest) BenchmarkRedisPing(c *C) { c.StopTimer() diff --git a/req.go b/req.go index 67b1dea9..014b8db1 100644 --- a/req.go +++ b/req.go @@ -66,6 +66,22 @@ func (r *BaseReq) ParseReply(rd reader) (interface{}, error) { //------------------------------------------------------------------------------ +type IfaceReq struct { + *BaseReq +} + +func NewIfaceReq(args ...string) *IfaceReq { + return &IfaceReq{ + BaseReq: NewBaseReq(args...), + } +} + +func (r *IfaceReq) Val() interface{} { + return r.val +} + +//------------------------------------------------------------------------------ + type StatusReq struct { *BaseReq } @@ -187,10 +203,6 @@ func NewIfaceSliceReq(args ...string) *IfaceSliceReq { } } -func (r *IfaceSliceReq) ParseReply(rd reader) (interface{}, error) { - return parseIfaceSliceReply(rd) -} - func (r *IfaceSliceReq) Val() []interface{} { if r.val == nil { return nil @@ -210,9 +222,36 @@ func NewStringSliceReq(args ...string) *StringSliceReq { } } +func (r *StringSliceReq) ParseReply(rd reader) (interface{}, error) { + return parseStringSliceReply(rd) +} + func (r *StringSliceReq) Val() []string { if r.val == nil { return nil } return r.val.([]string) } + +//------------------------------------------------------------------------------ + +type BoolSliceReq struct { + *BaseReq +} + +func NewBoolSliceReq(args ...string) *BoolSliceReq { + return &BoolSliceReq{ + BaseReq: NewBaseReq(args...), + } +} + +func (r *BoolSliceReq) ParseReply(rd reader) (interface{}, error) { + return parseBoolSliceReply(rd) +} + +func (r *BoolSliceReq) Val() []bool { + if r.val == nil { + return nil + } + return r.val.([]bool) +}