add json functions: json_{set/insert/replace} and json_merge. (pingca…

…p#3374)

util/types/json/functions.go

@@ -131,7 +131,11 @@ func extract(j JSON, pathExpr PathExpression) (ret []JSON) {
 		return []JSON{j}
 	}
 	currentLeg, subPathExpr := pathExpr.popOneLeg()
-	if currentLeg.typ == pathLegIndex && j.typeCode == typeCodeArray {
+	if currentLeg.typ == pathLegIndex {
+		// If j is not an array, autowrap that into array.
+		if j.typeCode != typeCodeArray {
+			j = autoWrapAsArray(j, 1)
+		}
 		if currentLeg.arrayIndex == arrayIndexAsterisk {
 			for _, child := range j.array {
 				ret = append(ret, extract(child, subPathExpr)...)
@@ -165,3 +169,135 @@ func extract(j JSON, pathExpr PathExpression) (ret []JSON) {
 	}
 	return
 }
+
+// autoWrapAsArray wraps input JSON into an array if needed.
+func autoWrapAsArray(j JSON, hintLength int) JSON {
+	jnew := CreateJSON(nil)
+	jnew.typeCode = typeCodeArray
+	jnew.array = make([]JSON, 0, hintLength)
+	jnew.array = append(jnew.array, j)
+	return jnew
+}
+
+// Merge merges suffixes into j according the following rules:
+// 1) adjacent arrays are merged to a single array;
+// 2) adjacent object are merged to a single object;
+// 3) a scalar value is autowrapped as an array before merge;
+// 4) an adjacent array and object are merged by autowrapping the object as an array.
+func (j JSON) Merge(suffixes []JSON) JSON {
+	if j.typeCode != typeCodeArray && j.typeCode != typeCodeObject {
+		j = autoWrapAsArray(j, len(suffixes)+1)
+	}
+	for i := 0; i < len(suffixes); i++ {
+		suffix := suffixes[i]
+		switch j.typeCode {
+		case typeCodeArray:
+			if suffix.typeCode == typeCodeArray {
+				// rule (1)
+				j.array = append(j.array, suffix.array...)
+			} else {
+				// rule (3), (4)
+				j.array = append(j.array, suffix)
+			}
+		case typeCodeObject:
+			if suffix.typeCode == typeCodeObject {
+				// rule (2)
+				for key := range suffix.object {
+					if child, ok := j.object[key]; ok {
+						j.object[key] = child.Merge([]JSON{suffix.object[key]})
+					} else {
+						j.object[key] = suffix.object[key]
+					}
+				}
+			} else {
+				// rule (4)
+				j = autoWrapAsArray(j, len(suffixes)+1-i)
+				if suffix.typeCode == typeCodeArray {
+					j.array = append(j.array, suffix.array...)
+				} else {
+					j.array = append(j.array, suffix)
+				}
+			}
+		}
+	}
+	return j
+}
+
+// ModifyType is for modify a JSON. There are three valid values:
+// ModifyInsert, ModifyReplace and ModifySet.
+type ModifyType byte
+
+const (
+	// ModifyInsert is for insert a new element into a JSON.
+	ModifyInsert ModifyType = 0x01
+	// ModifyReplace is for replace an old elemList from a JSON.
+	ModifyReplace ModifyType = 0x02
+	// ModifySet = ModifyInsert | ModifyReplace
+	ModifySet ModifyType = 0x03
+)
+
+// Modify modifies a JSON object by insert, replace or set.
+// All path expressions cannot contain * or ** wildcard.
+// If any error occurs, the input won't be changed.
+func (j JSON) Modify(pathExprList []PathExpression, values []JSON, mt ModifyType) (retj JSON, err error) {
+	if len(pathExprList) != len(values) {
+		// TODO should return 1582(42000)
+		return retj, errors.New("Incorrect parameter count")
+	}
+	for _, pathExpr := range pathExprList {
+		if pathExpr.flags.containsAnyAsterisk() {
+			// TODO should return 3149(42000)
+			return retj, errors.New("Invalid path expression")
+		}
+	}
+	for i := 0; i < len(pathExprList); i++ {
+		pathExpr, value := pathExprList[i], values[i]
+		j = set(j, pathExpr, value, mt)
+	}
+	return j, nil
+}
+
+// set is for Modify. The result JSON maybe share something with input JSON.
+func set(j JSON, pathExpr PathExpression, value JSON, mt ModifyType) JSON {
+	if len(pathExpr.legs) == 0 {
+		if mt&ModifyReplace != 0 {
+			return value
+		}
+		return j
+	}
+	currentLeg, subPathExpr := pathExpr.popOneLeg()
+	if currentLeg.typ == pathLegIndex {
+		// If j is not an array, we should autowrap that as array.
+		// Then if its length equals to 1, we unwrap it back.
+		var shouldUnwrap = false
+		if j.typeCode != typeCodeArray {
+			j = autoWrapAsArray(j, 1)
+			shouldUnwrap = true
+		}
+		var index = currentLeg.arrayIndex
+		if len(j.array) > index {
+			// e.g. json_replace('[1, 2, 3]', '$[0]', "x") => '["x", 2, 3]'
+			j.array[index] = set(j.array[index], subPathExpr, value, mt)
+		} else if len(subPathExpr.legs) == 0 && mt&ModifyInsert != 0 {
+			// e.g. json_insert('[1, 2, 3]', '$[3]', "x") => '[1, 2, 3, "x"]'
+			j.array = append(j.array, value)
+		}
+		if len(j.array) == 1 && shouldUnwrap {
+			j = j.array[0]
+		}
+	} else if currentLeg.typ == pathLegKey && j.typeCode == typeCodeObject {
+		var key = currentLeg.dotKey
+		if child, ok := j.object[key]; ok {
+			// e.g. json_replace('{"a": 1}', '$.a', 2) => '{"a": 2}'
+			j.object[key] = set(child, subPathExpr, value, mt)
+		} else if len(subPathExpr.legs) == 0 && mt&ModifyInsert != 0 {
+			// e.g. json_insert('{"a": 1}', '$.b', 2) => '{"a": 1, "b": 2}'
+			j.object[key] = value
+		}
+	}
+	// For these cases, we just return the input JSON back without any change:
+	// 1) we want to insert a new element, but the full path has already exists;
+	// 2) we want to replace an old element, but the full path doesn't exist;
+	// 3) we want to insert or replace something, but the path without last leg doesn't exist.
+	return j
+}

util/types/json/functions_test.go

@@ -38,7 +38,7 @@ func (s *testJSONSuite) TestJSONType(c *C) {
 }
 
 func (s *testJSONSuite) TestJSONExtract(c *C) {
-	j1 := mustParseFromString(`{"a": [1, "2", {"aa": "bb"}, 4.0, {"aa": "cc"}], "b": true, "c": ["d"], "\"hello\"": "world"}`)
+	j1 := mustParseFromString(`{"\"hello\"": "world", "a": [1, "2", {"aa": "bb"}, 4.0, {"aa": "cc"}], "b": true, "c": ["d"]}`)
 	j2 := mustParseFromString(`[{"a": 1, "b": true}, 3, 3.5, "hello, world", null, true]`)
 
 	var tests = []struct {
@@ -51,17 +51,17 @@ func (s *testJSONSuite) TestJSONExtract(c *C) {
 		// test extract with only one path expression.
 		{j1, []string{"$.a"}, j1.object["a"], true, nil},
 		{j2, []string{"$.a"}, CreateJSON(nil), false, nil},
-		{j1, []string{"$[0]"}, CreateJSON(nil), false, nil},
+		{j1, []string{"$[0]"}, j1, true, nil}, // in Extract, autowraped j1 as an array.
 		{j2, []string{"$[0]"}, j2.array[0], true, nil},
 		{j1, []string{"$.a[2].aa"}, CreateJSON("bb"), true, nil},
 		{j1, []string{"$.a[*].aa"}, mustParseFromString(`["bb", "cc"]`), true, nil},
-		{j1, []string{"$.*[0]"}, mustParseFromString(`[1, "d"]`), true, nil},
+		{j1, []string{"$.*[0]"}, mustParseFromString(`["world", 1, true, "d"]`), true, nil},
 		{j1, []string{`$.a[*]."aa"`}, mustParseFromString(`["bb", "cc"]`), true, nil},
 		{j1, []string{`$."\"hello\""`}, mustParseFromString(`"world"`), true, nil},
-		{j1, []string{`$**[0]`}, mustParseFromString(`[1, "d"]`), true, nil},
+		{j1, []string{`$**[1]`}, mustParseFromString(`"2"`), true, nil},
 
 		// test extract with multi path expressions.
-		{j1, []string{"$.a", "$[0]"}, mustParseFromString(`[[1, "2", {"aa": "bb"}, 4.0, {"aa": "cc"}]]`), true, nil},
+		{j1, []string{"$.a", "$[5]"}, mustParseFromString(`[[1, "2", {"aa": "bb"}, 4.0, {"aa": "cc"}]]`), true, nil},
 		{j2, []string{"$.a", "$[0]"}, mustParseFromString(`[{"a": 1, "b": true}]`), true, nil},
 	}
 
@@ -103,3 +103,77 @@ func (s *testJSONSuite) TestJSONUnquote(c *C) {
 		c.Assert(unquoted, Equals, tt.unquoted)
 	}
 }
+
+func (s *testJSONSuite) TestJSONMerge(c *C) {
+	var tests = []struct {
+		base     string
+		suffixes []string
+		expected string
+	}{
+		{`{"a": 1}`, []string{`{"b": 2}`}, `{"a": 1, "b": 2}`},
+		{`{"a": 1}`, []string{`{"a": 2}`}, `{"a": [1, 2]}`},
+		{`[1]`, []string{`[2]`}, `[1, 2]`},
+		{`{"a": 1}`, []string{`[1]`}, `[{"a": 1}, 1]`},
+		{`[1]`, []string{`{"a": 1}`}, `[1, {"a": 1}]`},
+		{`{"a": 1}`, []string{`4`}, `[{"a": 1}, 4]`},
+		{`[1]`, []string{`4`}, `[1, 4]`},
+		{`4`, []string{`{"a": 1}`}, `[4, {"a": 1}]`},
+		{`4`, []string{`1`}, `[4, 1]`},
+	}
+
+	for _, tt := range tests {
+		base := mustParseFromString(tt.base)
+		suffixes := make([]JSON, 0, len(tt.suffixes))
+		for _, s := range tt.suffixes {
+			suffixes = append(suffixes, mustParseFromString(s))
+		}
+		base = base.Merge(suffixes)
+		cmp, err := CompareJSON(base, mustParseFromString(tt.expected))
+		c.Assert(err, IsNil)
+		c.Assert(cmp, Equals, 0)
+	}
+}
+
+func (s *testJSONSuite) TestJSONModify(c *C) {
+	var tests = []struct {
+		base     string
+		setField string
+		setValue string
+		mt       ModifyType
+		expected string
+	}{
+		{`null`, "$", `{}`, ModifySet, `{}`},
+		{`{}`, "$.a", `3`, ModifySet, `{"a": 3}`},
+		{`{"a": 3}`, "$.a", `[]`, ModifyReplace, `{"a": []}`},
+		{`{"a": []}`, "$.a[0]", `3`, ModifySet, `{"a": [3]}`},
+		{`{"a": [3]}`, "$.a[1]", `4`, ModifyInsert, `{"a": [3, 4]}`},
+		{`{"a": [3]}`, "$[0]", `4`, ModifySet, `4`},
+		{`{"a": [3]}`, "$[1]", `4`, ModifySet, `[{"a": [3]}, 4]`},
+
+		// nothing changed because the path is empty and we want to insert.
+		{`{}`, "$", `1`, ModifyInsert, `{}`},
+		// nothing changed because the path without last leg doesn't exist.
+		{`{"a": [3, 4]}`, "$.b[1]", `3`, ModifySet, `{"a": [3, 4]}`},
+		// nothing changed because the path without last leg doesn't exist.
+		{`{"a": [3, 4]}`, "$.a[2].b", `3`, ModifySet, `{"a": [3, 4]}`},
+		// nothing changed because we want to insert but the full path exists.
+		{`{"a": [3, 4]}`, "$.a[0]", `30`, ModifyInsert, `{"a": [3, 4]}`},
+		// nothing changed because we want to replace but the full path doesn't exist.
+		{`{"a": [3, 4]}`, "$.a[2]", `30`, ModifyReplace, `{"a": [3, 4]}`},
+	}
+	for _, tt := range tests {
+		pathExpr, err := ParseJSONPathExpr(tt.setField)
+		c.Assert(err, IsNil)
+
+		base := mustParseFromString(tt.base)
+		value := mustParseFromString(tt.setValue)
+		expected := mustParseFromString(tt.expected)
+
+		obtain, err := base.Modify([]PathExpression{pathExpr}, []JSON{value}, tt.mt)
+		c.Assert(err, IsNil)
+
+		cmp, err := CompareJSON(obtain, expected)
+		c.Assert(err, IsNil)
+		c.Assert(cmp, Equals, 0)
+	}
+}

util/types/json/json.go

@@ -66,7 +66,8 @@ func CreateJSON(in interface{}) JSON {
 
 // ParseFromString parses a json from string.
 func ParseFromString(s string) (j JSON, err error) {
-	// TODO: implement the decoder directly.
+	// TODO: implement the decoder directly. It's important for keeping
+	// keys in object have same order with the original string.
 	if len(s) == 0 {
 		err = ErrInvalidJSONText.GenByArgs("The document is empty")
 		return