From 708611d898f4268d0a549736cf5e54f66544a753 Mon Sep 17 00:00:00 2001 From: Wenlu Wang Date: Thu, 18 Oct 2018 23:54:00 +0800 Subject: [PATCH] expression: add builtin json_keys (#7776) --- expression/builtin_json.go | 85 ++++++++++++++++++++++++++++++++- expression/builtin_json_test.go | 70 +++++++++++++++++++++++++++ expression/integration_test.go | 9 ++++ types/json/binary.go | 10 ++++ 4 files changed, 173 insertions(+), 1 deletion(-) diff --git a/expression/builtin_json.go b/expression/builtin_json.go index 9a1fdb318546b..5ca057e9f99af 100644 --- a/expression/builtin_json.go +++ b/expression/builtin_json.go @@ -61,6 +61,8 @@ var ( _ builtinFunc = &builtinJSONRemoveSig{} _ builtinFunc = &builtinJSONMergeSig{} _ builtinFunc = &builtinJSONContainsSig{} + _ builtinFunc = &builtinJSONKeysSig{} + _ builtinFunc = &builtinJSONKeys2ArgsSig{} _ builtinFunc = &builtinJSONLengthSig{} ) @@ -766,7 +768,88 @@ type jsonKeysFunctionClass struct { } func (c *jsonKeysFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) { - return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "JSON_KEYS") + if err := c.verifyArgs(args); err != nil { + return nil, errors.Trace(err) + } + argTps := []types.EvalType{types.ETJson} + if len(args) == 2 { + argTps = append(argTps, types.ETString) + } + bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETJson, argTps...) + var sig builtinFunc + switch len(args) { + case 1: + sig = &builtinJSONKeysSig{bf} + sig.setPbCode(tipb.ScalarFuncSig_JsonKeysSig) + case 2: + sig = &builtinJSONKeys2ArgsSig{bf} + sig.setPbCode(tipb.ScalarFuncSig_JsonKeys2ArgsSig) + } + return sig, nil +} + +type builtinJSONKeysSig struct { + baseBuiltinFunc +} + +func (b *builtinJSONKeysSig) Clone() builtinFunc { + newSig := &builtinJSONKeysSig{} + newSig.cloneFrom(&b.baseBuiltinFunc) + return newSig +} + +func (b *builtinJSONKeysSig) evalJSON(row chunk.Row) (res json.BinaryJSON, isNull bool, err error) { + res, isNull, err = b.args[0].EvalJSON(b.ctx, row) + if isNull || err != nil { + return res, isNull, errors.Trace(err) + } + if res.TypeCode != json.TypeCodeObject { + return res, true, json.ErrInvalidJSONData + } + return res.GetKeys(), false, nil +} + +type builtinJSONKeys2ArgsSig struct { + baseBuiltinFunc +} + +func (b *builtinJSONKeys2ArgsSig) Clone() builtinFunc { + newSig := &builtinJSONKeys2ArgsSig{} + newSig.cloneFrom(&b.baseBuiltinFunc) + return newSig +} + +func (b *builtinJSONKeys2ArgsSig) evalJSON(row chunk.Row) (res json.BinaryJSON, isNull bool, err error) { + res, isNull, err = b.args[0].EvalJSON(b.ctx, row) + if isNull || err != nil { + return res, isNull, errors.Trace(err) + } + if res.TypeCode != json.TypeCodeObject { + return res, true, json.ErrInvalidJSONData + } + + path, isNull, err := b.args[1].EvalString(b.ctx, row) + if isNull || err != nil { + return res, isNull, errors.Trace(err) + } + + pathExpr, err := json.ParseJSONPathExpr(path) + if err != nil { + return res, true, errors.Trace(err) + } + if pathExpr.ContainsAnyAsterisk() { + return res, true, json.ErrInvalidJSONPathWildcard + } + + res, exists := res.Extract([]json.PathExpression{pathExpr}) + if !exists { + return res, true, nil + } + if res.TypeCode != json.TypeCodeObject { + return res, true, json.ErrInvalidJSONData + } + + return res.GetKeys(), false, nil } type jsonLengthFunctionClass struct { diff --git a/expression/builtin_json_test.go b/expression/builtin_json_test.go index ab4be5790ad9f..6dc1bf1521197 100644 --- a/expression/builtin_json_test.go +++ b/expression/builtin_json_test.go @@ -485,6 +485,7 @@ func (s *testEvaluatorSuite) TestJSONLength(c *C) { d, err := evalBuiltinFunc(f, chunk.Row{}) if t.success { c.Assert(err, IsNil) + if t.expected == nil { c.Assert(d.IsNull(), IsTrue) } else { @@ -495,3 +496,72 @@ func (s *testEvaluatorSuite) TestJSONLength(c *C) { } } } + +func (s *testEvaluatorSuite) TestJSONKeys(c *C) { + defer testleak.AfterTest(c)() + fc := funcs[ast.JSONKeys] + tbl := []struct { + input []interface{} + expected interface{} + success bool + }{ + // Tests nil arguments + {[]interface{}{nil}, nil, true}, + {[]interface{}{nil, "$.c"}, nil, true}, + {[]interface{}{`{"a": 1}`, nil}, nil, true}, + {[]interface{}{nil, nil}, nil, true}, + + // Tests with other type + {[]interface{}{`1`}, nil, false}, + {[]interface{}{`"str"`}, nil, false}, + {[]interface{}{`true`}, nil, false}, + {[]interface{}{`null`}, nil, false}, + {[]interface{}{`[1, 2]`}, nil, false}, + {[]interface{}{`["1", "2"]`}, nil, false}, + + // Tests without path expression + {[]interface{}{`{}`}, `[]`, true}, + {[]interface{}{`{"a": 1}`}, `["a"]`, true}, + {[]interface{}{`{"a": 1, "b": 2}`}, `["a", "b"]`, true}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`}, `["a", "b"]`, true}, + + // Tests with path expression + {[]interface{}{`{"a": 1}`, "$.a"}, nil, false}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`, "$.a"}, `["c"]`, true}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`, "$.a.c"}, nil, false}, + + // Tests path expression contains any asterisk + {[]interface{}{`{}`, "$.*"}, nil, false}, + {[]interface{}{`{"a": 1}`, "$.*"}, nil, false}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`, "$.*"}, nil, false}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`, "$.a.*"}, nil, false}, + + // Tests path expression does not identify a section of the target document + {[]interface{}{`{"a": 1}`, "$.b"}, nil, true}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`, "$.c"}, nil, true}, + {[]interface{}{`{"a": {"c": 3}, "b": 2}`, "$.a.d"}, nil, true}, + } + 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 json.BinaryJSON + j1, err = json.ParseBinaryFromString(x) + c.Assert(err, IsNil) + j2 := d.GetMysqlJSON() + var cmp int + cmp = json.CompareBinary(j1, j2) + c.Assert(cmp, Equals, 0) + case nil: + c.Assert(d.IsNull(), IsTrue) + } + } else { + c.Assert(err, NotNil) + } + } +} diff --git a/expression/integration_test.go b/expression/integration_test.go index 20e1a6cec3fcc..c81203e1518b3 100644 --- a/expression/integration_test.go +++ b/expression/integration_test.go @@ -3367,6 +3367,15 @@ func (s *testIntegrationSuite) TestFuncJSON(c *C) { `) r.Check(testkit.Rows("1 0 1 0")) + r = tk.MustQuery(`select + + json_keys('{}'), + json_keys('{"a": 1, "b": 2}'), + json_keys('{"a": {"c": 3}, "b": 2}'), + json_keys('{"a": {"c": 3}, "b": 2}', "$.a") + `) + r.Check(testkit.Rows(`[] ["a", "b"] ["a", "b"] ["c"]`)) + r = tk.MustQuery(`select json_length('1'), json_length('{}'), diff --git a/types/json/binary.go b/types/json/binary.go index 0d0cd022d0dbb..a495c042cb344 100644 --- a/types/json/binary.go +++ b/types/json/binary.go @@ -169,6 +169,16 @@ func (bj BinaryJSON) GetString() []byte { return bj.Value[lenLen : lenLen+int(strLen)] } +// GetKeys gets the keys of the object +func (bj BinaryJSON) GetKeys() BinaryJSON { + count := bj.GetElemCount() + ret := make([]BinaryJSON, 0, count) + for i := 0; i < count; i++ { + ret = append(ret, CreateBinary(string(bj.objectGetKey(i)))) + } + return buildBinaryArray(ret) +} + // GetElemCount gets the count of Object or Array. func (bj BinaryJSON) GetElemCount() int { return int(endian.Uint32(bj.Value))