add zunionstorecommand & zinterstorecommand in server pkg

This commit is contained in:
wenyekui 2014-08-13 11:59:43 +08:00
parent 87fdc9125e
commit cb908a859c
3 changed files with 162 additions and 1 deletions

View File

@ -315,7 +315,6 @@ func TestZUnionStore(t *testing.T) {
t.Fatal("invalid value ", v)
}
pairs, _ := db.ZRange(out, 0, -1)
n, err = db.ZCount(out, 0, 0XFFFE)
if err != nil {

View File

@ -520,6 +520,120 @@ func zpersistCommand(req *requestContext) error {
return nil
}
func zparseZsetoptStore(args [][]byte) (destKey []byte, srcKeys [][]byte, weights []int64, aggregate byte, err error) {
destKey = args[0]
nKeys, err := strconv.Atoi(ledis.String(args[1]))
if err != nil {
err = ErrValue
return
}
args = args[2:]
if len(args) < nKeys {
err = ErrSyntax
return
}
srcKeys = args[:nKeys]
args = args[nKeys:]
var weightsFlag = false
var aggregateFlag = false
for len(args) > 0 {
if strings.ToLower(ledis.String(args[0])) == "weights" {
if weightsFlag {
err = ErrSyntax
return
}
args = args[1:]
if len(args) < nKeys {
err = ErrSyntax
return
}
weights = make([]int64, nKeys)
for i, arg := range args[:nKeys] {
if weights[i], err = ledis.StrInt64(arg, nil); err != nil {
err = ErrValue
return
}
}
args = args[nKeys:]
weightsFlag = true
} else if strings.ToLower(ledis.String(args[0])) == "aggregate" {
if aggregateFlag {
err = ErrSyntax
return
}
if len(args) < 2 {
err = ErrSyntax
return
}
if strings.ToLower(ledis.String(args[1])) == "sum" {
aggregate = ledis.AggregateSum
} else if strings.ToLower(ledis.String(args[1])) == "min" {
aggregate = ledis.AggregateMin
} else if strings.ToLower(ledis.String(args[1])) == "max" {
aggregate = ledis.AggregateMax
} else {
err = ErrSyntax
return
}
args = args[2:]
aggregateFlag = true
} else {
err = ErrSyntax
return
}
}
if !aggregateFlag {
aggregate = ledis.AggregateSum
}
return
}
func zunionstoreCommand(req *requestContext) error {
args := req.args
if len(args) < 2 {
return ErrCmdParams
}
destKey, srcKeys, weights, aggregate, err := zparseZsetoptStore(args)
if err != nil {
return err
}
if n, err := req.db.ZUnionStore(destKey, srcKeys, weights, aggregate); err != nil {
return err
} else {
req.resp.writeInteger(n)
}
return nil
}
func zinterstoreCommand(req *requestContext) error {
args := req.args
if len(args) < 2 {
return ErrCmdParams
}
destKey, srcKeys, weights, aggregate, err := zparseZsetoptStore(args)
if err != nil {
return err
}
if n, err := req.db.ZInterStore(destKey, srcKeys, weights, aggregate); err != nil {
return err
} else {
req.resp.writeInteger(n)
}
return nil
}
func init() {
register("zadd", zaddCommand)
register("zcard", zcardCommand)
@ -536,6 +650,9 @@ func init() {
register("zrevrangebyscore", zrevrangebyscoreCommand)
register("zscore", zscoreCommand)
register("zunionstore", zunionstoreCommand)
register("zinterstore", zinterstoreCommand)
//ledisdb special command
register("zclear", zclearCommand)

View File

@ -599,3 +599,48 @@ func TestZsetErrorParams(t *testing.T) {
}
}
func TestZUnionStore(t *testing.T) {
c := getTestConn()
defer c.Close()
if _, err := c.Do("zadd", "k1", "1", "one"); err != nil {
t.Fatal(err.Error())
}
if _, err := c.Do("zadd", "k1", "2", "two"); err != nil {
t.Fatal(err.Error())
}
if _, err := c.Do("zadd", "k2", "1", "two"); err != nil {
t.Fatal(err.Error())
}
if _, err := c.Do("zadd", "k2", "2", "three"); err != nil {
t.Fatal(err.Error())
}
if n, err := ledis.Int64(c.Do("zunionstore", "out", "2", "k1", "k2", "weights", "1", "2")); err != nil {
t.Fatal(err.Error())
} else {
if n != 3 {
t.Fatal("invalid value ", n)
}
}
if n, err := ledis.Int64(c.Do("zunionstore", "out", "2", "k1", "k2", "weights", "1", "2", "aggregate", "min")); err != nil {
t.Fatal(err.Error())
} else {
if n != 3 {
t.Fatal("invalid value ", n)
}
}
if n, err := ledis.Int64(c.Do("zunionstore", "out", "2", "k1", "k2", "aggregate", "max")); err != nil {
t.Fatal(err.Error())
} else {
if n != 3 {
t.Fatal("invalid value ", n)
}
}
}