jsonpath: add support for key access wildcards

This commit adds support for wildcard key accessors within jsonpath
queries.

Epic: None
Release note (sql change): Add support for wildcard key accessors in
JSONPath queries. For example, `SELECT jsonb_path_query('{"a": 1, "b": true}', '$.*');`.

Author: normanchenn

pkg/sql/logictest/testdata/logic_test/jsonb_path_query

@@ -925,5 +925,36 @@ SELECT jsonb_path_query('{"a": [1, 2], "b": "hello"}', 'strict $.a ? ($.b == "he
 ----
 [1, 2]
 
+query T rowsort
+SELECT jsonb_path_query('{"a": "world", "b": 2, "c": true}', '$.*');
+----
+"world"
+2
+true
+
+query T rowsort
+SELECT jsonb_path_query('{"a": ["hello", "world"], "b": [2, 5], "c": [true, false], "d": "non-array"}', '$.*[1]');
+----
+"world"
+5
+false
+
+query T
+SELECT jsonb_path_query('{"a": {"ab": 1}, "b": {"bc": 2}, "c": {"cd": 3}}', '$.*.bc');
+----
+2
+
+query empty
+SELECT jsonb_path_query('{}', '$.*');
+
+query empty
+SELECT jsonb_path_query('[1, 2, 3, 4, 5]', '$.*');
+
+query T rowsort
+SELECT jsonb_path_query('{"a": {"x": {"y": 1}}, "b": {"x": {"z": 2}}}', '$.*.x.*');
+----
+1
+2
+
 # select jsonb_path_query('[1, 2, 3, 4, 5]', '$[-1]');
 # select jsonb_path_query('[1, 2, 3, 4, 5]', 'strict $[-1]');

pkg/util/jsonpath/eval/eval.go

@@ -91,6 +91,8 @@ func (ctx *jsonpathCtx) eval(
 		return []json.JSON{jsonValue}, nil
 	case jsonpath.Key:
 		return ctx.evalKey(path, jsonValue, unwrap)
+	case jsonpath.AnyKey:
+		return ctx.evalAnyKey(path, jsonValue, unwrap)
 	case jsonpath.Wildcard:
 		return ctx.evalArrayWildcard(jsonValue)
 	case jsonpath.ArrayList:
@@ -149,20 +151,35 @@ func (ctx *jsonpathCtx) executeAnyItem(
 	}
 	var agg []json.JSON
 	for _, item := range childItems {
-		// The case when this will happen is if jsonValue is an empty array,
-		// in which case we just skip the evaluation.
+		// The case when this will happen is if jsonValue is an empty array or empty
+		// object, in which case we just skip the evaluation.
 		if item.Len() == 0 {
 			continue
 		}
 		if item.Len() != 1 {
 			return nil, errors.AssertionFailedf("unexpected path length")
 		}
-		unwrappedItem, err := item.FetchValIdx(0 /* idx */)
-		if err != nil {
-			return nil, err
-		}
-		if unwrappedItem == nil {
-			return nil, errors.AssertionFailedf("unwrapping json element")
+
+		var unwrappedItem json.JSON
+		switch item.Type() {
+		case json.ArrayJSONType:
+			unwrappedItem, err = item.FetchValIdx(0 /* idx */)
+			if err != nil {
+				return nil, err
+			}
+			if unwrappedItem == nil {
+				return nil, errors.AssertionFailedf("unwrapping json element")
+			}
+		case json.ObjectJSONType:
+			iter, _ := item.ObjectIter()
+			// Guaranteed to have one item.
+			ok := iter.Next()
+			if !ok {
+				return nil, errors.AssertionFailedf("unexpected empty json object")
+			}
+			unwrappedItem = iter.Value()
+		default:
+			panic(errors.AssertionFailedf("unexpected json type"))
 		}
 		if jsonPath == nil {
 			agg = append(agg, unwrappedItem)

pkg/util/jsonpath/eval/key.go

@@ -33,3 +33,17 @@ func (ctx *jsonpathCtx) evalKey(
 	}
 	return []json.JSON{}, nil
 }
+
+func (ctx *jsonpathCtx) evalAnyKey(
+	anyKey jsonpath.AnyKey, jsonValue json.JSON, unwrap bool,
+) ([]json.JSON, error) {
+	if jsonValue.Type() == json.ObjectJSONType {
+		return ctx.executeAnyItem(nil /* jsonPath */, jsonValue, !ctx.strict /* unwrapNext */)
+	} else if unwrap && jsonValue.Type() == json.ArrayJSONType {
+		return ctx.unwrapCurrentTargetAndEval(anyKey, jsonValue, false /* unwrapNext */)
+	} else if ctx.strict {
+		return nil, pgerror.Newf(pgcode.SQLJSONObjectNotFound,
+			"jsonpath wildcard member accessor can only be applied to an object")
+	}
+	return []json.JSON{}, nil
+}

pkg/util/jsonpath/expr.go

@@ -129,3 +129,9 @@ var _ tree.RegexpCacheKey = Regex{}
 func (r Regex) Pattern() (string, error) {
 	return r.Regex, nil
 }
+
+type AnyKey struct{}
+
+var _ Path = AnyKey{}
+
+func (a AnyKey) String() string { return ".*" }

pkg/util/jsonpath/parser/jsonpath.y

@@ -311,6 +311,10 @@ accessor_op:
   {
     $$.val = jsonpath.Filter{Condition: $3.path()}
   }
+| '.' '*'
+  {
+    $$.val = jsonpath.AnyKey{}
+  }
 ;
 
 key:
@@ -430,6 +434,7 @@ comp_op:
   }
 ;
 
+// TODO(normanchenn): support negative numbers.
 scalar_value:
   VARIABLE
   {

pkg/util/jsonpath/parser/testdata/jsonpath

@@ -459,6 +459,17 @@ $.a ? (@.b like_regex "^[aeiou]")
 ----
 $."a"?((@."b" like_regex "^[aeiou]")) -- normalized!
 
+parse
+$.*
+----
+$.*
+
+parse
+$.abc.*.def.*
+----
+$."abc".*."def".* -- normalized!
+
+
 # postgres allows floats as array indexes
 # parse
 # $.abc[1.0]