DiceDB#722: Adds support for ZADD and ZRANGE commands (DiceDB#760)

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/dgryski/go-metro v0.0.0-20180109044635-280f6062b5bc // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
+	github.com/google/btree v1.1.3 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/klauspost/cpuid/v2 v2.0.9 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect

go.sum

@@ -37,6 +37,8 @@ github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7z
 github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
 github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
+github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=
+github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=

integration_tests/commands/async/hvals_test.go

@@ -1,7 +1,7 @@
 package async
 
 import (
-	"reflect"
+	testifyAssert "github.com/stretchr/testify/assert"
 	"testing"
 
 	"gotest.tools/v3/assert"
@@ -10,51 +10,57 @@ import (
 func TestHvals(t *testing.T) {
 	conn := getLocalConnection()
 	defer conn.Close()
-	defer FireCommand(conn, "DEL key_hVals key_hVals02")
+	defer FireCommand(conn, "DEL hvalsKey hvalsKey01 hvalsKey02")
 
 	testCases := []TestCase{
 		{
-			commands: []string{"HSET key_hVals field value", "HSET key_hVals field2 value_new", "HVALS key_hVals"},
+			name:     "HVALS with multiple fields",
+			commands: []string{"HSET hvalsKey field value", "HSET hvalsKey field2 value_new", "HVALS hvalsKey"},
 			expected: []interface{}{ONE, ONE, []string{"value", "value_new"}},
 		},
 		{
-			commands: []string{"HVALS key_hVals01"},
-			expected: []interface{}{[]interface{}{}},
+			name:     "HVALS with non-existing key",
+			commands: []string{"HVALS hvalsKey01"},
+			expected: []interface{}{[]string{}},
 		},
 		{
-			commands: []string{"SET key_hVals02 field", "HVALS key_hVals02"},
+			name:     "HVALS on wrong key type",
+			commands: []string{"SET hvalsKey02 field", "HVALS hvalsKey02"},
 			expected: []interface{}{"OK", "WRONGTYPE Operation against a key holding the wrong kind of value"},
 		},
 		{
-			commands: []string{"HVALS key_hVals03 x", "HVALS"},
-			expected: []interface{}{"ERR wrong number of arguments for 'hvals' command",
-				"ERR wrong number of arguments for 'hvals' command"},
+			name:     "HVALS with wrong number of arguments",
+			commands: []string{"HVALS hvalsKey03 x", "HVALS"},
+			expected: []interface{}{"ERR wrong number of arguments for 'hvals' command", "ERR wrong number of arguments for 'hvals' command"},
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			for i, cmd := range tc.commands {
 				result := FireCommand(conn, cmd)
-				expectedResults, ok := tc.expected[i].([]string)
-				results, ok2 := result.([]interface{})
 
-				if ok && ok2 && len(results) == len(expectedResults) {
-					expectedResultsMap := make(map[string]string)
-					resultsMap := make(map[string]string)
-
-					for i := 0; i < len(results); i += 2 {
-						expectedResultsMap[expectedResults[i]] = expectedResults[i+1]
-						resultsMap[results[i].(string)] = results[i+1].(string)
-					}
-					if !reflect.DeepEqual(resultsMap, expectedResultsMap) {
-						t.Fatalf("Assertion failed: expected true, got false")
-					}
+				// Type check for expected value and result
+				expectedList, isExpectedList := tc.expected[i].([]string)
+				resultList, isResultList := result.([]interface{})
 
+				// If both are lists, compare them unordered
+				if isExpectedList && isResultList && len(resultList) == len(expectedList) {
+					testifyAssert.ElementsMatch(t, expectedList, convertToStringSlice(resultList))
 				} else {
+					// Otherwise, do a deep comparison
 					assert.DeepEqual(t, tc.expected[i], result)
 				}
 			}
 		})
 	}
 }
+
+// Helper function to convert []interface{} to []string for easier comparison
+func convertToStringSlice(input []interface{}) []string {
+	output := make([]string, len(input))
+	for i, v := range input {
+		output[i] = v.(string)
+	}
+	return output
+}

integration_tests/commands/async/mget_test.go

@@ -10,6 +10,10 @@ import (
 func TestMGET(t *testing.T) {
 	conn := getLocalConnection()
 	defer conn.Close()
+
+	defer FireCommand(conn, "DEL k1")
+	defer FireCommand(conn, "DEL k2")
+
 	testCases := []struct {
 		name     string
 		commands []string
@@ -39,7 +43,6 @@ func TestMGET(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			// deleteTestKeys([]string{"k1", "k2"}, store)
 			FireCommand(conn, "DEL k1")
 			FireCommand(conn, "DEL k2")
 

integration_tests/commands/async/zset_test.go

@@ -0,0 +1,119 @@
+package async
+
+import (
+	"gotest.tools/v3/assert"
+	"testing"
+)
+
+func TestZADD(t *testing.T) {
+	conn := getLocalConnection()
+	defer conn.Close()
+
+	FireCommand(conn, "DEL key")
+	defer FireCommand(conn, "DEL key")
+
+	testCases := []TestCase{
+		{
+			name:     "ZADD with two new members",
+			commands: []string{"ZADD key 1 member1 2 member2"},
+			expected: []interface{}{int64(2)},
+		},
+		{
+			name:     "ZADD with three new members",
+			commands: []string{"ZADD key 3 member3 4 member4 5 member5"},
+			expected: []interface{}{int64(3)},
+		},
+		{
+			name:     "ZADD with existing members",
+			commands: []string{"ZADD key 1 member1 2 member2 3 member3 4 member4 5 member5"},
+			expected: []interface{}{int64(0)},
+		},
+		{
+			name:     "ZADD with mixed new and existing members",
+			commands: []string{"ZADD key 1 member1 2 member2 3 member3 4 member4 5 member5 6 member6"},
+			expected: []interface{}{int64(1)},
+		},
+		{
+			name:     "ZADD without any members",
+			commands: []string{"ZADD key 1"},
+			expected: []interface{}{"ERR wrong number of arguments for 'zadd' command"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			for i, cmd := range tc.commands {
+				result := FireCommand(conn, cmd)
+				assert.DeepEqual(t, tc.expected[i], result)
+			}
+		})
+	}
+}
+
+func TestZRANGE(t *testing.T) {
+	conn := getLocalConnection()
+	defer conn.Close()
+
+	FireCommand(conn, "DEL key")
+	defer FireCommand(conn, "DEL key")
+
+	FireCommand(conn, "ZADD key 1 member1 2 member2 3 member3 4 member4 5 member5 6 member6")
+	defer FireCommand(conn, "DEL key")
+
+	testCases := []TestCase{
+		{
+			name:     "ZRANGE with mixed indices",
+			commands: []string{"ZRANGE key 0 -1"},
+			expected: []interface{}{[]interface{}{"member1", "member2", "member3", "member4", "member5", "member6"}},
+		},
+		{
+			name:     "ZRANGE with positive indices #1",
+			commands: []string{"ZRANGE key 0 2"},
+			expected: []interface{}{[]interface{}{"member1", "member2", "member3"}},
+		},
+		{
+			name:     "ZRANGE with positive indices #2",
+			commands: []string{"ZRANGE key 2 4"},
+			expected: []interface{}{[]interface{}{"member3", "member4", "member5"}},
+		},
+		{
+			name:     "ZRANGE with all positive indices",
+			commands: []string{"ZRANGE key 0 10"},
+			expected: []interface{}{[]interface{}{"member1", "member2", "member3", "member4", "member5", "member6"}},
+		},
+		{
+			name:     "ZRANGE with out of bound indices",
+			commands: []string{"ZRANGE key 10 20"},
+			expected: []interface{}{[]interface{}{}},
+		},
+		{
+			name:     "ZRANGE with positive indices and scores",
+			commands: []string{"ZRANGE key 0 10 WITHSCORES"},
+			expected: []interface{}{[]interface{}{"member1", "1", "member2", "2", "member3", "3", "member4", "4", "member5", "5", "member6", "6"}},
+		},
+		{
+			name:     "ZRANGE with positive indices and scores in reverse order",
+			commands: []string{"ZRANGE key 0 10 REV WITHSCORES"},
+			expected: []interface{}{[]interface{}{"member6", "6", "member5", "5", "member4", "4", "member3", "3", "member2", "2", "member1", "1"}},
+		},
+		{
+			name:     "ZRANGE with negative indices",
+			commands: []string{"ZRANGE key -1 -1"},
+			expected: []interface{}{[]interface{}{"member6"}},
+		},
+		{
+			name:     "ZRANGE with negative indices and scores",
+			commands: []string{"ZRANGE key -8 -5 WITHSCORES"},
+			expected: []interface{}{[]interface{}{"member1", "1", "member2", "2"}},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			for i, cmd := range tc.commands {
+				result := FireCommand(conn, cmd)
+				assert.DeepEqual(t, tc.expected[i], result)
+			}
+		})
+	}
+}

internal/errors/errors.go

@@ -7,26 +7,28 @@ import (
 )
 
 const (
-	ArityErr             = "wrong number of arguments for '%s' command"
-	SyntaxErr            = "syntax error"
-	ExpiryErr            = "invalid expire time in '%s' command"
-	AuthErr              = "AUTH failed"
-	IntOrOutOfRangeErr   = "value is not an integer or out of range"
-	IntOrFloatErr        = "value is not an integer or a float"
-	ValOutOfRangeErr     = "value is out of range"
-	IncrDecrOverflowErr  = "increment or decrement would overflow"
-	ElementPeekErr       = "number of elements to peek should be a positive number less than %d"
-	NoKeyErr             = "no such key"
-	ErrDefault           = "-ERR %s"
-	WrongTypeErr         = "-WRONGTYPE Operation against a key holding the wrong kind of value"
-	WrongTypeHllErr      = "-WRONGTYPE Key is not a valid HyperLogLog string value."
-	InvalidHllErr        = "-INVALIDOBJ Corrupted HLL object detected"
-	WorkerNotFoundErr    = "worker with ID %s not found"
-	JSONPathNotExistErr  = "-ERR Path '%s' does not exist"
-	JSONPathValueTypeErr = "-WRONGTYPE wrong type of path value - expected string but found integer"
-	InvalidExpireTime    = "-ERR invalid expire time"
+	ArityErr               = "wrong number of arguments for '%s' command"
+	SyntaxErr              = "syntax error"
+	ExpiryErr              = "invalid expire time in '%s' command"
+	AuthErr                = "AUTH failed"
+	IntOrOutOfRangeErr     = "value is not an integer or out of range"
+	IntOrFloatErr          = "value is not an integer or a float"
+	ValOutOfRangeErr       = "value is out of range"
+	IncrDecrOverflowErr    = "increment or decrement would overflow"
+	ElementPeekErr         = "number of elements to peek should be a positive number less than %d"
+	NoKeyErr               = "no such key"
+	ErrDefault             = "-ERR %s"
+	WrongTypeErr           = "-WRONGTYPE Operation against a key holding the wrong kind of value"
+	WrongTypeHllErr        = "-WRONGTYPE Key is not a valid HyperLogLog string value."
+	InvalidHllErr          = "-INVALIDOBJ Corrupted HLL object detected"
+	WorkerNotFoundErr      = "worker with ID %s not found"
+	JSONPathNotExistErr    = "-ERR Path '%s' does not exist"
+	JSONPathValueTypeErr   = "-WRONGTYPE wrong type of path value - expected string but found integer"
+	InvalidExpireTime      = "-ERR invalid expire time"
 	HashValueNotIntegerErr = "hash value is not an integer"
-	InternalServerError  = "-ERR: Internal server error, unable to process command"
+	InternalServerError    = "-ERR: Internal server error, unable to process command"
+	InvalidFloatErr        = "-ERR value is not a valid float"
+	InvalidIntErr          = "-ERR value is not a valid integer"
 )
 
 var (

internal/eval/commands.go

@@ -926,6 +926,28 @@ var (
 		Arity:    -2,
 		KeySpecs: KeySpecs{BeginIndex: 1},
 	}
+	zaddCmdMeta = DiceCmdMeta{
+		Name: "ZADD",
+		Info: `ZADD key [NX|XX] [CH] [INCR] score member [score member ...]
+		Adds all the specified members with the specified scores to the sorted set stored at key.
+		Options: NX, XX, CH, INCR
+		Returns the number of elements added to the sorted set, not including elements already existing for which the score was updated.`,
+		Eval:     evalZADD,
+		Arity:    -4,
+		KeySpecs: KeySpecs{BeginIndex: 1},
+	}
+	zrangeCmdMeta = DiceCmdMeta{
+		Name: "ZRANGE",
+		Info: `ZRANGE key start stop [WithScores]
+		Returns the specified range of elements in the sorted set stored at key.
+		The elements are considered to be ordered from the lowest to the highest score.
+		Both start and stop are 0-based indexes, where 0 is the first element, 1 is the next element and so on.
+		These indexes can also be negative numbers indicating offsets from the end of the sorted set, with -1 being the last element of the sorted set, -2 the penultimate element and so on.
+		Returns the specified range of elements in the sorted set.`,
+		Eval:     evalZRANGE,
+		Arity:    -4,
+		KeySpecs: KeySpecs{BeginIndex: 1},
+	}
 )
 
 func init() {
@@ -1030,6 +1052,8 @@ func init() {
 	DiceCmds["HRANDFIELD"] = hrandfieldCmdMeta
 	DiceCmds["HDEL"] = hdelCmdMeta
 	DiceCmds["HVALS"] = hValsCmdMeta
+	DiceCmds["ZADD"] = zaddCmdMeta
+	DiceCmds["ZRANGE"] = zrangeCmdMeta
 }
 
 // Function to convert DiceCmdMeta to []interface{}

internal/eval/constants.go

@@ -15,12 +15,9 @@ const (
 	Exat       string = "EXAT"
 	XX         string = "XX"
 	NX         string = "NX"
-	Xx         string = "xx"
-	Nx         string = "nx"
 	GT         string = "GT"
 	LT         string = "LT"
-	KEEPTTL    string = "KEEPTTL"
-	Keepttl    string = "keepttl"
+	KeepTTL    string = "KEEPTTL"
 	Sync       string = "SYNC"
 	Async      string = "ASYNC"
 	Help       string = "HELP"
@@ -29,8 +26,8 @@ const (
 	GetKeys    string = "GETKEYS"
 	List       string = "LIST"
 	Info       string = "INFO"
-	Null       string = "null"
 	null       string = "null"
-	NULL       string = "null"
-	WITHVALUES string = "WITHVALUES"
+	WithValues string = "WITHVALUES"
+	WithScores string = "WITHSCORES"
+	REV        string = "REV"
 )