Skip to content

Commit e7dbdda

Browse files
monkey92tmonkey92t
committed
add HRANDFIELD and ZRANDMEMBER commands
Signed-off-by: monkey <[email protected]>
1 parent cb1be72 commit e7dbdda

File tree

2 files changed

+91
-64
lines changed

2 files changed

+91
-64
lines changed

‎commands.go

Lines changed: 58 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,18 @@ import (
99
"github.com/go-redis/redis/v8/internal"
1010
)
1111

12-
// KeepTTL is an option for Set command to keep key's existing TTL.
13-
// For example:
14-
//
15-
// rdb.Set(ctx, key, value, redis.KeepTTL)
16-
const KeepTTL = -1
12+
const (
13+
// KeepTTL is an option for Set command to keep key's existing TTL.
14+
// For example:
15+
//
16+
// rdb.Set(ctx, key, value, redis.KeepTTL)
17+
KeepTTL = -1
18+
19+
// Persist is remove the time to live associated with the key.
20+
// For example:
21+
// rdb.GetEX(ctx, key, redis.Persist)
22+
Persist = -2
23+
)
1724

1825
func usePrecise(dur time.Duration) bool {
1926
return dur < time.Second || dur%time.Second != 0
@@ -117,7 +124,7 @@ type Cmdable interface {
117124
Get(ctx context.Context, key string) *StringCmd
118125
GetRange(ctx context.Context, key string, start, end int64) *StringCmd
119126
GetSet(ctx context.Context, key string, value interface{}) *StringCmd
120-
GetEX(ctx context.Context, key string, ttl *SetTTL) *StringCmd
127+
GetEX(ctx context.Context, key string, expiration time.Duration) *StringCmd
121128
GetDel(ctx context.Context, key string) *StringCmd
122129
Incr(ctx context.Context, key string) *IntCmd
123130
IncrBy(ctx context.Context, key string, value int64) *IntCmd
@@ -162,6 +169,7 @@ type Cmdable interface {
162169
HMSet(ctx context.Context, key string, values ...interface{}) *BoolCmd
163170
HSetNX(ctx context.Context, key, field string, value interface{}) *BoolCmd
164171
HVals(ctx context.Context, key string) *StringSliceCmd
172+
HRandField(ctx context.Context, key string, count int, withValues bool) *StringSliceCmd
165173

166174
BLPop(ctx context.Context, timeout time.Duration, keys ...string) *StringSliceCmd
167175
BRPop(ctx context.Context, timeout time.Duration, keys ...string) *StringSliceCmd
@@ -265,6 +273,7 @@ type Cmdable interface {
265273
ZRevRank(ctx context.Context, key, member string) *IntCmd
266274
ZScore(ctx context.Context, key, member string) *FloatCmd
267275
ZUnionStore(ctx context.Context, dest string, store *ZStore) *IntCmd
276+
ZRandMember(ctx context.Context, key string, count int, withScores bool) *StringSliceCmd
268277

269278
PFAdd(ctx context.Context, key string, els ...interface{}) *IntCmd
270279
PFCount(ctx context.Context, keys ...string) *IntCmd
@@ -363,58 +372,6 @@ type statefulCmdable func(ctx context.Context, cmd Cmder) error
363372

364373
//------------------------------------------------------------------------------
365374

366-
type ttlAttr int
367-
368-
const (
369-
TExpire ttlAttr = 1 << iota
370-
TExpireAT
371-
TKeepTTL
372-
TPersist
373-
)
374-
375-
// TTL related parameters, not all commands support all ttl attributes.
376-
// priority: Expire > ExpireAt > KeepTTL > Persist
377-
type SetTTL struct {
378-
// set the specified expire time.
379-
// Expire > time.Second AND Expire % time.Second == 0: set key EX Expire/time.Second
380-
// Expire < time.Second OR Expire % time.Second != 0: set key PX Expire/time.Millisecond
381-
Expire time.Duration
382-
383-
// set the specified Unix time at which the key will expire.
384-
// Example: set key EXAT ExpireAt.Unix()
385-
// Don't consider milliseconds for now(PXAT)
386-
ExpireAt time.Time
387-
388-
// Retain the time to live associated with the key.
389-
KeepTTL bool
390-
391-
// Remove the time to live associated with the key, Change to never expire
392-
Persist bool
393-
}
394-
395-
func appendTTL(ctx context.Context, args []interface{}, t *SetTTL, attr ttlAttr) []interface{} {
396-
if t == nil {
397-
return args
398-
}
399-
400-
switch {
401-
case attr&TExpire == 1 && t.Expire > 0:
402-
if usePrecise(t.Expire) {
403-
args = append(args, "px", formatMs(ctx, t.Expire))
404-
} else {
405-
args = append(args, "ex", formatSec(ctx, t.Expire))
406-
}
407-
case attr&TExpireAT == 1 && !t.ExpireAt.IsZero():
408-
args = append(args, "exat", t.ExpireAt.Unix())
409-
case attr&TKeepTTL == 1 && t.KeepTTL:
410-
args = append(args, "keepttl")
411-
case attr&TPersist == 1 && t.Persist:
412-
args = append(args, "persist")
413-
}
414-
415-
return args
416-
}
417-
418375
func (c statefulCmdable) Auth(ctx context.Context, password string) *StatusCmd {
419376
cmd := NewStatusCmd(ctx, "auth", password)
420377
_ = c(ctx, cmd)
@@ -764,17 +721,29 @@ func (c cmdable) GetSet(ctx context.Context, key string, value interface{}) *Str
764721
return cmd
765722
}
766723

767-
// redis-server version >= 6.2.0
768-
func (c cmdable) GetEX(ctx context.Context, key string, ttl *SetTTL) *StringCmd {
769-
args := make([]interface{}, 2, 4)
724+
// redis-server version >= 6.2.0.
725+
//
726+
// A value of zero means that the expiration time will not be changed.
727+
// Persist(-2) Remove the time to live associated with the key.
728+
func (c cmdable) GetEX(ctx context.Context, key string, expiration time.Duration) *StringCmd {
729+
args := make([]interface{}, 0, 4)
770730
args = append(args, "getex", key)
771-
args = appendTTL(ctx, args, ttl, TExpire|TExpireAT|TPersist)
731+
if expiration > 0 {
732+
if usePrecise(expiration) {
733+
args = append(args, "px", formatMs(ctx, expiration))
734+
} else {
735+
args = append(args, "ex", formatSec(ctx, expiration))
736+
}
737+
} else if expiration == Persist {
738+
args = append(args, "persist")
739+
}
740+
772741
cmd := NewStringCmd(ctx, args...)
773742
_ = c(ctx, cmd)
774743
return cmd
775744
}
776745

777-
// redis-server version >= 6.2.0
746+
// redis-server version >= 6.2.0.
778747
func (c cmdable) GetDel(ctx context.Context, key string) *StringCmd {
779748
cmd := NewStringCmd(ctx, "getdel", key)
780749
_ = c(ctx, cmd)
@@ -1253,6 +1222,20 @@ func (c cmdable) HVals(ctx context.Context, key string) *StringSliceCmd {
12531222
return cmd
12541223
}
12551224

1225+
func (c cmdable) HRandField(ctx context.Context, key string, count int, withValues bool) *StringSliceCmd {
1226+
args := make([]interface{}, 0, 4)
1227+
1228+
// Although count=0 is meaningless, redis accepts count=0.
1229+
args = append(args, "hrandfield", key, count)
1230+
if withValues {
1231+
args = append(args, "withvalues")
1232+
}
1233+
1234+
cmd := NewStringSliceCmd(ctx, args...)
1235+
_ = c(ctx, cmd)
1236+
return cmd
1237+
}
1238+
12561239
//------------------------------------------------------------------------------
12571240

12581241
func (c cmdable) BLPop(ctx context.Context, timeout time.Duration, keys ...string) *StringSliceCmd {
@@ -2327,6 +2310,18 @@ func (c cmdable) ZUnionStore(ctx context.Context, dest string, store *ZStore) *I
23272310
return cmd
23282311
}
23292312

2313+
func (c cmdable) ZRandMember(ctx context.Context, key string, count int, withScores bool) *StringSliceCmd {
2314+
args := make([]interface{}, 0, 4)
2315+
args = append(args, "zrandmember", key, count)
2316+
if withScores {
2317+
args = append(args, "withscores")
2318+
}
2319+
2320+
cmd := NewStringSliceCmd(ctx, args...)
2321+
_ = c(ctx, cmd)
2322+
return cmd
2323+
}
2324+
23302325
//------------------------------------------------------------------------------
23312326

23322327
func (c cmdable) PFAdd(ctx context.Context, key string, els ...interface{}) *IntCmd {

‎commands_test.go

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1092,7 +1092,7 @@ var _ = Describe("Commands", func() {
10921092
Expect(ttl.Err()).NotTo(HaveOccurred())
10931093
Expect(ttl.Val()).To(BeNumerically("~", 100*time.Second, 3*time.Second))
10941094

1095-
getEX := client.GetEX(ctx, "key", &redis.SetTTL{Expire: 200 * time.Second})
1095+
getEX := client.GetEX(ctx, "key", 200*time.Second)
10961096
Expect(getEX.Err()).NotTo(HaveOccurred())
10971097
Expect(getEX.Val()).To(Equal("value"))
10981098

@@ -1832,6 +1832,22 @@ var _ = Describe("Commands", func() {
18321832
Expect(err).NotTo(HaveOccurred())
18331833
Expect(slice).To(Equal([]string{"hello1", "hello2"}))
18341834
})
1835+
1836+
It("should HRandField", func() {
1837+
err := client.HSet(ctx, "hash", "key1", "hello1").Err()
1838+
Expect(err).NotTo(HaveOccurred())
1839+
err = client.HSet(ctx, "hash", "key2", "hello2").Err()
1840+
Expect(err).NotTo(HaveOccurred())
1841+
1842+
v := client.HRandField(ctx, "hash", 1, false)
1843+
Expect(v.Err()).NotTo(HaveOccurred())
1844+
Expect(v.Result()).To(Or(Equal([]string{"key1"}), Equal([]string{"key2"})))
1845+
1846+
var slice []string
1847+
err = client.HRandField(ctx, "hash", 1, true).ScanSlice(&slice)
1848+
Expect(err).NotTo(HaveOccurred())
1849+
Expect(slice).To(Or(Equal([]string{"key1", "hello1"}), Equal([]string{"key2", "hello2"})))
1850+
})
18351851
})
18361852

18371853
Describe("hyperloglog", func() {
@@ -3870,6 +3886,22 @@ var _ = Describe("Commands", func() {
38703886
Member: "two",
38713887
}}))
38723888
})
3889+
3890+
It("should ZRandMember", func() {
3891+
err := client.ZAdd(ctx, "zset", &redis.Z{Score: 1, Member: "one"}).Err()
3892+
Expect(err).NotTo(HaveOccurred())
3893+
err = client.ZAdd(ctx, "zset", &redis.Z{Score: 2, Member: "two"}).Err()
3894+
Expect(err).NotTo(HaveOccurred())
3895+
3896+
v := client.ZRandMember(ctx, "zset", 1, false)
3897+
Expect(v.Err()).NotTo(HaveOccurred())
3898+
Expect(v.Val()).To(Or(Equal([]string{"one"}), Equal([]string{"two"})))
3899+
3900+
var slice []string
3901+
err = client.ZRandMember(ctx, "zset", 1, true).ScanSlice(&slice)
3902+
Expect(err).NotTo(HaveOccurred())
3903+
Expect(slice).To(Or(Equal([]string{"one", "1"}), Equal([]string{"two", "2"})))
3904+
})
38733905
})
38743906

38753907
Describe("streams", func() {

0 commit comments

Comments
 (0)