Add support for command JSON.STRAPPEND DiceDB#501

integration_tests/commands/async/json_test.go

@@ -1378,3 +1378,62 @@ func TestJsonARRTRIM(t *testing.T) {
 		})
 	}
 }
+
+func TestJsonSTRAPPEND(t *testing.T) {
+	conn := getLocalConnection()
+	defer conn.Close()
+
+	simpleJSON := `{"name":"John","age":30}`
+
+	testCases := []struct {
+		name     string
+		setCmd   string
+		getCmd   string
+		expected interface{}
+	}{
+		{
+			name:     "STRAPPEND to nested string",
+			setCmd:   `JSON.SET doc $ ` + simpleJSON,
+			getCmd:   `JSON.STRAPPEND doc $.name " Doe"`,
+			expected: []interface{}{int64(8)},
+		},
+		{
+			name:     "STRAPPEND to multiple paths",
+			setCmd:   `JSON.SET doc $ ` + simpleJSON,
+			getCmd:   `JSON.STRAPPEND doc $..name "baz"`,
+			expected: []interface{}{int64(7)},
+		},
+		{
+			name:     "STRAPPEND to non-string",
+			setCmd:   `JSON.SET doc $ ` + simpleJSON,
+			getCmd:   `JSON.STRAPPEND doc $.age " years"`,
+			expected: []interface{}{"(nil)"},
+		},
+		{
+			name:     "STRAPPEND with empty string",
+			setCmd:   `JSON.SET doc $ ` + simpleJSON,
+			getCmd:   `JSON.STRAPPEND doc $.name ""`,
+			expected: []interface{}{int64(4)},
+		},
+		{
+			name:     "STRAPPEND to non-existent path",
+			setCmd:   `JSON.SET doc $ ` + simpleJSON,
+			getCmd:   `JSON.STRAPPEND doc $.nonexistent " test"`,
+			expected: []interface{}{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			FireCommand(conn, "DEL doc")
+
+			result := FireCommand(conn, tc.setCmd)
+
+			assert.Equal(t, "OK", result)
+
+			result = FireCommand(conn, tc.getCmd)
+			assert.DeepEqual(t, tc.expected, result)
+		})
+	}
+
+}

internal/eval/commands.go

@@ -971,6 +971,16 @@ var (
 		Arity:    -4,
 		KeySpecs: KeySpecs{BeginIndex: 1},
 	}
