expression: add builtin json_array_append (#9609)

expression/builtin_json.go

@@ -54,6 +54,7 @@ var (
 	_ builtinFunc = &builtinJSONQuoteSig{}
 	_ builtinFunc = &builtinJSONUnquoteSig{}
 	_ builtinFunc = &builtinJSONArraySig{}
+	_ builtinFunc = &builtinJSONArrayAppendSig{}
 	_ builtinFunc = &builtinJSONObjectSig{}
 	_ builtinFunc = &builtinJSONExtractSig{}
 	_ builtinFunc = &builtinJSONSetSig{}
@@ -702,8 +703,94 @@ type jsonArrayAppendFunctionClass struct {
 	baseFunctionClass
 }
 
+type builtinJSONArrayAppendSig struct {
+	baseBuiltinFunc
+}
+
+func (c *jsonArrayAppendFunctionClass) verifyArgs(args []Expression) error {
+	if len(args) < 3 || (len(args)&1 != 1) {
+		return ErrIncorrectParameterCount.GenWithStackByArgs(c.funcName)
+	}
+	return nil
+}
+
 func (c *jsonArrayAppendFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
-	return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "JSON_ARRAY_APPEND")
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+	argTps := make([]types.EvalType, 0, len(args))
+	argTps = append(argTps, types.ETJson)
+	for i := 1; i < len(args)-1; i += 2 {
+		argTps = append(argTps, types.ETString, types.ETJson)
+	}
+	bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETJson, argTps...)
+	for i := 2; i < len(args); i += 2 {
+		DisableParseJSONFlag4Expr(args[i])
+	}
+	sig := &builtinJSONArrayAppendSig{bf}
+	sig.setPbCode(tipb.ScalarFuncSig_JsonArrayAppendSig)
+	return sig, nil
+}
+
+func (b *builtinJSONArrayAppendSig) Clone() builtinFunc {
+	newSig := &builtinJSONArrayAppendSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinJSONArrayAppendSig) evalJSON(row chunk.Row) (res json.BinaryJSON, isNull bool, err error) {
+	res, isNull, err = b.args[0].EvalJSON(b.ctx, row)
+	if err != nil || isNull {
+		return res, true, err
+	}
+
+	for i := 1; i < len(b.args)-1; i += 2 {
+		// If JSON path is NULL, MySQL breaks and returns NULL.
+		s, isNull, err := b.args[i].EvalString(b.ctx, row)
+		if isNull || err != nil {
+			return res, true, err
+		}
+
+		// We should do the following checks to get correct values in res.Extract
+		pathExpr, err := json.ParseJSONPathExpr(s)
+		if err != nil {
+			return res, true, json.ErrInvalidJSONPath.GenWithStackByArgs(s)
+		}
+		if pathExpr.ContainsAnyAsterisk() {
+			return res, true, json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(s)
+		}
+
+		obj, exists := res.Extract([]json.PathExpression{pathExpr})
+		if !exists {
+			// If path not exists, just do nothing and no errors.
+			continue
+		}
+
+		if obj.TypeCode != json.TypeCodeArray {
+			// res.Extract will return a json object instead of an array if there is an object at path pathExpr.
+			// JSON_ARRAY_APPEND({"a": "b"}, "$", {"b": "c"}) => [{"a": "b"}, {"b", "c"}]
+			// We should wrap them to a single array first.
+			obj = json.CreateBinary([]interface{}{obj})
+		}
+
+		value, isnull, err := b.args[i+1].EvalJSON(b.ctx, row)
+		if err != nil {
+			return res, true, err
+		}
+
+		if isnull {
+			value = json.CreateBinary(nil)
+		}
+
+		obj = json.MergeBinary([]json.BinaryJSON{obj, value})
+		res, err = res.Modify([]json.PathExpression{pathExpr}, []json.BinaryJSON{obj}, json.ModifySet)
+		if err != nil {
+			// We checked pathExpr in the same way as res.Modify do.
+			// So err should always be nil, the function should never return here.
+			return res, true, err
+		}
+	}
+	return res, false, nil
 }
 
 type jsonArrayInsertFunctionClass struct {

expression/builtin_json_test.go

@@ -16,6 +16,7 @@ package expression
 import (
 	. "github.com/pingcap/check"
 	"github.com/pingcap/parser/ast"
+	"github.com/pingcap/parser/terror"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
@@ -686,3 +687,82 @@ func (s *testEvaluatorSuite) TestJSONDepth(c *C) {
 		}
 	}
 }
