expression: add builtin function json_pretty

expression/builtin_json.go

@@ -14,7 +14,8 @@
 package expression
 
 import (
-	json2 "encoding/json"
+	"bytes"
+	goJSON "encoding/json"
 	"strconv"
 	"strings"
 
@@ -841,7 +842,7 @@ func (b *builtinJSONValidStringSig) evalInt(row chunk.Row) (res int64, isNull bo
 	}
 
 	data := hack.Slice(val)
-	if json2.Valid(data) {
+	if goJSON.Valid(data) {
 		res = 1
 	} else {
 		res = 0
@@ -1072,8 +1073,45 @@ type jsonPrettyFunctionClass struct {
 	baseFunctionClass
 }
 
+type builtinJSONSPrettySig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinJSONSPrettySig) Clone() builtinFunc {
+	newSig := &builtinJSONSPrettySig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
 func (c *jsonPrettyFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
-	return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "JSON_PRETTY")
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString, types.ETJson)
+	if err != nil {
+		return nil, err
+	}
+	sig := &builtinJSONSPrettySig{bf}
+	sig.setPbCode(tipb.ScalarFuncSig_JsonPrettySig)
+	return sig, nil
+}
+
+func (b *builtinJSONSPrettySig) evalString(row chunk.Row) (res string, isNull bool, err error) {
+	obj, isNull, err := b.args[0].EvalJSON(b.ctx, row)
+	if isNull || err != nil {
+		return res, isNull, err
+	}
+
+	buf, err := obj.MarshalJSON()
+	if err != nil {
+		return res, isNull, err
+	}
+	var resBuf bytes.Buffer
+	if err = goJSON.Indent(&resBuf, buf, "", "  "); err != nil {
+		return res, isNull, err
+	}
+	return resBuf.String(), false, nil
 }
 
 type jsonQuoteFunctionClass struct {

expression/builtin_json_test.go

@@ -995,3 +995,76 @@ func (s *testEvaluatorSuite) TestJSONStorageSize(c *C) {
 		}
 	}
 }
+
+func (s *testEvaluatorSuite) TestJSONPretty(c *C) {
+	fc := funcs[ast.JSONPretty]
+	tbl := []struct {
+		input    []interface{}
+		expected interface{}
+		success  bool
+	}{
+		// Tests scalar arguments
+		{[]interface{}{nil}, nil, true},
+		{[]interface{}{`true`}, "true", true},
+		{[]interface{}{`false`}, "false", true},
+		{[]interface{}{`2223`}, "2223", true},
+		// Tests simple json
+		{[]interface{}{`{"a":1}`}, `{
+  "a": 1
+}`, true},
+		{[]interface{}{`[1]`}, `[
+  1
+]`, true},
+		// Test complex json
+		{[]interface{}{`{"a":1,"b":[{"d":1},{"e":2},{"f":3}],"c":"eee"}`}, `{
+  "a": 1,
+  "b": [
+    {
+      "d": 1
+    },
+    {
+      "e": 2
+    },
+    {
+      "f": 3
+    }
+  ],
+  "c": "eee"
+}`, true},
+		{[]interface{}{`{"a":1,"b":"qwe","c":[1,2,3,"123",null],"d":{"d1":1,"d2":2}}`}, `{
+  "a": 1,
+  "b": "qwe",
+  "c": [
+    1,
+    2,
+    3,
+    "123",
+    null
+  ],
+  "d": {
+    "d1": 1,
+    "d2": 2
+  }
+}`, true},
+		// Tests invalid json data
+		{[]interface{}{`{1}`}, nil, false},
+		{[]interface{}{`[1,3,4,5]]`}, nil, false},
+	}
+	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)
+
+			if t.expected == nil {
+				c.Assert(d.IsNull(), IsTrue)
+			} else {
+				c.Assert(d.GetString(), Equals, t.expected.(string))
+			}
+		} else {
+			c.Assert(err, NotNil)
+		}
+	}
+}

expression/builtin_json_vec.go

@@ -14,6 +14,8 @@
 package expression
 
 import (
+	"bytes"
+	goJSON "encoding/json"
 	"strconv"
 	"strings"
 
@@ -1176,3 +1178,39 @@ func (b *builtinJSONUnquoteSig) vecEvalString(input *chunk.Chunk, result *chunk.
 	}
 	return nil
 }
