Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

expression: add builtin function json_pretty #24675

Merged
merged 8 commits into from
May 20, 2021
44 changes: 41 additions & 3 deletions expression/builtin_json.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@
package expression

import (
json2 "encoding/json"
"bytes"
goJSON "encoding/json"
"strconv"
"strings"

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
73 changes: 73 additions & 0 deletions expression/builtin_json_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}
38 changes: 38 additions & 0 deletions expression/builtin_json_vec.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
package expression

import (
"bytes"
goJSON "encoding/json"
"strconv"
"strings"

Expand Down Expand Up @@ -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
}
68 changes: 68 additions & 0 deletions expression/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down