expression/json: add builtin function JSON_SEARCH

expression/builtin_json.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/stringutil"
 	"github.com/pingcap/tipb/go-tipb"
 )
 
@@ -66,6 +67,7 @@ var (
 	_ builtinFunc = &builtinJSONMergeSig{}
 	_ builtinFunc = &builtinJSONContainsSig{}
 	_ builtinFunc = &builtinJSONDepthSig{}
+	_ builtinFunc = &builtinJSONSearchSig{}
 	_ builtinFunc = &builtinJSONKeysSig{}
 	_ builtinFunc = &builtinJSONKeys2ArgsSig{}
 	_ builtinFunc = &builtinJSONLengthSig{}
@@ -876,8 +878,121 @@ type jsonSearchFunctionClass struct {
 	baseFunctionClass
 }
 
+type builtinJSONSearchSig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinJSONSearchSig) Clone() builtinFunc {
+	newSig := &builtinJSONSearchSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
 func (c *jsonSearchFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
-	return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "JSON_SEARCH")
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+	// json_doc, one_or_all, search_str[, escape_char[, path] ...])
+	argTps := make([]types.EvalType, 0, len(args))
+	argTps = append(argTps, types.ETJson)
+	for range args[1:] {
+		argTps = append(argTps, types.ETString)
+	}
+	bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETJson, argTps...)
+	sig := &builtinJSONSearchSig{bf}
+	sig.setPbCode(tipb.ScalarFuncSig_JsonSearchSig)
+	return sig, nil
+}
+
+func (b *builtinJSONSearchSig) evalJSON(row chunk.Row) (res json.BinaryJSON, isNull bool, err error) {
+	// json_doc
+	var obj json.BinaryJSON
+	obj, isNull, err = b.args[0].EvalJSON(b.ctx, row)
+	if isNull || err != nil {
+		return res, isNull, err
+	}
+
+	// one_or_all
+	var containType string
+	containType, isNull, err = b.args[1].EvalString(b.ctx, row)
+	if isNull || err != nil {
+		return res, isNull, err
+	}
+	if containType != json.ContainsPathAll && containType != json.ContainsPathOne {
+		return res, true, errors.AddStack(json.ErrInvalidJSONContainsPathType)
+	}
+
+	// search_str & escape_char
+	var searchStr string
+	searchStr, isNull, err = b.args[2].EvalString(b.ctx, row)
+	if isNull || err != nil {
+		return res, isNull, err
+	}
+	escape := byte('\\')
+	if len(b.args) >= 4 {
+		var escapeStr string
+		escapeStr, isNull, err = b.args[3].EvalString(b.ctx, row)
+		if err != nil {
+			return res, isNull, err
+		}
+		if isNull || len(escapeStr) == 0 {
+			escape = byte('\\')
+		} else if len(escapeStr) == 1 {
+			escape = byte(escapeStr[0])
+		} else {
+			return res, true, errIncorrectArgs.GenWithStackByArgs("ESCAPE")
+		}
+	}
+	patChars, patTypes := stringutil.CompilePattern(searchStr, escape)
+
+	// result
+	result := make([]interface{}, 0)
+
+	// walk json_doc
+	walkFn := func(fullpath json.PathExpression, bj json.BinaryJSON) (stop bool, err error) {
+		if bj.TypeCode == json.TypeCodeString && stringutil.DoMatch(string(bj.GetString()), patChars, patTypes) {
+			result = append(result, fullpath.String())
+			if containType == json.ContainsPathOne {
+				return true, nil
+			}
+		}
+		return false, nil
+	}
+	if len(b.args) >= 5 { // path...
+		pathExprs := make([]json.PathExpression, 0, len(b.args)-4)
+		for i := 4; i < len(b.args); i++ {
+			var s string
+			s, isNull, err = b.args[i].EvalString(b.ctx, row)
+			if isNull || err != nil {
+				return res, isNull, err
+			}
+			var pathExpr json.PathExpression
+			pathExpr, err = json.ParseJSONPathExpr(s)
+			if err != nil {
+				return res, true, err
+			}
+			pathExprs = append(pathExprs, pathExpr)
+		}
+		err = obj.Walk(walkFn, pathExprs...)
+		if err != nil {
+			return res, true, err
+		}
+	} else {
+		err = obj.Walk(walkFn)
+		if err != nil {
+			return res, true, err
+		}
+	}
+
+	// return
+	switch len(result) {
+	case 0:
+		return res, true, nil
+	case 1:
+		return json.CreateBinary(result[0]), false, nil
+	default:
+		return json.CreateBinary(result), false, nil
+	}
 }
 
 type jsonStorageSizeFunctionClass struct {

expression/builtin_json_test.go

@@ -773,3 +773,77 @@ func (s *testEvaluatorSuite) TestJSONArrayAppend(c *C) {
 		c.Assert(json.CompareBinary(j1, d.GetMysqlJSON()), Equals, 0, comment)
 	}
 }
