Skip to content

Commit 621c02c

Browse files
SoulPancakechayim
SoulPancake
and
chayim
authored
Adding support for ZMPOP command (#2408)
* feat: add ZMPOP command * fix: reply reading string * fix: evaluating a test tweak * fix: test fix * fix: reverting to debug * fix: test fix * fix: remove old implementation * feat: adding ZMPOP and tests * feat: modifying ZMpopCmd * fix: fix test * fix: test removal check * fix: fix testS * Adding more tests * fix: using redis.Nil instead of string * fix: renaming command to ZArrayWithKeyCmd to match the standard * feat: updated ZArrayWithKeyCmd to ZSliceWithKeyCmd * feat: adding help strings --------- Co-authored-by: Anuragkillswitch <70265851+Anuragkillswitch@users.noreply.github.com> Co-authored-by: Chayim <chayim@users.noreply.github.com>
1 parent f95bdb8 commit 621c02c

File tree

3 files changed

+179
-0
lines changed

3 files changed

+179
-0
lines changed

‎command.go‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3752,3 +3752,86 @@ func (cmd *KeyValuesCmd) readReply(rd *proto.Reader) (err error) {
37523752

37533753
return nil
37543754
}
3755+
3756+
//------------------------------------------------------------------------------
3757+
3758+
type ZSliceWithKeyCmd struct {
3759+
baseCmd
3760+
3761+
key string
3762+
val []Z
3763+
}
3764+
3765+
var _ Cmder = (*ZSliceWithKeyCmd)(nil)
3766+
3767+
func NewZSliceWithKeyCmd(ctx context.Context, args ...interface{}) *ZSliceWithKeyCmd {
3768+
return &ZSliceWithKeyCmd{
3769+
baseCmd: baseCmd{
3770+
ctx: ctx,
3771+
args: args,
3772+
},
3773+
}
3774+
}
3775+
3776+
func (cmd *ZSliceWithKeyCmd) SetVal(key string, val []Z) {
3777+
cmd.key = key
3778+
cmd.val = val
3779+
}
3780+
3781+
func (cmd *ZSliceWithKeyCmd) Val() (string, []Z) {
3782+
return cmd.key, cmd.val
3783+
}
3784+
3785+
func (cmd *ZSliceWithKeyCmd) Result() (string, []Z, error) {
3786+
return cmd.key, cmd.val, cmd.err
3787+
}
3788+
3789+
func (cmd *ZSliceWithKeyCmd) String() string {
3790+
return cmdString(cmd, cmd.val)
3791+
}
3792+
3793+
func (cmd *ZSliceWithKeyCmd) readReply(rd *proto.Reader) (err error) {
3794+
if err = rd.ReadFixedArrayLen(2); err != nil {
3795+
return err
3796+
}
3797+
3798+
cmd.key, err = rd.ReadString()
3799+
if err != nil {
3800+
return err
3801+
}
3802+
3803+
n, err := rd.ReadArrayLen()
3804+
if err != nil {
3805+
return err
3806+
}
3807+
3808+
typ, err := rd.PeekReplyType()
3809+
if err != nil {
3810+
return err
3811+
}
3812+
array := typ == proto.RespArray
3813+
3814+
if array {
3815+
cmd.val = make([]Z, n)
3816+
} else {
3817+
cmd.val = make([]Z, n/2)
3818+
}
3819+
3820+
for i := 0; i < len(cmd.val); i++ {
3821+
if array {
3822+
if err = rd.ReadFixedArrayLen(2); err != nil {
3823+
return err
3824+
}
3825+
}
3826+
3827+
if cmd.val[i].Member, err = rd.ReadString(); err != nil {
3828+
return err
3829+
}
3830+
3831+
if cmd.val[i].Score, err = rd.ReadFloat(); err != nil {
3832+
return err
3833+
}
3834+
}
3835+
3836+
return nil
3837+
}

‎commands.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,7 @@ type Cmdable interface {
313313
ZInterWithScores(ctx context.Context, store *ZStore) *ZSliceCmd
314314
ZInterCard(ctx context.Context, limit int64, keys ...string) *IntCmd
315315
ZInterStore(ctx context.Context, destination string, store *ZStore) *IntCmd
316+
ZMPop(ctx context.Context, order string, count int64, keys ...string) *ZSliceWithKeyCmd
316317
ZMScore(ctx context.Context, key string, members ...string) *FloatSliceCmd
317318
ZPopMax(ctx context.Context, key string, count ...int64) *ZSliceCmd
318319
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
24732474
return cmd
24742475
}
24752476

2477+
// 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.
2478+
// direction: "max" (highest score) or "min" (lowest score), count: > 0
2479+
// example: client.ZMPop(ctx, "max", 5, "set1", "set2")
2480+
func (c cmdable) ZMPop(ctx context.Context, order string, count int64, keys ...string) *ZSliceWithKeyCmd {
2481+
args := make([]interface{}, 2+len(keys), 5+len(keys))
2482+
args[0] = "zmpop"
2483+
args[1] = len(keys)
2484+
for i, key := range keys {
2485+
args[2+i] = key
2486+
}
2487+
args = append(args, strings.ToLower(order), "count", count)
2488+
cmd := NewZSliceWithKeyCmd(ctx, args...)
2489+
_ = c(ctx, cmd)
2490+
return cmd
2491+
}
2492+
24762493
func (c cmdable) ZMScore(ctx context.Context, key string, members ...string) *FloatSliceCmd {
24772494
args := make([]interface{}, 2+len(members))
24782495
args[0] = "zmscore"

‎commands_test.go‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3752,6 +3752,85 @@ var _ = Describe("Commands", func() {
37523752
}}))
37533753
})
37543754