+
+func (b *builtinJSONSPrettySig) vectorized() bool {
+	return true
+}
+
+func (b *builtinJSONSPrettySig) vecEvalString(input *chunk.Chunk, result *chunk.Column) error {
+	n := input.NumRows()
+	buf, err := b.bufAllocator.get(types.ETJson, n)
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(buf)
+	if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil {
+		return err
+	}
+	result.ReserveString(n)
+	for i := 0; i < n; i++ {
+		if buf.IsNull(i) {
+			result.AppendNull()
+			continue
+		}
+
+		jb, err := buf.GetJSON(i).MarshalJSON()
+		if err != nil {
+			return err
+		}
+
+		var resBuf bytes.Buffer
+		if err = goJSON.Indent(&resBuf, jb, "", "  "); err != nil {
+			return err
+		}
+
+		result.AppendString(resBuf.String())
+	}
+	return nil
+}

expression/integration_test.go

@@ -213,6 +213,74 @@ func (s *testIntegrationSuite) TestFuncLpadAndRpad(c *C) {
 	result.Check(testkit.Rows("<nil> <nil>"))
 }
 
+func (s *testIntegrationSuite) TestBuiltinFuncJsonPretty(c *C) {
+	ctx := context.Background()
+	tk := testkit.NewTestKit(c, s.store)
+	defer s.cleanEnv(c)
+
+	tk.MustExec(`use test;`)
+	tk.MustExec(`drop table if exists t;`)
+	tk.MustExec("CREATE TABLE t  (`id` int NOT NULL AUTO_INCREMENT, `j` json,vc VARCHAR(500) ,  PRIMARY KEY (`id`));")
+	tk.MustExec(`INSERT INTO t ( id, j, vc ) VALUES
+	( 1, '{"a":1,"b":"qwe","c":[1,2,3,"123",null],"d":{"d1":1,"d2":2}}', '{"a":1,"b":"qwe","c":[1,2,3,"123",null],"d":{"d1":1,"d2":2}}' ),
+	( 2, '[1,2,34]', '{' );`)
+
+	// valid json format in json and varchar
+	checkResult := []string{
+		`{
+  "a": 1,
+  "b": "qwe",
+  "c": [
+    1,
+    2,
+    3,
+    "123",
+    null
+  ],
+  "d": {
+    "d1": 1,
+    "d2": 2
+  }
+}`,
+		`{
+  "a": 1,
+  "b": "qwe",
+  "c": [
+    1,
+    2,
+    3,
+    "123",
+    null
+  ],
+  "d": {
+    "d1": 1,
+    "d2": 2
+  }
+}`,
+	}
+	tk.
+		MustQuery("select JSON_PRETTY(t.j),JSON_PRETTY(vc) from  t where id = 1;").
+		Check(testkit.Rows(strings.Join(checkResult, " ")))
+
+	// invalid json format in varchar
+	rs, _ := tk.Exec("select JSON_PRETTY(t.j),JSON_PRETTY(vc) from  t where id = 2;")
+	_, err := session.GetRows4Test(ctx, tk.Se, rs)
+	terr := errors.Cause(err).(*terror.Error)
+	c.Assert(terr.Code(), Equals, errors.ErrCode(mysql.ErrInvalidJSONText))
+
+	// invalid json format in one row
+	rs, _ = tk.Exec("select JSON_PRETTY(t.j),JSON_PRETTY(vc) from  t where id in (1,2);")
+	_, err = session.GetRows4Test(ctx, tk.Se, rs)
+	terr = errors.Cause(err).(*terror.Error)
+	c.Assert(terr.Code(), Equals, errors.ErrCode(mysql.ErrInvalidJSONText))
+
+	// invalid json string
+	rs, _ = tk.Exec(`select JSON_PRETTY("[1,2,3]}");`)
+	_, err = session.GetRows4Test(ctx, tk.Se, rs)
+	terr = errors.Cause(err).(*terror.Error)
+	c.Assert(terr.Code(), Equals, errors.ErrCode(mysql.ErrInvalidJSONText))
+}
+
 func (s *testIntegrationSuite) TestMiscellaneousBuiltin(c *C) {
 	ctx := context.Background()
 	defer s.cleanEnv(c)