+
+func (s *testEvaluatorSuite) TestJSONSearch(c *C) {
+	defer testleak.AfterTest(c)()
+	fc := funcs[ast.JSONSearch]
+	jsonString := `["abc", [{"k": "10"}, "def"], {"x":"abc"}, {"y":"bcd"}]`
+	jsonString2 := `["abc", [{"k": "10"}, "def"], {"x":"ab%d"}, {"y":"abcd"}]`
+	tbl := []struct {
+		input    []interface{}
+		expected interface{}
+		success  bool
+	}{
+		// simple case
+		{[]interface{}{jsonString, `one`, `abc`}, `"$[0]"`, true},
+		{[]interface{}{jsonString, `all`, `abc`}, `["$[0]", "$[2].x"]`, true},
+		{[]interface{}{jsonString, `all`, `ghi`}, nil, true},
+		{[]interface{}{jsonString, `all`, `10`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `10`, nil, `$`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `10`, nil, `$[*]`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `10`, nil, `$**.k`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `10`, nil, `$[*][0].k`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `10`, nil, `$[1]`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `10`, nil, `$[1][0]`}, `"$[1][0].k"`, true},
+		{[]interface{}{jsonString, `all`, `abc`, nil, `$[2]`}, `"$[2].x"`, true},
+		{[]interface{}{jsonString, `all`, `abc`, nil, `$[2]`, `$[0]`}, `["$[2].x", "$[0]"]`, true},
+		{[]interface{}{jsonString, `all`, `abc`, nil, `$[2]`, `$[2]`}, `"$[2].x"`, true},
+
+		// search pattern
+		{[]interface{}{jsonString, `all`, `%a%`}, `["$[0]", "$[2].x"]`, true},
+		{[]interface{}{jsonString, `all`, `%b%`}, `["$[0]", "$[2].x", "$[3].y"]`, true},
+		{[]interface{}{jsonString, `all`, `%b%`, nil, `$[0]`}, `"$[0]"`, true},
+		{[]interface{}{jsonString, `all`, `%b%`, nil, `$[2]`}, `"$[2].x"`, true},
+		{[]interface{}{jsonString, `all`, `%b%`, nil, `$[1]`}, nil, true},
+		{[]interface{}{jsonString, `all`, `%b%`, ``, `$[1]`}, nil, true},
+		{[]interface{}{jsonString, `all`, `%b%`, nil, `$[3]`}, `"$[3].y"`, true},
+		{[]interface{}{jsonString2, `all`, `ab_d`}, `["$[2].x", "$[3].y"]`, true},
+
+		// escape char
+		{[]interface{}{jsonString2, `all`, `ab%d`}, `["$[2].x", "$[3].y"]`, true},
+		{[]interface{}{jsonString2, `all`, `ab\%d`}, `"$[2].x"`, true},
+		{[]interface{}{jsonString2, `all`, `ab|%d`, `|`}, `"$[2].x"`, true},
+
+		// error handle
+		{[]interface{}{nil, `all`, `abc`}, nil, true},                     // NULL json
+		{[]interface{}{`a`, `all`, `abc`}, nil, false},                    // non json
+		{[]interface{}{jsonString, `wrong`, `abc`}, nil, false},           // wrong one_or_all
+		{[]interface{}{jsonString, `all`, nil}, nil, true},                // NULL search_str
+		{[]interface{}{jsonString, `all`, `abc`, `??`}, nil, false},       // wrong escape_char
+		{[]interface{}{jsonString, `all`, `abc`, nil, nil}, nil, true},    // NULL path
+		{[]interface{}{jsonString, `all`, `abc`, nil, `$xx`}, nil, false}, // wrong path
+	}
+	for _, t := range tbl {
+		args := types.MakeDatums(t.input...)
+		f, err := fc.getFunction(s.ctx, s.datumsToConstants(args))
+		c.Assert(err, IsNil)
+		d, err := evalBuiltinFunc(f, chunk.Row{})
+		if t.success {
+			c.Assert(err, IsNil)
+			switch x := t.expected.(type) {
+			case string:
+				var j1, j2 json.BinaryJSON
+				j1, err = json.ParseBinaryFromString(x)
+				c.Assert(err, IsNil)
+				j2 = d.GetMysqlJSON()
+				cmp := json.CompareBinary(j1, j2)
+				//fmt.Println(j1, j2)
+				c.Assert(cmp, Equals, 0)
+			case nil:
+				c.Assert(d.IsNull(), IsTrue)
+			}
+		} else {
+			c.Assert(err, NotNil)
+		}
+	}
+}

