Skip to content

Commit fdcf275

Browse files
vitess-bot[bot]Jun Wang
vitess-bot[bot]
authored and
Jun Wang
committed
[release-23.0] evalengine: make JSON_EXTRACT work with non-static arguments (vitessio#19035) (vitessio#19254)
Signed-off-by: Arthur Schreiber <arthur@planetscale.com> Co-authored-by: vitess-bot[bot] <108069721+vitess-bot[bot]@users.noreply.github.com> Signed-off-by: Jun Wang <jun.wang@demonware.net>
1 parent d447af9 commit fdcf275

4 files changed

Lines changed: 164 additions & 34 deletions

File tree

‎go/vt/vtgate/evalengine/compiler_asm.go‎

Lines changed: 65 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ import (
4747
"vitess.io/vitess/go/mysql/hex"
4848
"vitess.io/vitess/go/mysql/icuregex"
4949
"vitess.io/vitess/go/mysql/json"
50-
"vitess.io/vitess/go/slice"
5150
"vitess.io/vitess/go/sqltypes"
5251
querypb "vitess.io/vitess/go/vt/proto/query"
5352
"vitess.io/vitess/go/vt/proto/vtrpc"
@@ -2215,40 +2214,89 @@ func (asm *assembler) Fn_JSON_CONTAINS_PATH(match jsonMatch, paths []*json.Path)
22152214
}
22162215
}
22172216

