DiceDB#496: Feature JSON.NUMINCRBY (DiceDB#585)

Co-authored-by: ayadav16 <[email protected]>

go.mod

@@ -55,7 +55,7 @@ require (
 	github.com/cockroachdb/swiss v0.0.0-20240612210725-f4de07ae6964
 	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13
 	github.com/dicedb/go-dice v0.0.0-20240820180649-d97f15fca831
-	github.com/ohler55/ojg v1.23.0
+	github.com/ohler55/ojg v1.24.0
 	github.com/pelletier/go-toml/v2 v2.2.3
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0

go.sum

@@ -71,6 +71,8 @@ github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo
 github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8=
 github.com/ohler55/ojg v1.23.0 h1:xjJasLaKf4dKkyJq0CNXQMRdL7F1172tms885aPKcS0=
 github.com/ohler55/ojg v1.23.0/go.mod h1:gQhDVpQLqrmnd2eqGAvJtn+NfKoYJbe/A4Sj3/Vro4o=
+github.com/ohler55/ojg v1.24.0 h1:y2AVez6fPTszK/jPhaAYMCAzAoSleConMqSDD5wJKJg=
+github.com/ohler55/ojg v1.24.0/go.mod h1:gQhDVpQLqrmnd2eqGAvJtn+NfKoYJbe/A4Sj3/Vro4o=
 github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
 github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

integration_tests/commands/json_test.go

@@ -2,11 +2,12 @@ package commands
 
 import (
 	"fmt"
-	"github.com/bytedance/sonic"
 	"net"
 	"strings"
 	"testing"
 
+	"github.com/bytedance/sonic"
+
 	"github.com/dicedb/dice/testutils"
 	"gotest.tools/v3/assert"
 )
@@ -606,15 +607,13 @@ func TestJSONForgetOperations(t *testing.T) {
 		})
 	}
 }
-
-func arraysArePermutations(a, b []interface{}) bool {
-	// If lengths are different, they cannot be permutations
+func arraysArePermutations[T comparable](a, b []T) bool {
 	if len(a) != len(b) {
 		return false
 	}
 
 	// Count occurrences of each element in array 'a'
-	countA := make(map[interface{}]int)
+	countA := make(map[T]int)
 	for _, elem := range a {
 		countA[elem]++
 	}
@@ -841,3 +840,98 @@ func TestJsonARRAPPEND(t *testing.T) {
 		})
 	}
 }
+func convertToArray(input string) []string {
+	input = strings.Trim(input, `"[`)
+	input = strings.Trim(input, `]"`)
+	elements := strings.Split(input, ",")
+	for i, element := range elements {
+		elements[i] = strings.TrimSpace(element)
+	}
+	return elements
+}
+func TestJSONNumIncrBy(t *testing.T) {
+	conn := getLocalConnection()
+	defer conn.Close()
+	invalidArgMessage := "ERR wrong number of arguments for 'json.numincrby' command"
+	testCases := []struct {
+		name        string
+		setupData   string
+		commands    []string
+		expected    []interface{}
+		assert_type []string
+		cleanUp     []string
+	}{
+		{
+			name:        "Invalid number of arguments",
+			setupData:   "",
+			commands:    []string{"JSON.NUMINCRBY ", "JSON.NUMINCRBY foo", "JSON.NUMINCRBY foo $"},
+			expected:    []interface{}{invalidArgMessage, invalidArgMessage, invalidArgMessage},
+			assert_type: []string{"equal", "equal", "equal"},
+			cleanUp:     []string{},
+		},
+		{
+			name:        "Non-existant key",
+			setupData:   "",
+			commands:    []string{"JSON.NUMINCRBY foo $ 1"},
+			expected:    []interface{}{"ERR could not perform this operation on a key that doesn't exist"},
+			assert_type: []string{"equal"},
+			cleanUp:     []string{},
+		},
+		{
+			name:        "Invalid value of increment",
+			setupData:   "JSON.SET foo $ 1",
+			commands:    []string{"JSON.GET foo $", "JSON.NUMINCRBY foo $ @", "JSON.NUMINCRBY foo $ 122@"},
+			expected:    []interface{}{"1", "ERR expected value at line 1 column 1", "ERR trailing characters at line 1 column 4"},
+			assert_type: []string{"equal", "equal", "equal"},
+			cleanUp:     []string{"DEL foo"},
+		},
+		{
+			name:        "incrby at non root path",
+			setupData:   fmt.Sprintf("JSON.SET %s $ %s", "foo", `{"a":"b","b":[{"a":2.2},{"a":5},{"a":"c"}]}`),
+			commands:    []string{"JSON.NUMINCRBY foo $..a 2", "JSON.NUMINCRBY foo $.a 2", "JSON.GET foo", "JSON.NUMINCRBY foo $..a -2", "JSON.GET foo"},
+			expected:    []interface{}{"[null,4.2,7,null]", "[null]", "{\"a\":\"b\",\"b\":[{\"a\":4.2},{\"a\":7},{\"a\":\"c\"}]}", "[null,2.2,5,null]", "{\"a\":\"b\",\"b\":[{\"a\":2.2},{\"a\":5},{\"a\":\"c\"}]}"},
+			assert_type: []string{"perm_equal", "perm_equal", "equal", "perm_equal", "equal"},
+			cleanUp:     []string{"DEL foo"},
+		},
+		{
+			name:        "incrby at root path",
+			setupData:   "JSON.SET foo $ 1",
+			commands:    []string{"JSON.NUMINCRBY foo $ 1", "JSON.GET foo $", "JSON.NUMINCRBY foo $ -1", "JSON.GET foo $"},
+			expected:    []interface{}{"[2]", "2", "[1]", "1"},
+			assert_type: []string{"equal", "equal", "equal", "equal"},
+			cleanUp:     []string{"DEL foo"},
+		},
+		{
+			name:        "incrby at root path",
+			setupData:   "JSON.SET foo $ 1",
+			commands:    []string{"expire foo 10", "JSON.NUMINCRBY foo $ 1", "ttl foo", "JSON.GET foo $", "JSON.NUMINCRBY foo $ -1", "JSON.GET foo $"},
+			expected:    []interface{}{int64(1), "[2]", int64(10), "2", "[1]", "1"},
+			assert_type: []string{"equal", "equal", "range", "equal", "equal", "equal"},
+			cleanUp:     []string{"DEL foo"},
+		},
+	}
+
+	for _, tc := range testCases {
+		FireCommand(conn, "DEL foo")
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.setupData != "" {
+				assert.Equal(t, FireCommand(conn, tc.setupData), "OK")
+			}
+			for i := 0; i < len(tc.commands); i++ {
+				cmd := tc.commands[i]
+				out := tc.expected[i]
+				result := FireCommand(conn, cmd)
+				if tc.assert_type[i] == "equal" {
+					assert.Equal(t, out, result)
+				} else if tc.assert_type[i] == "perm_equal" {
+					assert.Assert(t, arraysArePermutations(convertToArray(out.(string)), convertToArray(result.(string))))
+				} else if tc.assert_type[i] == "range" {
+					assert.Assert(t, result.(int64) <= tc.expected[i].(int64) && result.(int64) > 0, "Expected %v to be within 0 to %v", result, tc.expected[i])
+				}
+			}
+			for i := 0; i < len(tc.cleanUp); i++ {
+				FireCommand(conn, tc.cleanUp[i])
+			}
+		})
+	}
+}

internal/eval/commands.go

@@ -721,6 +721,13 @@ var (
 		Eval:  evalSELECT,
 		Arity: 1,
 	}
+	jsonnumincrbyCmdMeta = DiceCmdMeta{
+		Name:     "JSON.NUMINCRBY",
+		Info:     `Increment the number value stored at path by number.`,
+		Eval:     evalJSONNUMINCRBY,
+		Arity:    3,
+		KeySpecs: KeySpecs{BeginIndex: 1},
+	}
 )
 
 func init() {
@@ -806,6 +813,7 @@ func init() {
 	DiceCmds["JSON.MGET"] = jsonMGetCmdMeta
 	DiceCmds["HLEN"] = hlenCmdMeta
 	DiceCmds["SELECT"] = selectCmdMeta
+	DiceCmds["JSON.NUMINCRBY"] = jsonnumincrbyCmdMeta
 }
 
 // Function to convert DiceCmdMeta to []interface{}

internal/eval/constants.go

@@ -25,4 +25,7 @@ const (
 	Async   string = "ASYNC"
 	Help    string = "HELP"
 	Memory  string = "MEMORY"
+	Null    string = "null"
+	null    string = "null"
+	NULL    string = "null"
 )

internal/eval/eval.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 	"syscall"
 	"time"
+	"unicode"
 	"unsafe"
 
 	"github.com/dicedb/dice/internal/object"
@@ -3174,3 +3175,131 @@ func evalSELECT(args []string, store *dstore.Store) []byte {
 
 	return clientio.RespOK
 }
+
+func formatFloat(f float64, b bool) string {
+	formatted := strconv.FormatFloat(f, 'f', -1, 64)
+	if b {
+		parts := strings.Split(formatted, ".")
+		if len(parts) == 1 {
+			formatted += ".0"
+		}
+	}
+	return formatted
+}
+
+// takes original value, increment values (float or int), a flag representing if increment is float
+// returns new value, string representation, a boolean representing if the value was modified
+func incrementValue(value any, isIncrFloat bool, incrFloat float64, incrInt int64) (newVal interface{}, stringRepresentation string, isModified bool) {
+	switch utils.GetJSONFieldType(value) {
+	case utils.NumberType:
+		oldVal := value.(float64)
+		var newVal float64
+		if isIncrFloat {
+			newVal = oldVal + incrFloat
+		} else {
+			newVal = oldVal + float64(incrInt)
+		}
+		resultString := formatFloat(newVal, isIncrFloat)
+		return newVal, resultString, true
+	case utils.IntegerType:
+		if isIncrFloat {
+			oldVal := float64(value.(int64))
+			newVal := oldVal + incrFloat
+			resultString := formatFloat(newVal, isIncrFloat)
+			return newVal, resultString, true
+		} else {
+			oldVal := value.(int64)
+			newVal := oldVal + incrInt
+			resultString := fmt.Sprintf("%d", newVal)
+			return newVal, resultString, true
+		}
+	default:
+		return value, null, false
+	}
+}
+
+func evalJSONNUMINCRBY(args []string, store *dstore.Store) []byte {
+	if len(args) < 3 {
+		return diceerrors.NewErrArity("JSON.NUMINCRBY")
+	}
+	key := args[0]
+	obj := store.Get(key)
+
+	if obj == nil {
+		return diceerrors.NewErrWithFormattedMessage("-ERR could not perform this operation on a key that doesn't exist")
+	}
+
+	// Check if the object is of JSON type
+	errWithMessage := object.AssertTypeAndEncoding(obj.TypeEncoding, object.ObjTypeJSON, object.ObjEncodingJSON)
+	if errWithMessage != nil {
+		return errWithMessage
+	}
+
+	path := args[1]
+
+	jsonData := obj.Value
+	// Parse the JSONPath expression
+	expr, err := jp.ParseString(path)
+
+	if err != nil {
+		return diceerrors.NewErrWithMessage("invalid JSONPath")
+	}
+
+	isIncrFloat := false
+
+	for i, r := range args[2] {
+		if !unicode.IsDigit(r) && r != '.' && r != '-' {
+			if i == 0 {
+				return diceerrors.NewErrWithFormattedMessage("-ERR expected value at line 1 column %d", i+1)
+			}
+			return diceerrors.NewErrWithFormattedMessage("-ERR trailing characters at line 1 column %d", i+1)
+		}
+		if r == '.' {
+			isIncrFloat = true
+		}
+	}
+	var incrFloat float64
+	var incrInt int64
+	if isIncrFloat {
+		incrFloat, err = strconv.ParseFloat(args[2], 64)
+		if err != nil {
+			return diceerrors.NewErrWithMessage(diceerrors.IntOrOutOfRangeErr)
+		}
+	} else {
+		incrInt, err = strconv.ParseInt(args[2], 10, 64)
+		if err != nil {
+			return diceerrors.NewErrWithMessage(diceerrors.IntOrOutOfRangeErr)
+		}
+	}
+	results := expr.Get(jsonData)
+
+	if len(results) == 0 {
+		respString := "[]"
+		return clientio.Encode(respString, false)
+	}
+
+	resultArray := make([]string, 0, len(results))
+
+	if path == defaultRootPath {
+		newValue, resultString, isModified := incrementValue(jsonData, isIncrFloat, incrFloat, incrInt)
+		if isModified {
+			jsonData = newValue
+		}
+		resultArray = append(resultArray, resultString)
+	} else {
+		// Execute the JSONPath query
+		_, err := expr.Modify(jsonData, func(value any) (interface{}, bool) {
+			newValue, resultString, isModified := incrementValue(value, isIncrFloat, incrFloat, incrInt)
+			resultArray = append(resultArray, resultString)
+			return newValue, isModified
+		})
+		if err != nil {
+			return diceerrors.NewErrWithMessage("invalid JSONPath")
+		}
+	}
+
+	resultString := `[` + strings.Join(resultArray, ",") + `]`
+
+	obj.Value = jsonData
+	return clientio.Encode(resultString, false)
+}