3755+
It("should ZMPop", func() {
3756+
3757+
err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
3758+
Expect(err).NotTo(HaveOccurred())
3759+
err = client.ZAdd(ctx, "zset", redis.Z{Score: 2, Member: "two"}).Err()
3760+
Expect(err).NotTo(HaveOccurred())
3761+
err = client.ZAdd(ctx, "zset", redis.Z{Score: 3, Member: "three"}).Err()
3762+
Expect(err).NotTo(HaveOccurred())
3763+
3764+
err = client.ZAdd(ctx, "zset2", redis.Z{Score: 1, Member: "one"}).Err()
3765+
Expect(err).NotTo(HaveOccurred())
3766+
err = client.ZAdd(ctx, "zset2", redis.Z{Score: 2, Member: "two"}).Err()
3767+
Expect(err).NotTo(HaveOccurred())
3768+
err = client.ZAdd(ctx, "zset2", redis.Z{Score: 3, Member: "three"}).Err()
3769+
Expect(err).NotTo(HaveOccurred())
3770+
3771+
key, elems, err := client.ZMPop(ctx, "min", 1, "zset").Result()
3772+
Expect(err).NotTo(HaveOccurred())
3773+
Expect(key).To(Equal("zset"))
3774+
Expect(elems).To(Equal([]redis.Z{{
3775+
Score: 1,
3776+
Member: "one",
3777+
}}))
3778+
3779+
_, _, err = client.ZMPop(ctx, "min", 1, "nosuchkey").Result()
3780+
Expect(err).To(Equal(redis.Nil))
3781+
3782+
err = client.ZAdd(ctx, "myzset", redis.Z{Score: 1, Member: "one"}).Err()
3783+
Expect(err).NotTo(HaveOccurred())
3784+
err = client.ZAdd(ctx, "myzset", redis.Z{Score: 2, Member: "two"}).Err()
3785+
Expect(err).NotTo(HaveOccurred())
3786+
err = client.ZAdd(ctx, "myzset", redis.Z{Score: 3, Member: "three"}).Err()
3787+
Expect(err).NotTo(HaveOccurred())
3788+
3789+
key, elems, err = client.ZMPop(ctx, "min", 1, "myzset").Result()
3790+
Expect(err).NotTo(HaveOccurred())
3791+
Expect(key).To(Equal("myzset"))
3792+
Expect(elems).To(Equal([]redis.Z{{
3793+
Score:1,
3794+
Member:"one",
3795+
}}))
3796+
3797+
key, elems, err = client.ZMPop(ctx, "max", 10, "myzset").Result()
3798+
Expect(err).NotTo(HaveOccurred())
3799+
Expect(key).To(Equal("myzset"))
3800+
Expect(elems).To(Equal([]redis.Z{{
3801+
Score:3,
3802+
Member:"three",
3803+
},{
3804+
Score: 2,
3805+
Member: "two",
3806+
}}))
3807+
3808+
3809+
err = client.ZAdd(ctx, "myzset2", redis.Z{Score: 4, Member: "four"}).Err()
3810+
Expect(err).NotTo(HaveOccurred())
3811+
err = client.ZAdd(ctx, "myzset2", redis.Z{Score: 5, Member: "five"}).Err()
3812+
Expect(err).NotTo(HaveOccurred())
3813+
err = client.ZAdd(ctx, "myzset2", redis.Z{Score: 6, Member: "six"}).Err()
3814+
Expect(err).NotTo(HaveOccurred())
3815+
3816+
key, elems, err = client.ZMPop(ctx, "min", 10, "myzset","myzset2").Result()
3817+
Expect(err).NotTo(HaveOccurred())
3818+
Expect(key).To(Equal("myzset2"))
3819+
Expect(elems).To(Equal([]redis.Z{{
3820+
Score:4,
3821+
Member:"four",
3822+
},{
3823+
Score: 5,
3824+
Member: "five",
3825+
},{
3826+
Score:6,
3827+
Member: "six",
3828+
}}))
3829+
3830+
3831+
3832+
})
3833+
37553834
It("should ZMScore", func() {
37563835
zmScore := client.ZMScore(ctx, "zset", "one", "three")
37573836
Expect(zmScore.Err()).NotTo(HaveOccurred())

0 commit comments

Comments
 (0)