+
+func (s *testEvaluatorSuite) TestJSONArrayAppend(c *C) {
+	defer testleak.AfterTest(c)()
+	sampleJSON, err := json.ParseBinaryFromString(`{"b": 2}`)
+	c.Assert(err, IsNil)
+	fc := funcs[ast.JSONArrayAppend]
+	tbl := []struct {
+		input    []interface{}
+		expected interface{}
+		err      *terror.Error
+	}{
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$.d`, `z`}, `{"a": 1, "b": [2, 3], "c": 4}`, nil},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$`, `w`}, `[{"a": 1, "b": [2, 3], "c": 4}, "w"]`, nil},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$`, nil}, `[{"a": 1, "b": [2, 3], "c": 4}, null]`, nil},
+		{[]interface{}{`{"a": 1}`, `$`, `{"b": 2}`}, `[{"a": 1}, "{\"b\": 2}"]`, nil},
+		{[]interface{}{`{"a": 1}`, `$`, sampleJSON}, `[{"a": 1}, {"b": 2}]`, nil},
+		{[]interface{}{`{"a": 1}`, `$.a`, sampleJSON}, `{"a": [1, {"b": 2}]}`, nil},
+
+		{[]interface{}{`{"a": 1}`, `$.a`, sampleJSON, `$.a[1]`, sampleJSON}, `{"a": [1, [{"b": 2}, {"b": 2}]]}`, nil},
+		{[]interface{}{nil, `$`, nil}, nil, nil},
+		{[]interface{}{nil, `$`, `a`}, nil, nil},
+		{[]interface{}{`null`, `$`, nil}, `[null, null]`, nil},
+		{[]interface{}{`[]`, `$`, nil}, `[null]`, nil},
+		{[]interface{}{`{}`, `$`, nil}, `[{}, null]`, nil},
+		// Bad arguments.
+		{[]interface{}{`asdf`, `$`, nil}, nil, json.ErrInvalidJSONText},
+		{[]interface{}{``, `$`, nil}, nil, json.ErrInvalidJSONText},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$.d`}, nil, ErrIncorrectParameterCount},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$.c`, `y`, `$.b`}, nil, ErrIncorrectParameterCount},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, nil, nil}, nil, nil},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `asdf`, nil}, nil, json.ErrInvalidJSONPath},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, 42, nil}, nil, json.ErrInvalidJSONPath},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$.*`, nil}, nil, json.ErrInvalidJSONPathWildcard},
+		// Following tests come from MySQL doc.
+		{[]interface{}{`["a", ["b", "c"], "d"]`, `$[1]`, 1}, `["a", ["b", "c", 1], "d"]`, nil},
+		{[]interface{}{`["a", ["b", "c"], "d"]`, `$[0]`, 2}, `[["a", 2], ["b", "c"], "d"]`, nil},
+		{[]interface{}{`["a", ["b", "c"], "d"]`, `$[1][0]`, 3}, `["a", [["b", 3], "c"], "d"]`, nil},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$.b`, `x`}, `{"a": 1, "b": [2, 3, "x"], "c": 4}`, nil},
+		{[]interface{}{`{"a": 1, "b": [2, 3], "c": 4}`, `$.c`, `y`}, `{"a": 1, "b": [2, 3], "c": [4, "y"]}`, nil},
+		// Following tests come from MySQL test.
+		{[]interface{}{`[1,2,3, {"a":[4,5,6]}]`, `$`, 7}, `[1, 2, 3, {"a": [4, 5, 6]}, 7]`, nil},
+		{[]interface{}{`[1,2,3, {"a":[4,5,6]}]`, `$`, 7, `$[3].a`, 3.14}, `[1, 2, 3, {"a": [4, 5, 6, 3.14]}, 7]`, nil},
+		{[]interface{}{`[1,2,3, {"a":[4,5,6]}]`, `$`, 7, `$[3].b`, 8}, `[1, 2, 3, {"a": [4, 5, 6]}, 7]`, nil},
+	}
+
+	for i, t := range tbl {
+		args := types.MakeDatums(t.input...)
+		s.ctx.GetSessionVars().StmtCtx.SetWarnings(nil)
+		f, err := fc.getFunction(s.ctx, s.datumsToConstants(args))
+		// No error should return in getFunction if t.err is nil.
+		if err != nil {
+			c.Assert(t.err, NotNil)
+			c.Assert(t.err.Equal(err), Equals, true)
+			continue
+		}
+
+		c.Assert(f, NotNil)
+		d, err := evalBuiltinFunc(f, chunk.Row{})
+		comment := Commentf("case:%v \n input:%v \n output: %s \n expected: %v \n warnings: %v \n expected error %v", i, t.input, d.GetMysqlJSON(), t.expected, s.ctx.GetSessionVars().StmtCtx.GetWarnings(), t.err)
+
+		if t.err != nil {
+			c.Assert(t.err.Equal(err), Equals, true, comment)
+			continue
+		}
+
+		c.Assert(err, IsNil, comment)
+		c.Assert(int(s.ctx.GetSessionVars().StmtCtx.WarningCount()), Equals, 0, comment)
+
+		if t.expected == nil {
+			c.Assert(d.IsNull(), IsTrue, comment)
+			continue
+		}
+
+		j1, err := json.ParseBinaryFromString(t.expected.(string))
+
+		c.Assert(err, IsNil, comment)
+		c.Assert(json.CompareBinary(j1, d.GetMysqlJSON()), Equals, 0, comment)
+	}
+}

expression/distsql_builtin.go

@@ -420,6 +420,8 @@ func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *ti
 		f = &builtinJSONUnquoteSig{base}
 	case tipb.ScalarFuncSig_JsonArraySig:
 		f = &builtinJSONArraySig{base}
+	case tipb.ScalarFuncSig_JsonArrayAppendSig:
+		f = &builtinJSONArrayAppendSig{base}
 	case tipb.ScalarFuncSig_JsonObjectSig:
 		f = &builtinJSONObjectSig{base}
 	case tipb.ScalarFuncSig_JsonExtractSig: