update json_keys

expression/builtin_json.go

@@ -61,6 +61,8 @@ var (
 	_ builtinFunc = &builtinJSONRemoveSig{}
 	_ builtinFunc = &builtinJSONMergeSig{}
 	_ builtinFunc = &builtinJSONContainsSig{}
+	_ builtinFunc = &builtinJSONKeysSig{}
+	_ builtinFunc = &builtinJSONKeys2ArgsSig{}
 	_ builtinFunc = &builtinJSONLengthSig{}
 )
 
@@ -766,7 +768,87 @@ 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}
+	case 2:
+		sig = &builtinJSONKeys2ArgsSig{bf}
+	}
+	sig.setPbCode(tipb.ScalarFuncSig_JsonKeysSig)
+	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 {

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)
+		}
+	}
+}

expression/integration_test.go

@@ -3317,6 +3317,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('{}'),

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))