Skip to content

Commit

Permalink
Support COUNT argument in SCAN and ZSCAN (#346)
Browse files Browse the repository at this point in the history 
* Implementing Count functionality in SCAN
* Implementing Count functionality in ZSCAN
  • Loading branch information
@BarakSilverfort
BarakSilverfort committed Oct 11, 2023
1 parent 4f4940e commit 279adef
Show file tree
Hide file tree
Showing 6 changed files with 171 additions and 27 deletions.
44 changes: 31 additions & 13 deletions cmd_generic.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -552,6 +552,7 @@ func (m *Miniredis) cmdScan(c *server.Peer, cmd string, args []string) {

var opts struct {
cursor int
count int
withMatch bool
match string
withType bool
Expand All @@ -566,17 +567,23 @@ func (m *Miniredis) cmdScan(c *server.Peer, cmd string, args []string) {
// MATCH, COUNT and TYPE options
for len(args) > 0 {
if strings.ToLower(args[0]) == "count" {
// we do nothing with count
if len(args) < 2 {
setDirty(c)
c.WriteError(msgSyntaxError)
return
}
if _, err := strconv.Atoi(args[1]); err != nil {
count, err := strconv.Atoi(args[1])
if err != nil || count < 0 {
setDirty(c)
c.WriteError(msgInvalidInt)
return
}
if count == 0 {
setDirty(c)
c.WriteError(msgSyntaxError)
return
}
opts.count = count
args = args[2:]
continue
}
Expand Down Expand Up @@ -608,15 +615,6 @@ func (m *Miniredis) cmdScan(c *server.Peer, cmd string, args []string) {
withTx(m, c, func(c *server.Peer, ctx *connCtx) {
db := m.db(ctx.selectedDB)
// We return _all_ (matched) keys every time.

if opts.cursor != 0 {
// Invalid cursor.
c.WriteLen(2)
c.WriteBulk("0") // no next cursor
c.WriteLen(0) // no elements
return
}

var keys []string

if opts.withType {
Expand All @@ -627,17 +625,37 @@ func (m *Miniredis) cmdScan(c *server.Peer, cmd string, args []string) {
keys = append(keys, k)
}
}
sort.Strings(keys) // To make things deterministic.
} else {
keys = db.allKeys()
}

sort.Strings(keys) // To make things deterministic.

if opts.withMatch {
keys, _ = matchKeys(keys, opts.match)
}

low := opts.cursor
high := low + opts.count
// validate high is correct
if high > len(keys) || high == 0 {
high = len(keys)
}
if opts.cursor > high {
// invalid cursor
c.WriteLen(2)
c.WriteBulk("0") // no next cursor
c.WriteLen(0) // no elements
return
}
cursorValue := low + opts.count
if cursorValue >= len(keys) {
cursorValue = 0 // no next cursor
}
keys = keys[low:high]

c.WriteLen(2)
c.WriteBulk("0") // no next cursor
c.WriteBulk(fmt.Sprintf("%d", cursorValue))
c.WriteLen(len(keys))
for _, k := range keys {
c.WriteBulk(k)
Expand Down
44 changes: 43 additions & 1 deletion cmd_generic_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -655,7 +655,7 @@ func TestScan(t *testing.T) {
)
})

t.Run("count (ignored)", func(t *testing.T) {
t.Run("count", func(t *testing.T) {
mustDo(t, c,
"SCAN", "0", "COUNT", "200",
proto.Array(
Expand All @@ -665,6 +665,40 @@ func TestScan(t *testing.T) {
),
),
)

s.Set("v1", "value")
s.Set("v2", "value")
s.Set("v3", "value")
s.Set("v4", "value")
s.Set("v5", "value")
s.Set("v6", "value")
s.Set("v7", "value")
s.Set("v8", "value")
s.Set("v9", "value")

mustDo(t, c,
"SCAN", "0", "COUNT", "3",
proto.Array(
proto.String("3"),
proto.Array(
proto.String("key"),
proto.String("v1"),
proto.String("v2"),
),
),
)

mustDo(t, c,
"SCAN", "3", "COUNT", "3",
proto.Array(
proto.String("6"),
proto.Array(
proto.String("v3"),
proto.String("v4"),
proto.String("v5"),
),
),
)
})

t.Run("match", func(t *testing.T) {
Expand Down Expand Up @@ -726,6 +760,14 @@ func TestScan(t *testing.T) {
"SCAN", "1", "COUNT", "noint",
proto.Error("ERR value is not an integer or out of range"),
)
mustDo(t, c,
"SCAN", "0", "COUNT", "0",
proto.Error("ERR syntax error"),
)
mustDo(t, c,
"SCAN", "0", "COUNT", "-1",
proto.Error("ERR value is not an integer or out of range"),
)
mustDo(t, c,
"SCAN", "1", "TYPE",
proto.Error("ERR syntax error"),
Expand Down
41 changes: 30 additions & 11 deletions cmd_sorted_set.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ package miniredis

import (
"errors"
"fmt"
"math"
"sort"
"strconv"
Expand Down Expand Up @@ -1455,6 +1456,7 @@ func (m *Miniredis) cmdZscan(c *server.Peer, cmd string, args []string) {
var opts struct {
key string
cursor int
count int
withMatch bool
match string
}
Expand All @@ -1472,12 +1474,18 @@ func (m *Miniredis) cmdZscan(c *server.Peer, cmd string, args []string) {
c.WriteError(msgSyntaxError)
return
}
if _, err := strconv.Atoi(args[1]); err != nil {
count, err := strconv.Atoi(args[1])
if err != nil {
setDirty(c)
c.WriteError(msgInvalidInt)
return
}
// We do nothing with count.
if count <= 0 {
setDirty(c)
c.WriteError(msgSyntaxError)
return
}
opts.count = count
args = args[2:]
continue
}
Expand All @@ -1499,14 +1507,6 @@ func (m *Miniredis) cmdZscan(c *server.Peer, cmd string, args []string) {

withTx(m, c, func(c *server.Peer, ctx *connCtx) {
db := m.db(ctx.selectedDB)
// Paging is not implementend, all results are returned for cursor 0.
if opts.cursor != 0 {
// Invalid cursor.
c.WriteLen(2)
c.WriteBulk("0") // no next cursor
c.WriteLen(0) // no elements
return
}
if db.exists(opts.key) && db.t(opts.key) != "zset" {
c.WriteError(ErrWrongType.Error())
return
Expand All @@ -1517,8 +1517,27 @@ func (m *Miniredis) cmdZscan(c *server.Peer, cmd string, args []string) {
members, _ = matchKeys(members, opts.match)
}

low := opts.cursor
high := low + opts.count
// validate high is correct
if high > len(members) || high == 0 {
high = len(members)
}
if opts.cursor > high {
// invalid cursor
c.WriteLen(2)
c.WriteBulk("0") // no next cursor
c.WriteLen(0) // no elements
return
}
cursorValue := low + opts.count
if cursorValue >= len(members) {
cursorValue = 0 // no next cursor
}
members = members[low:high]

c.WriteLen(2)
c.WriteBulk("0") // no next cursor
c.WriteBulk(fmt.Sprintf("%d", cursorValue))
// HSCAN gives key, values.
c.WriteLen(len(members) * 2)
for _, k := range members {
Expand Down
59 changes: 58 additions & 1 deletion cmd_sorted_set_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1563,17 +1563,74 @@ func TestZscan(t *testing.T) {
"ZSCAN", "set", "0", "COUNT",
proto.Error(msgSyntaxError),
)
mustDo(t, c,
"ZSCAN", "set", "0", "COUNT", "0",
proto.Error(msgSyntaxError),
)
mustDo(t, c,
"ZSCAN", "set", "0", "COUNT", "noint",
proto.Error(msgInvalidInt),
)

mustDo(t, c,
"ZSCAN", "set", "0", "COUNT", "-3",
proto.Error(msgSyntaxError),
)
s.Set("str", "value")
mustDo(t, c,
"ZSCAN", "str", "0",
proto.Error(msgWrongType),
)
})

s.ZAdd("largeset", 1.0, "v1")
s.ZAdd("largeset", 2.0, "v2")
s.ZAdd("largeset", 3.0, "v3")
s.ZAdd("largeset", 4.0, "v4")
s.ZAdd("largeset", 5.0, "v5")
s.ZAdd("largeset", 6.0, "v6")
s.ZAdd("largeset", 7.0, "v7")
s.ZAdd("largeset", 8.0, "v8")

mustDo(t, c,
"ZSCAN", "largeset", "0", "COUNT", "3",
proto.Array(
proto.String("3"),
proto.Array(
proto.String("v1"),
proto.String("1"),
proto.String("v2"),
proto.String("2"),
proto.String("v3"),
proto.String("3"),
),
),
)
mustDo(t, c,
"ZSCAN", "largeset", "3", "COUNT", "3",
proto.Array(
proto.String("6"),
proto.Array(
proto.String("v4"),
proto.String("4"),
proto.String("v5"),
proto.String("5"),
proto.String("v6"),
proto.String("6"),
),
),
)
mustDo(t, c,
"ZSCAN", "largeset", "6", "COUNT", "3",
proto.Array(
proto.String("0"),
proto.Array(
proto.String("v7"),
proto.String("7"),
proto.String("v8"),
proto.String("8"),
),
),
)
}

func TestZunionstore(t *testing.T) {
Expand Down
4 changes: 4 additions & 0 deletions integration/generic_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,9 @@ func TestScan(t *testing.T) {
c.Do("SCAN", "0", "TYPE", "set", "COUNT", "100")
c.Do("SCAN", "0", "TYPE", "set", "MATCH", "setkey", "COUNT", "100")

// SCAN may return a higher count of items than requested (See https://redis.io/docs/manual/keyspace/), so we must query all items.
c.DoLoosely("SCAN", "0", "COUNT", "100") // cursor differs

// Can't really test multiple keys.
// c.Do("SET", "key2", "value2")
// c.Do("SCAN", "0")
Expand All @@ -173,6 +176,7 @@ func TestScan(t *testing.T) {
c.Error("wrong number", "SCAN")
c.Error("invalid cursor", "SCAN", "noint")
c.Error("not an integer", "SCAN", "0", "COUNT", "noint")
c.Error("syntax error", "SCAN", "0", "COUNT", "0")
c.Error("syntax error", "SCAN", "0", "COUNT")
c.Error("syntax error", "SCAN", "0", "MATCH")
c.Error("syntax error", "SCAN", "0", "garbage")
Expand Down
6 changes: 5 additions & 1 deletion integration/sorted_set_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -636,7 +636,6 @@ func TestSortedSetIncyby(t *testing.T) {
}

func TestZscan(t *testing.T) {
skip(t)
testRaw(t, func(c *client) {
// No set yet
c.Do("ZSCAN", "h", "0")
Expand All @@ -646,6 +645,9 @@ func TestZscan(t *testing.T) {
c.Do("ZSCAN", "h", "0", "COUNT", "12")
c.Do("ZSCAN", "h", "0", "cOuNt", "12")

// ZSCAN may return a higher count of items than requested (See https://redis.io/docs/manual/keyspace/), so we must query all items.
c.Do("ZSCAN", "h", "0", "COUNT", "10") // cursor differs

c.Do("ZADD", "h", "2.0", "anotherkey")
c.Do("ZSCAN", "h", "0", "MATCH", "anoth*")
c.Do("ZSCAN", "h", "0", "MATCH", "anoth*", "COUNT", "100")
Expand All @@ -660,6 +662,8 @@ func TestZscan(t *testing.T) {
c.Error("wrong number", "ZSCAN", "noint")
c.Error("not an integer", "ZSCAN", "h", "0", "COUNT", "noint")
c.Error("syntax error", "ZSCAN", "h", "0", "COUNT")
c.Error("syntax error", "ZSCAN", "h", "0", "COUNT", "0")
c.Error("syntax error", "ZSCAN", "h", "0", "COUNT", "-1")
c.Error("syntax error", "ZSCAN", "h", "0", "MATCH")
c.Error("syntax error", "ZSCAN", "h", "0", "garbage")
c.Error("syntax error", "ZSCAN", "h", "0", "COUNT", "12", "MATCH", "foo", "garbage")
Expand Down

0 comments on commit 279adef

Please sign in to comment.