From 9b144d60b86fc181d9cfb85e205793cb921e1213 Mon Sep 17 00:00:00 2001 From: Karan Dixit Date: Mon, 30 Sep 2024 16:18:12 +0530 Subject: [PATCH] Add support for command JSON.STRAPPEND #501 --- integration_tests/commands/async/json_test.go | 59 ++++++++ internal/eval/commands.go | 11 ++ internal/eval/eval.go | 137 +++++++++++++++++- internal/eval/eval_test.go | 82 +++++++++++ 4 files changed, 281 insertions(+), 8 deletions(-) diff --git a/integration_tests/commands/async/json_test.go b/integration_tests/commands/async/json_test.go index a92e60b656..10c4df024e 100644 --- a/integration_tests/commands/async/json_test.go +++ b/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) + }) + } + +} diff --git a/internal/eval/commands.go b/internal/eval/commands.go index c5157f3176..3eb52c6045 100644 --- a/internal/eval/commands.go +++ b/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{} diff --git a/internal/eval/eval.go b/internal/eval/eval.go index a4ad191626..aebaa806ab 100644 --- a/internal/eval/eval.go +++ b/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) +} diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 81847ef60e..46824930d4 100644 --- a/internal/eval/eval_test.go +++ b/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) + } +}