HRANDFIELD implementation (#351)

Adds HRANDFIELD
alicebob · Dec 12, 2023 · 794a09d · 794a09d
1 parent feeb189
commit 794a09d
Show file tree

Hide file tree

Showing 5 changed files with 196 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -103,6 +103,7 @@ Implemented commands:
    - HLEN
    - HMGET
    - HMSET
+   - HRANDFIELD
    - HSET
    - HSETNX
    - HSTRLEN

diff --git a/cmd_hash.go b/cmd_hash.go
@@ -27,6 +27,7 @@ func commandsHash(m *Miniredis) {
 	m.srv.Register("HSTRLEN", m.cmdHstrlen)
 	m.srv.Register("HVALS", m.cmdHvals)
 	m.srv.Register("HSCAN", m.cmdHscan)
+	m.srv.Register("HRANDFIELD", m.cmdHrandfield)
 }
 
 // HSET
@@ -681,3 +682,83 @@ func (m *Miniredis) cmdHscan(c *server.Peer, cmd string, args []string) {
 		}
 	})
 }
+
+// HRANDFIELD
+func (m *Miniredis) cmdHrandfield(c *server.Peer, cmd string, args []string) {
+	if len(args) > 3 || len(args) < 1 {
+		setDirty(c)
+		c.WriteError(errWrongNumber(cmd))
+		return
+	}
+	if !m.handleAuth(c) {
+		return
+	}
+	if m.checkPubsub(c, cmd) {
+		return
+	}
+
+	opts := struct {
+		key        string
+		count      int
+		withValues bool
+	}{
+		key:   args[0],
+		count: 1,
+	}
+
+	if len(args) > 1 {
+		if ok := optIntErr(c, args[1], &opts.count, msgInvalidInt); !ok {
+			return
+		}
+	}
+
+	if len(args) == 3 {
+		if strings.ToLower(args[2]) == "withvalues" {
+			opts.withValues = true
+		} else {
+			setDirty(c)
+			c.WriteError(msgSyntaxError)
+			return
+		}
+	}
+
+	withTx(m, c, func(peer *server.Peer, ctx *connCtx) {
+		if opts.count == 0 {
+			peer.WriteLen(0)
+			return
+		}
+		db := m.db(ctx.selectedDB)
+		members := db.hashFields(opts.key)
+		// if cnt positive
+		cnt := 0
+		iterCnt := len(members)
+		if opts.count > 0 {
+			if opts.count > len(members) {
+				cnt = len(members)
+			} else {
+				cnt = opts.count
+			}
+		} else {
+			cnt = -opts.count
+			iterCnt *= cnt
+		}
+
+		p := m.randPerm(iterCnt)
+		if opts.withValues {
+			peer.WriteMapLen(cnt)
+			for i := 0; i < cnt; i++ {
+				idx := p[i] % len(members)
+				peer.WriteBulk(members[idx])
+				peer.WriteBulk(db.hashGet(opts.key, members[idx]))
+			}
+			return
+		} else {
+			peer.WriteLen(cnt)
+			for i := 0; i < cnt; i++ {
+				idx := p[i] % len(members)
+				peer.WriteBulk(members[idx])
+			}
+			return
+		}
+	})
+}
diff --git a/cmd_hash_test.go b/cmd_hash_test.go
@@ -1,6 +1,7 @@
 package miniredis
 
 import (
+	"sort"
 	"testing"
 	"time"
 
@@ -633,3 +634,98 @@ func TestHstrlen(t *testing.T) {
 		)
 	})
 }
