diff --git a/command.go b/command.go index fa948489..6536f33c 100644 --- a/command.go +++ b/command.go @@ -3752,3 +3752,86 @@ func (cmd *KeyValuesCmd) readReply(rd *proto.Reader) (err error) { return nil } + +//------------------------------------------------------------------------------ + +type ZSliceWithKeyCmd struct { + baseCmd + + key string + val []Z +} + +var _ Cmder = (*ZSliceWithKeyCmd)(nil) + +func NewZSliceWithKeyCmd(ctx context.Context, args ...interface{}) *ZSliceWithKeyCmd { + return &ZSliceWithKeyCmd{ + baseCmd: baseCmd{ + ctx: ctx, + args: args, + }, + } +} + +func (cmd *ZSliceWithKeyCmd) SetVal(key string, val []Z) { + cmd.key = key + cmd.val = val +} + +func (cmd *ZSliceWithKeyCmd) Val() (string, []Z) { + return cmd.key, cmd.val +} + +func (cmd *ZSliceWithKeyCmd) Result() (string, []Z, error) { + return cmd.key, cmd.val, cmd.err +} + +func (cmd *ZSliceWithKeyCmd) String() string { + return cmdString(cmd, cmd.val) +} + +func (cmd *ZSliceWithKeyCmd) readReply(rd *proto.Reader) (err error) { + if err = rd.ReadFixedArrayLen(2); err != nil { + return err + } + + cmd.key, err = rd.ReadString() + if err != nil { + return err + } + + n, err := rd.ReadArrayLen() + if err != nil { + return err + } + + typ, err := rd.PeekReplyType() + if err != nil { + return err + } + array := typ == proto.RespArray + + if array { + cmd.val = make([]Z, n) + } else { + cmd.val = make([]Z, n/2) + } + + for i := 0; i < len(cmd.val); i++ { + if array { + if err = rd.ReadFixedArrayLen(2); err != nil { + return err + } + } + + if cmd.val[i].Member, err = rd.ReadString(); err != nil { + return err + } + + if cmd.val[i].Score, err = rd.ReadFloat(); err != nil { + return err + } + } + + return nil +} diff --git a/commands.go b/commands.go index ca6b1074..bba0bb76 100644 --- a/commands.go +++ b/commands.go @@ -313,6 +313,7 @@ type Cmdable interface { ZInterWithScores(ctx context.Context, store *ZStore) *ZSliceCmd ZInterCard(ctx context.Context, limit int64, keys ...string) *IntCmd ZInterStore(ctx context.Context, destination string, store *ZStore) *IntCmd + ZMPop(ctx context.Context, order string, count int64, keys ...string) *ZSliceWithKeyCmd ZMScore(ctx context.Context, key string, members ...string) *FloatSliceCmd ZPopMax(ctx context.Context, key string, count ...int64) *ZSliceCmd ZPopMin(ctx context.Context, key string, count ...int64) *ZSliceCmd @@ -2473,6 +2474,22 @@ func (c cmdable) ZInterCard(ctx context.Context, limit int64, keys ...string) *I return cmd } +// ZMPop Pops one or more elements with the highest or lowest score from the first non-empty sorted set key from the list of provided key names. +// direction: "max" (highest score) or "min" (lowest score), count: > 0 +// example: client.ZMPop(ctx, "max", 5, "set1", "set2") +func (c cmdable) ZMPop(ctx context.Context, order string, count int64, keys ...string) *ZSliceWithKeyCmd { + args := make([]interface{}, 2+len(keys), 5+len(keys)) + args[0] = "zmpop" + args[1] = len(keys) + for i, key := range keys { + args[2+i] = key + } + args = append(args, strings.ToLower(order), "count", count) + cmd := NewZSliceWithKeyCmd(ctx, args...) + _ = c(ctx, cmd) + return cmd +} + func (c cmdable) ZMScore(ctx context.Context, key string, members ...string) *FloatSliceCmd { args := make([]interface{}, 2+len(members)) args[0] = "zmscore" diff --git a/commands_test.go b/commands_test.go index d34a30fb..7b7dfde2 100644 --- a/commands_test.go +++ b/commands_test.go @@ -3752,6 +3752,85 @@ var _ = Describe("Commands", func() { }})) }) + It("should ZMPop", func() { + + err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zset", redis.Z{Score: 2, Member: "two"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zset", redis.Z{Score: 3, Member: "three"}).Err() + Expect(err).NotTo(HaveOccurred()) + + err = client.ZAdd(ctx, "zset2", redis.Z{Score: 1, Member: "one"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zset2", redis.Z{Score: 2, Member: "two"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "zset2", redis.Z{Score: 3, Member: "three"}).Err() + Expect(err).NotTo(HaveOccurred()) + + key, elems, err := client.ZMPop(ctx, "min", 1, "zset").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("zset")) + Expect(elems).To(Equal([]redis.Z{{ + Score: 1, + Member: "one", + }})) + + _, _, err = client.ZMPop(ctx, "min", 1, "nosuchkey").Result() + Expect(err).To(Equal(redis.Nil)) + + err = client.ZAdd(ctx, "myzset", redis.Z{Score: 1, Member: "one"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "myzset", redis.Z{Score: 2, Member: "two"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "myzset", redis.Z{Score: 3, Member: "three"}).Err() + Expect(err).NotTo(HaveOccurred()) + + key, elems, err = client.ZMPop(ctx, "min", 1, "myzset").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("myzset")) + Expect(elems).To(Equal([]redis.Z{{ + Score:1, + Member:"one", + }})) + + key, elems, err = client.ZMPop(ctx, "max", 10, "myzset").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("myzset")) + Expect(elems).To(Equal([]redis.Z{{ + Score:3, + Member:"three", + },{ + Score: 2, + Member: "two", + }})) + + + err = client.ZAdd(ctx, "myzset2", redis.Z{Score: 4, Member: "four"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "myzset2", redis.Z{Score: 5, Member: "five"}).Err() + Expect(err).NotTo(HaveOccurred()) + err = client.ZAdd(ctx, "myzset2", redis.Z{Score: 6, Member: "six"}).Err() + Expect(err).NotTo(HaveOccurred()) + + key, elems, err = client.ZMPop(ctx, "min", 10, "myzset","myzset2").Result() + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("myzset2")) + Expect(elems).To(Equal([]redis.Z{{ + Score:4, + Member:"four", + },{ + Score: 5, + Member: "five", + },{ + Score:6, + Member: "six", + }})) + + + + }) + It("should ZMScore", func() { zmScore := client.ZMScore(ctx, "zset", "one", "three") Expect(zmScore.Err()).NotTo(HaveOccurred())