+	jsonstrappendCmdMeta = DiceCmdMeta{
+		Name: "JSON.STRAPPEND",
+		Info: `JSON.STRAPPEND key [path] value
+		Append the JSON string values to the string at path
+		Returns an array of integer replies for each path, the string's new length, or nil, if the matching JSON value is not a string. 
+		Error reply: If the value at path is not a string or if the key doesn't exist.`,
+		Eval:     evalJSONSTRAPPEND,
+		Arity:    -3,
+		KeySpecs: KeySpecs{BeginIndex: 1},
+	}
 )
 
 func init() {
@@ -1079,6 +1089,7 @@ func init() {
 	DiceCmds["ZADD"] = zaddCmdMeta
 	DiceCmds["ZRANGE"] = zrangeCmdMeta
 	DiceCmds["HINCRBYFLOAT"] = hincrbyFloatCmdMeta
+	DiceCmds["JSON.STRAPPEND"] = jsonstrappendCmdMeta
 }
 
 // Function to convert DiceCmdMeta to []interface{}

internal/eval/eval.go

@@ -1078,16 +1078,68 @@ func evalJSONGET(args []string, store *dstore.Store) []byte {
 	}
 
 	key := args[0]
-	// Default path is root if not specified
-	path := defaultRootPath
-	if len(args) > 1 {
-		path = args[1]
+	obj := store.Get(key)
+
+	if obj == nil {
+		return clientio.RespNIL
 	}
-	result, err := jsonGETHelper(store, path, key)
-	if err != nil {
-		return err
+
+	// Check if the object is of JSON type
+	errWithMessage := object.AssertTypeAndEncoding(obj.TypeEncoding, object.ObjTypeJSON, object.ObjEncodingJSON)
+	if errWithMessage != nil {
+		return errWithMessage
 	}
-	return clientio.Encode(result, false)
+
+	jsonData := obj.Value
+
+	// If no path is specified, return the entire JSON document
+	if len(args) == 1 {
+		jsonString, err := sonic.MarshalString(jsonData)
+		if err != nil {
+			return diceerrors.NewErrWithMessage("Failed to marshal JSON")
+		}
+		return clientio.Encode(jsonString, false)
+	}
+
+	// Handle paths
+	results := make([]interface{}, 0)
+	for _, path := range args[1:] {
+		if path == "$" {
+			// Root path, return the entire document
+			jsonString, err := sonic.MarshalString(jsonData)
+			if err != nil {
+				return diceerrors.NewErrWithMessage("Failed to marshal JSON")
+			}
+			results = append(results, jsonString)
+		} else {
+			// Parse the JSONPath expression
+			expr, err := jp.ParseString(path)
+			if err != nil {
+				results = append(results, clientio.RespNIL)
+				continue
+			}
+
+			// Execute the JSONPath query
+			pathResults := expr.Get(jsonData)
+			if len(pathResults) == 0 {
+				results = append(results, clientio.RespNIL)
+			} else {
+				jsonString, err := sonic.MarshalString(pathResults[0])
+				if err != nil {
+					return diceerrors.NewErrWithMessage("Failed to marshal JSON")
+				}
+				results = append(results, jsonString)
+			}
+		}
+	}
+
+	// If only one path was specified, return the result directly
+	if len(results) == 1 {
+		return clientio.Encode(results[0], false)
+	}
+
+	// Otherwise, return an array of results
+	return clientio.Encode(results, false)
 }
 
 // helper function used by evalJSONGET and evalJSONMGET to prepare the results
@@ -4534,3 +4586,72 @@ func evalHINCRBYFLOAT(args []string, store *dstore.Store) []byte {
 
 	return clientio.Encode(numkey, false)
 }
+
+// evalJSONSTRAPPEND appends a string value to the JSON string value at the specified path
+// in the JSON object saved at the key in arguments.
+// Args must contain at least a key and the string value to append.
+// If the key does not exist or is expired, it returns an error response.
+// If the value at the specified path is not a string, it returns an error response.
+// Returns the new length of the string at the specified path if successful.
+func evalJSONSTRAPPEND(args []string, store *dstore.Store) []byte {
+	if len(args) != 3 {
+		return diceerrors.NewErrArity("JSON.STRAPPEND")
+	}
+
+	key := args[0]
+	path := args[1]
+	value := args[2]
+
+	obj := store.Get(key)
+	if obj == nil {
+		return clientio.Encode([]interface{}{}, false)
+	}
+
+	errWithMessage := object.AssertTypeAndEncoding(obj.TypeEncoding, object.ObjTypeJSON, object.ObjEncodingJSON)
+	if errWithMessage != nil {
+		return errWithMessage
+	}
+
+	jsonData := obj.Value
+
+	var resultsArray []interface{}
+
+	if path == "$" {
+		// Handle root-level string
+		if str, ok := jsonData.(string); ok {
+			unquotedValue := strings.Trim(value, "\"")
+			newValue := str + unquotedValue
+			resultsArray = append(resultsArray, int64(len(newValue)))
+			jsonData = newValue
+		} else {
+			return clientio.Encode([]interface{}{}, false)
+		}
+	} else {
+		expr, err := jp.ParseString(path)
+		if err != nil {
+			return clientio.Encode([]interface{}{}, false)
+		}
+
+		newData, modifyErr := expr.Modify(jsonData, func(data any) (interface{}, bool) {
+			switch v := data.(type) {
+			case string:
+				unquotedValue := strings.Trim(value, "\"")
+				newValue := v + unquotedValue
+				resultsArray = append([]interface{}{int64(len(newValue))}, resultsArray...)
+				return newValue, true
+			default:
+				resultsArray = append([]interface{}{clientio.RespNIL}, resultsArray...)
+				return data, false
+			}
+		})
+
+		if modifyErr != nil {
+			return clientio.Encode([]interface{}{}, false)
+		}
+		jsonData = newData
+	}
+
+	store.Put(key, store.NewObj(jsonData, -1, object.ObjTypeJSON, object.ObjEncodingJSON))
+
+	return clientio.Encode(resultsArray, false)
+}

internal/eval/eval_test.go

@@ -98,6 +98,7 @@ func TestEval(t *testing.T) {
 	testEvalZRANGE(t, store)
 	testEvalHVALS(t, store)
 	testEvalHINCRBYFLOAT(t, store)
+	testEvalJSONSTRAPPEND(t, store)
 }
 
 func testEvalPING(t *testing.T, store *dstore.Store) {
@@ -4726,3 +4727,84 @@ func BenchmarkEvalHINCRBYFLOAT(b *testing.B) {
 		})
 	}
 }