2218-
func (asm *assembler) Fn_JSON_EXTRACT0(jp []*json.Path) {
2219-
multi := len(jp) > 1 || slice.Any(jp, func(path *json.Path) bool { return path.ContainsWildcards() })
2217+
type staticPath struct {
2218+
p *json.Path
2219+
err error
2220+
}
22202221

2221-
if multi {
2222-
asm.emit(func(env *ExpressionEnv) int {
2222+
func (asm *assembler) Fn_JSON_EXTRACT(args int, staticPaths []staticPath) {
2223+
asm.adjustStack(-args)
2224+
asm.emit(func(env *ExpressionEnv) int {
2225+
paths := make([]*json.Path, 0, args)
2226+
2227+
multi := args > 1
2228+
2229+
doct := env.vm.stack[env.vm.sp-(args+1)].(*evalJSON)
2230+
2231+
for i := args; i > 0; i-- {
2232+
arg := env.vm.stack[env.vm.sp-i]
2233+
2234+
if arg == nil {
2235+
env.vm.sp -= args
2236+
env.vm.stack[env.vm.sp-1] = nil
2237+
return 1
2238+
}
2239+
2240+
staticPath := staticPaths[args-i]
2241+
if staticPath.err != nil {
2242+
env.vm.err = staticPath.err
2243+
return 1
2244+
}
2245+
2246+
var path *json.Path
2247+
if staticPath.p != nil {
2248+
path = staticPath.p
2249+
} else {
2250+
pathBytes, err := evalToVarchar(arg, collations.CollationUtf8mb4ID, true)
2251+
if err != nil {
2252+
env.vm.err = err
2253+
return 1
2254+
}
2255+
2256+
path, err = intoJSONPath(pathBytes)
2257+
if err != nil {
2258+
env.vm.err = err
2259+
return 1
2260+
}
2261+
}
2262+
2263+
if !multi {
2264+
multi = path.ContainsWildcards()
2265+
}
2266+
2267+
paths = append(paths, path)
2268+
}
2269+
2270+
env.vm.sp -= args
2271+
2272+
if multi {
22232273
matches := make([]*json.Value, 0, 4)
2224-
arg := env.vm.stack[env.vm.sp-1].(*evalJSON)
2225-
for _, jp := range jp {
2226-
jp.Match(arg, true, func(value *json.Value) {
2274+
for _, jp := range paths {
2275+
jp.Match(doct, true, func(value *json.Value) {
22272276
matches = append(matches, value)
22282277
})
22292278
}
2279+
22302280
if len(matches) == 0 {
22312281
env.vm.stack[env.vm.sp-1] = nil
22322282
} else {
22332283
env.vm.stack[env.vm.sp-1] = json.NewArray(matches)
22342284
}
2235-
return 1
2236-
}, "FN JSON_EXTRACT, SP-1, [static]")
2237-
} else {
2238-
asm.emit(func(env *ExpressionEnv) int {
2285+
} else {
22392286
var match *json.Value
2240-
arg := env.vm.stack[env.vm.sp-1].(*evalJSON)
2241-
jp[0].Match(arg, true, func(value *json.Value) {
2287+
paths[0].Match(doct, true, func(value *json.Value) {
22422288
match = value
22432289
})
2290+
22442291
if match == nil {
22452292
env.vm.stack[env.vm.sp-1] = nil
22462293
} else {
22472294
env.vm.stack[env.vm.sp-1] = match
22482295
}
2249-
return 1
2250-
}, "FN JSON_EXTRACT, SP-1, [static]")
2251-
}
2296+
}
2297+
2298+
return 1
2299+
}, "FN JSON_EXTRACT, SP-1, ..., SP-N")
22522300
}
22532301

22542302
func (asm *assembler) Fn_JSON_KEYS(jp *json.Path) {

‎go/vt/vtgate/evalengine/fn_json.go‎

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ func (call *builtinJSONExtract) eval(env *ExpressionEnv) (eval, error) {
8989
}
9090

9191
func builtin_JSON_EXTRACT(doc *json.Value, paths []eval) (eval, error) {
92-
matches := make([]*json.Value, 0, 4)
92+
matches := make([]*json.Value, 0, len(paths))
9393
multi := len(paths) > 1
9494

9595
for _, p := range paths {
@@ -125,27 +125,57 @@ func (call *builtinJSONExtract) compile(c *compiler) (ctype, error) {
125125
return ctype{}, err
126126
}
127127

128-
if slice.All(call.Arguments[1:], func(expr IR) bool { return expr.constant() }) {
129-
paths := make([]*json.Path, 0, len(call.Arguments[1:]))
128+
// Handle `NULL` arguments
129+
nullable := doct.nullable()
130+
skip := c.compileNullCheck1(doct)
130131

131-
for _, arg := range call.Arguments[1:] {
132-
jp, err := c.jsonExtractPath(arg)
133-
if err != nil {
134-
return ctype{}, err
135-
}
136-
paths = append(paths, jp)
132+
// TODO: `*compiler.compileParseJSON` should handle `sqltypes.Null`` properly but
133+
// we'll handle it here until all call sites are fixed.
134+
var jt ctype
135+
if doct.Type != sqltypes.Null {
136+
jt, err = c.compileParseJSON("JSON_EXTRACT", doct, 1)
137+
if err != nil {
138+
return ctype{}, err
137139
}
140+
} else {
141+
jt = ctype{Type: sqltypes.Null, Flag: flagNull | flagNullable, Col: collationNull}
142+
}
143+
144+
staticPaths := make([]staticPath, 0, len(call.Arguments[1:]))
138145

139-
jt, err := c.compileParseJSON("JSON_EXTRACT", doct, 1)
146+
for _, arg := range call.Arguments[1:] {
147+
argType, err := arg.compile(c)
140148
if err != nil {
141149
return ctype{}, err
142150
}
143151

144-
c.asm.Fn_JSON_EXTRACT0(paths)
145-
return jt, nil
152+
if !nullable {
153+
nullable = argType.nullable()
154+
}
155+
156+
if arg.constant() {
157+
staticEnv := EmptyExpressionEnv(c.env)
158+
arg, err = simplifyExpr(staticEnv, arg)
159+
if err != nil {
160+
return ctype{}, err
161+
}
162+
163+
p, err := c.jsonExtractPath(arg)
164+
staticPaths = append(staticPaths, staticPath{p, err})
165+
} else {
166+
staticPaths = append(staticPaths, staticPath{nil, nil})
167+
}
146168
}
147169

148-
return ctype{}, c.unsupported(call)
170+
c.asm.Fn_JSON_EXTRACT(len(call.Arguments[1:]), staticPaths)
171+
c.asm.jumpDestination(skip)
172+
173+
if nullable {
174+
// If any argument is nullable, the result is nullable too
175+
jt.Flag |= flagNullable
176+
}
177+
178+
return jt, nil
149179
}
150180

151181
func (call *builtinJSONUnquote) eval(env *ExpressionEnv) (eval, error) {

‎go/vt/vtgate/evalengine/integration/fuzz_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ var (
102102
regexp.MustCompile(`Illegal argument to a regular expression`),
103103
regexp.MustCompile(`Incorrect arguments to regexp_substr`),
104104
regexp.MustCompile(`Incorrect arguments to regexp_replace`),
105+
regexp.MustCompile(`Invalid JSON path expression\. The error is around character position (\d+)\.`),
105106
}
106107
)
107108

‎go/vt/vtgate/evalengine/testcases/cases.go‎

Lines changed: 55 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ import (
3030

3131
var Cases = []TestCase{
3232
{Run: JSONExtract, Schema: JSONExtract_Schema},
33-
{Run: JSONPathOperations},
33+
{Run: FnJSONKeys},
34+
{Run: FnJSONExtract},
35+
{Run: FnJSONContainsPath},
3436
{Run: JSONArray},
3537
{Run: JSONObject},
3638
{Run: CharsetConversionOperators},
@@ -176,18 +178,67 @@ var Cases = []TestCase{
176178
{Run: RegexpReplace},
177179
}
178180

179-
func JSONPathOperations(yield Query) {
181+
func FnJSONKeys(yield Query) {
180182
for _, obj := range inputJSONObjects {
181183
yield(fmt.Sprintf("JSON_KEYS('%s')", obj), nil, false)
182184

185+
for _, path1 := range inputJSONPaths {
186+
yield(fmt.Sprintf("JSON_KEYS('%s', '%s')", obj, path1), nil, false)
187+
}
188+
}
189+
}
190+
191+
func FnJSONExtract(yield Query) {
192+
for _, obj := range inputJSONObjects {
183193
for _, path1 := range inputJSONPaths {
184194
yield(fmt.Sprintf("JSON_EXTRACT('%s', '%s')", obj, path1), nil, false)
195+
196+
for _, path2 := range inputJSONPaths {
197+
yield(fmt.Sprintf("JSON_EXTRACT('%s', '%s', '%s')", obj, path1, path2), nil, false)
198+
}
199+
}
200+
}
201+
202+
yield(`JSON_EXTRACT('{"a": 1}', '$.a')`, nil, false)
203+
yield(`JSON_EXTRACT('{"a": 1}', '$.*')`, nil, false)
204+
yield(`JSON_EXTRACT('[1, 2, 3]', '$[0 to 2]')`, nil, false)
205+
yield(`JSON_EXTRACT('{"a": 1, "b": 2}', '$.a', '$.b')`, nil, false)
206+
yield(`JSON_EXTRACT('{"a": 1}', '$.a', '$.z')`, nil, false)
207+
208+
yield(`JSON_EXTRACT(CONCAT('{', '"a"', ':', ' ', '1', '}'), '$.a')`, nil, false)
209+
yield(`JSON_EXTRACT('{"a": 1}', CONCAT('$', '.', 'a'))`, nil, false)
210+
211+
yield(`JSON_EXTRACT(NULL, '$.a')`, nil, false)
212+
yield(`JSON_EXTRACT(NULL, NULL)`, nil, false)
213+
214+
yield(`JSON_EXTRACT('{"a": 1}', NULL)`, nil, false)
215+
yield(`JSON_EXTRACT('{"a": 1}', '$.a', NULL)`, nil, false)
216+
yield(`JSON_EXTRACT('{"a": 1}', NULL, '$.a')`, nil, false)
217+
218+
yield(`JSON_EXTRACT('{"a": 1}', '$.b')`, nil, false)
219+
yield(`JSON_EXTRACT('{"a": 1}', '$.b', '$.c')`, nil, false)
220+
yield(`JSON_EXTRACT('[1,2,3]', '$[10]')`, nil, false)
221+
222+
yield(`JSON_EXTRACT('{invalid}', '$.a')`, nil, false)
223+
yield(`JSON_EXTRACT('not json', '$.a')`, nil, false)
224+
yield(`JSON_EXTRACT('', '$.a')`, nil, false)
225+
yield(`JSON_EXTRACT('{invalid}', NULL)`, nil, false)
226+
227+
yield(`JSON_EXTRACT('{"a": 1}', '$.b[ 1 ].')`, nil, false)
228+
229+
yield(`JSON_EXTRACT(NULL, 'invalid-path')`, nil, false)
230+
yield(`JSON_EXTRACT('{"a": 1}', NULL, 'invalid-path')`, nil, false)
231+
yield(`JSON_EXTRACT('{"a": 1}', 'invalid-path', NULL)`, nil, false)
232+
yield(`JSON_EXTRACT('{"a": 1}', '$.a', 'invalid')`, nil, false)
233+
}
234+
235+
func FnJSONContainsPath(yield Query) {
236+
for _, obj := range inputJSONObjects {
237+
for _, path1 := range inputJSONPaths {
185238
yield(fmt.Sprintf("JSON_CONTAINS_PATH('%s', 'one', '%s')", obj, path1), nil, false)
186239
yield(fmt.Sprintf("JSON_CONTAINS_PATH('%s', 'all', '%s')", obj, path1), nil, false)
187-
yield(fmt.Sprintf("JSON_KEYS('%s', '%s')", obj, path1), nil, false)
188240

189241
for _, path2 := range inputJSONPaths {
190-
yield(fmt.Sprintf("JSON_EXTRACT('%s', '%s', '%s')", obj, path1, path2), nil, false)
191242
yield(fmt.Sprintf("JSON_CONTAINS_PATH('%s', 'one', '%s', '%s')", obj, path1, path2), nil, false)
192243
yield(fmt.Sprintf("JSON_CONTAINS_PATH('%s', 'all', '%s', '%s')", obj, path1, path2), nil, false)
193244
}

0 commit comments

Comments
 (0)