expression/distsql_builtin.go

@@ -444,6 +444,8 @@ func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *ti
 		f = &builtinJSONLengthSig{base}
 	case tipb.ScalarFuncSig_JsonDepthSig:
 		f = &builtinJSONDepthSig{base}
+	case tipb.ScalarFuncSig_JsonSearchSig:
+		f = &builtinJSONSearchSig{base}
 
 	case tipb.ScalarFuncSig_InInt:
 		f = &builtinInIntSig{base}

expression/distsql_builtin_test.go

@@ -17,6 +17,7 @@ import (
 	"time"
 
 	. "github.com/pingcap/check"
+	"github.com/pingcap/parser/charset"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/types"
@@ -108,6 +109,17 @@ func (s *testEvalSuite) TestEval(c *C) {
 			),
 			types.NewIntDatum(3),
 		},
+		{
+			scalarFunctionExpr(tipb.ScalarFuncSig_JsonSearchSig,
+				toPBFieldType(newJSONFieldType()),
+				jsonDatumExpr(c, `["abc", [{"k": "10"}, "def"], {"x":"abc"}, {"y":"bcd"}]`),
+				datumExpr(c, types.NewBytesDatum([]byte(`all`))),
+				datumExpr(c, types.NewBytesDatum([]byte(`10`))),
+				datumExpr(c, types.NewBytesDatum([]byte(`\`))),
+				datumExpr(c, types.NewBytesDatum([]byte(`$**.k`))),
+			),
+			newJSONDatum(c, `"$[1][0].k"`),
+		},
 	}
 	sc := new(stmtctx.StatementContext)
 	for _, tt := range tests {
@@ -178,12 +190,15 @@ func datumExpr(c *C, d types.Datum) *tipb.Expr {
 	return expr
 }
 
-func jsonDatumExpr(c *C, s string) *tipb.Expr {
-	var d types.Datum
+func newJSONDatum(c *C, s string) (d types.Datum) {
 	j, err := json.ParseBinaryFromString(s)
 	c.Assert(err, IsNil)
 	d.SetMysqlJSON(j)
-	return datumExpr(c, d)
+	return d
+}
+
+func jsonDatumExpr(c *C, s string) *tipb.Expr {
+	return datumExpr(c, newJSONDatum(c, s))
 }
 
 func columnExpr(columnID int64) *tipb.Expr {
@@ -214,6 +229,16 @@ func newIntFieldType() *types.FieldType {
 	}
 }
 
+func newJSONFieldType() *types.FieldType {
+	return &types.FieldType{
+		Tp:      mysql.TypeJSON,
+		Flen:    types.UnspecifiedLength,
+		Decimal: 0,
+		Charset: charset.CharsetBin,
+		Collate: charset.CollationBin,
+	}
+}
+
 func scalarFunctionExpr(sigCode tipb.ScalarFuncSig, retType *tipb.FieldType, args ...*tipb.Expr) *tipb.Expr {
 	return &tipb.Expr{
 		Tp:        tipb.ExprType_ScalarFunc,