+
+func testEvalJSONSTRAPPEND(t *testing.T, store *dstore.Store) {
+	tests := map[string]evalTestCase{
+		"append to multiple fields": {
+			setup: func() {
+				key := "doc1"
+				value := "{\"a\":\"foo\", \"nested1\": {\"a\": \"hello\"}, \"nested2\": {\"a\": 31}}"
+				var rootData interface{}
+				_ = sonic.Unmarshal([]byte(value), &rootData)
+				obj := store.NewObj(rootData, -1, object.ObjTypeJSON, object.ObjEncodingJSON)
+				store.Put(key, obj)
+			},
+			input:  []string{"doc1", "$..a", "\"bar\""},
+			output: []byte("*3\r\n:6\r\n:8\r\n$-1\r\n"), // Expected lengths after append
+		},
+		"append to single field": {
+			setup: func() {
+				key := "doc1"
+				value := "{\"a\":\"foo\", \"nested1\": {\"a\": \"hello\"}, \"nested2\": {\"a\": 31}}"
+				var rootData interface{}
+				_ = sonic.Unmarshal([]byte(value), &rootData)
+				obj := store.NewObj(rootData, -1, object.ObjTypeJSON, object.ObjEncodingJSON)
+				store.Put(key, obj)
+			},
+			input:  []string{"doc1", "$.nested1.a", "\"baz\""},
+			output: []byte("*1\r\n:11\r\n"), // Expected length after append
+		},
+		"append to non-existing key": {
+			setup: func() {
+				// No setup needed as we are testing a non-existing document.
+			},
+			input:  []string{"non_existing_doc", "$..a", "\"err\""},
+			output: []byte("*1\r\n$-1\r\n"), // Expect an error, but captured in the test framework
+		},
+		"legacy path append": {
+			setup: func() {
+				key := "doc1"
+				value := "{\"a\":\"foo\", \"nested1\": {\"a\": \"hello\"}, \"nested2\": {\"a\": 31}}"
+				var rootData interface{}
+				_ = sonic.Unmarshal([]byte(value), &rootData)
+				obj := store.NewObj(rootData, -1, object.ObjTypeJSON, object.ObjEncodingJSON)
+				store.Put(key, obj)
+			},
+			input:  []string{"doc1", ".*.a", "\"bar\""},
+			output: []byte("*1\r\n:8\r\n"), // Expected length after append for nested1.a
+		},
+		"append to root node": {
+			setup: func() {
+				key := "doc1"
+				value := "\"abcd\""
+				var rootData interface{}
+				_ = sonic.Unmarshal([]byte(value), &rootData)
+				obj := store.NewObj(rootData, -1, object.ObjTypeJSON, object.ObjEncodingJSON)
+				store.Put(key, obj)
+			},
+			input:  []string{"doc1", "$", "\"piu\""},
+			output: []byte("*1\r\n:7\r\n"), // Expected length after appending to "abcd"
+		},
+	}
+
+	// Run the tests
+	runEvalTests(t, tests, evalJSONSTRAPPEND, store)
+}
+
+func BenchmarkEvalJSONSTRAPPEND(b *testing.B) {
+	store := dstore.NewStore(nil)
+
+	// Setup a sample JSON document
+	key := "doc1"
+	value := "{\"a\":\"foo\", \"nested1\": {\"a\": \"hello\"}, \"nested2\": {\"a\": 31}}"
+	var rootData interface{}
+	_ = sonic.Unmarshal([]byte(value), &rootData)
+	obj := store.NewObj(rootData, -1, object.ObjTypeJSON, object.ObjEncodingJSON)
+	store.Put(key, obj)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// Benchmark appending to multiple fields
+		evalJSONSTRAPPEND([]string{"doc1", "$..a", "\"bar\""}, store)
+	}
+}