+
+func TestHashRandField(t *testing.T) {
+	s := RunT(t)
+	c, err := proto.Dial(s.Addr())
+	ok(t, err)
+	defer c.Close()
+
+	s.HSet("wim", "zus", "jet")
+	s.HSet("wim", "teun", "vuur")
+	s.HSet("wim", "gijs", "lam")
+	s.HSet("wim", "kees", "bok")
+
+	{
+		v, err := c.Do("HRANDFIELD", "wim", "1")
+		ok(t, err)
+		assert(t, v == proto.Strings("zus") || v == proto.Strings("teun") || v == proto.Strings("gijs") || v == proto.Strings("kees"), "HRANDFIELD looks sane")
+	}
+
+	{
+		v, err := c.Do("HRANDFIELD", "wim", "1", "WITHVALUES")
+		ok(t, err)
+		st, err := proto.Parse(v)
+		ok(t, err)
+		li := st.([]interface{})
+		keys := make([]string, len(li))
+		for i, v := range li {
+			keys[i] = v.(string)
+		}
+
+		assert(t, len(keys) == 2, "HRANDFIELD looks sane")
+		assert(t, keys[0] == "zus" || keys[0] == "teun" || keys[0] == "gijs" || keys[0] == "kees", "HRANDFIELD looks sane")
+		assert(t, keys[1] == "jet" || keys[1] == "vuur" || keys[1] == "lam" || keys[1] == "bok", "HRANDFIELD looks sane")
+	}
+
+	{
+		v, err := c.Do("HRANDFIELD", "wim", "4")
+		ok(t, err)
+		st, err := proto.Parse(v)
+		ok(t, err)
+		li := st.([]interface{})
+		keys := make([]string, len(li))
+		for i, v := range li {
+			keys[i] = v.(string)
+		}
+		sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+		assert(t, len(keys) == 4, "HRANDFIELD looks sane")
+		assert(t, keys[0] == "gijs", "HRANDFIELD looks sane")
+		assert(t, keys[1] == "kees", "HRANDFIELD looks sane")
+		assert(t, keys[2] == "teun", "HRANDFIELD looks sane")
+		assert(t, keys[3] == "zus", "HRANDFIELD looks sane")
+	}
+
+	{
+		v, err := c.Do("HRANDFIELD", "wim", "5")
+		ok(t, err)
+		st, err := proto.Parse(v)
+		ok(t, err)
+		li := st.([]interface{})
+		keys := make([]string, len(li))
+		for i, v := range li {
+			keys[i] = v.(string)
+		}
+		sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+		assert(t, len(keys) == 4, "HRANDFIELD looks sane")
+		assert(t, keys[0] == "gijs", "HRANDFIELD looks sane")
+		assert(t, keys[1] == "kees", "HRANDFIELD looks sane")
+		assert(t, keys[2] == "teun", "HRANDFIELD looks sane")
+		assert(t, keys[3] == "zus", "HRANDFIELD looks sane")
+	}
+
+	{
+		v, err := c.Do("HRANDFIELD", "wim", "-5")
+		ok(t, err)
+		st, err := proto.Parse(v)
+		ok(t, err)
+		li := st.([]interface{})
+		keys := make([]string, len(li))
+		for i, v := range li {
+			keys[i] = v.(string)
+		}
+
+		keyMap := make(map[string]bool)
+		for _, key := range keys {
+			keyMap[key] = true
+		}
+		assert(t, len(keys) == 5, "HRANDFIELD looks sane")
+		assert(t, len(keyMap) <= 4, "HRANDFIELD looks sane")
+	}
+
+	// Wrong key type
+	mustDo(t, c,
+		"HRANDFIELD", "wim", "zus",
+		proto.Error(msgInvalidInt),
+	)
+}
diff --git a/integration/hash_test.go b/integration/hash_test.go
@@ -229,3 +229,14 @@ func TestHstrlen(t *testing.T) {
 		c.Error("wrong kind", "HSTRLEN", "str", "bar")
 	})
 }
+
+func TestHrandfield(t *testing.T) {
+	skip(t)
+	testRaw(t, func(c *client) {
+		// A random key from a DB with a single key. We can test that.
+		c.Do("HSET", "one", "foo", "bar")
+		c.Do("HRANDFIELD", "one", "1")
+
+		c.Error("ERR syntax error", "HRANDFIELD", "foo", "1", "2")
+	})
+}
diff --git a/miniredis.go b/miniredis.go
@@ -634,6 +634,13 @@ func (m *Miniredis) randIntn(n int) int {
 	return m.rand.Intn(n)
 }
 
+func (m *Miniredis) randPerm(n int) []int {
+	if m.rand == nil {
+		return rand.Perm(n)
+	}
+	return m.rand.Perm(n)
+}
+
 // shuffle shuffles a list of strings. Kinda.
 func (m *Miniredis) shuffle(l []string) {
 	for range l {