From c5b2005706d378c60c8108fc498d59f088e7ca72 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Wed, 5 Apr 2017 22:53:12 +0300
Subject: [PATCH v1 3/8] Add jsonpath sequence constructors

---
 doc/src/sgml/func.sgml                       |  7 +++
 src/backend/utils/adt/jsonpath.c             | 68 ++++++++++++++++++++++++++--
 src/backend/utils/adt/jsonpath_exec.c        | 46 +++++++++++++++++++
 src/backend/utils/adt/jsonpath_gram.y        | 32 +++++++++++--
 src/include/utils/jsonpath.h                 | 12 +++++
 src/test/regress/expected/jsonb_jsonpath.out | 54 ++++++++++++++++++++++
 src/test/regress/expected/jsonpath.out       | 36 +++++++++++++++
 src/test/regress/sql/jsonb_jsonpath.sql      |  9 ++++
 src/test/regress/sql/jsonpath.sql            |  7 +++
 9 files changed, 264 insertions(+), 7 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index d344b95..9defd1a 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -12625,6 +12625,13 @@ table2-mapping
        </row>
       </thead>
       <tbody>
+       <row>
+        <entry>Sequence constructor</entry>
+        <entry>Construct a JSON sequence from a list of comma-separated expressions</entry>
+        <entry><literal>[1, 2, 3]</literal></entry>
+        <entry><literal>pg $[*], 4, 5</literal></entry>
+        <entry><literal>1, 2, 3, 4, 5</literal></entry>
+       </row>
       </tbody>
      </tgroup>
     </table>
diff --git a/src/backend/utils/adt/jsonpath.c b/src/backend/utils/adt/jsonpath.c
index 17c09a7..0fb4257 100644
--- a/src/backend/utils/adt/jsonpath.c
+++ b/src/backend/utils/adt/jsonpath.c
@@ -233,7 +233,7 @@ jsonPathToCstring(StringInfo out, JsonPath *in, int estimated_len)
 		appendBinaryStringInfo(out, "strict ", 7);
 
 	jspInit(&v, in);
-	printJsonPathItem(out, &v, false, true);
+	printJsonPathItem(out, &v, false, v.type != jpiSequence);
 
 	return out->data;
 }
@@ -448,6 +448,31 @@ flattenJsonPathParseItem(JsonPathEncodingContext *cxt, JsonPathParseItem *item,
 		case jpiDouble:
 		case jpiKeyValue:
 			break;
+		case jpiSequence:
+			{
+				int32		nelems = list_length(item->value.sequence.elems);
+				ListCell   *lc;
+				int			offset;
+
+				checkJsonPathExtensionsEnabled(cxt, item->type);
+
+				appendBinaryStringInfo(buf, (char *) &nelems, sizeof(nelems));
+
+				offset = buf->len;
+
+				appendStringInfoSpaces(buf, sizeof(int32) * nelems);
+
+				foreach(lc, item->value.sequence.elems)
+				{
+					int32		elempos =
+						flattenJsonPathParseItem(cxt, lfirst(lc), nestingLevel,
+												 insideArraySubscript);
+
+					*(int32 *) &buf->data[offset] = elempos - pos;
+					offset += sizeof(int32);
+				}
+			}
+			break;
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", item->type);
 	}
@@ -670,12 +695,12 @@ printJsonPathItem(StringInfo buf, JsonPathItem *v, bool inKey,
 				if (i)
 					appendStringInfoChar(buf, ',');
 
-				printJsonPathItem(buf, &from, false, false);
+				printJsonPathItem(buf, &from, false, from.type == jpiSequence);
 
 				if (range)
 				{
 					appendBinaryStringInfo(buf, " to ", 4);
-					printJsonPathItem(buf, &to, false, false);
+					printJsonPathItem(buf, &to, false, to.type == jpiSequence);
 				}
 			}
 			appendStringInfoChar(buf, ']');
@@ -736,6 +761,25 @@ printJsonPathItem(StringInfo buf, JsonPathItem *v, bool inKey,
 		case jpiKeyValue:
 			appendBinaryStringInfo(buf, ".keyvalue()", 11);
 			break;
+		case jpiSequence:
+			if (printBracketes || jspHasNext(v))
+				appendStringInfoChar(buf, '(');
+
+			for (i = 0; i < v->content.sequence.nelems; i++)
+			{
+				JsonPathItem elem;
+
+				if (i)
+					appendBinaryStringInfo(buf, ", ", 2);
+
+				jspGetSequenceElement(v, i, &elem);
+
+				printJsonPathItem(buf, &elem, false, elem.type == jpiSequence);
+			}
+
+			if (printBracketes || jspHasNext(v))
+				appendStringInfoChar(buf, ')');
+			break;
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", v->type);
 	}
@@ -808,6 +852,8 @@ operationPriority(JsonPathItemType op)
 {
 	switch (op)
 	{
+		case jpiSequence:
+			return -1;
 		case jpiOr:
 			return 0;
 		case jpiAnd:
@@ -944,6 +990,11 @@ jspInitByBuffer(JsonPathItem *v, char *base, int32 pos)
 			read_int32(v->content.anybounds.first, base, pos);
 			read_int32(v->content.anybounds.last, base, pos);
 			break;
+		case jpiSequence:
+			read_int32(v->content.sequence.nelems, base, pos);
+			read_int32_n(v->content.sequence.elems, base, pos,
+						 v->content.sequence.nelems);
+			break;
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", v->type);
 	}
@@ -1008,7 +1059,8 @@ jspGetNext(JsonPathItem *v, JsonPathItem *a)
 			   v->type == jpiDouble ||
 			   v->type == jpiDatetime ||
 			   v->type == jpiKeyValue ||
-			   v->type == jpiStartsWith);
+			   v->type == jpiStartsWith ||
+			   v->type == jpiSequence);
 
 		if (a)
 			jspInitByBuffer(a, v->base, v->nextPos);
@@ -1103,3 +1155,11 @@ jspGetArraySubscript(JsonPathItem *v, JsonPathItem *from, JsonPathItem *to,
 
 	return true;
 }
+
+void
+jspGetSequenceElement(JsonPathItem *v, int i, JsonPathItem *elem)
+{
+	Assert(v->type == jpiSequence);
+
+	jspInitByBuffer(elem, v->base, v->content.sequence.elems[i]);
+}
diff --git a/src/backend/utils/adt/jsonpath_exec.c b/src/backend/utils/adt/jsonpath_exec.c
index 1b05307..1e31d3c 100644
--- a/src/backend/utils/adt/jsonpath_exec.c
+++ b/src/backend/utils/adt/jsonpath_exec.c
@@ -1129,6 +1129,52 @@ executeItemOptUnwrapTarget(JsonPathExecContext *cxt, JsonPathItem *jsp,
 
 			return executeKeyValueMethod(cxt, jsp, jb, found);
 
+		case jpiSequence:
+			{
+				JsonPathItem next;
+				bool		hasNext = jspGetNext(jsp, &next);
+				JsonValueList list;
+				JsonValueList *plist = hasNext ? &list : found;
+				JsonValueListIterator it;
+				int			i;
+
+				for (i = 0; i < jsp->content.sequence.nelems; i++)
+				{
+					JsonbValue *v;
+
+					if (hasNext)
+						memset(&list, 0, sizeof(list));
+
+					jspGetSequenceElement(jsp, i, &elem);
+					res = executeItem(cxt, &elem, jb, plist);
+
+					if (jperIsError(res))
+						break;
+
+					if (!hasNext)
+					{
+						if (!found && res == jperOk)
+							break;
+						continue;
+					}
+
+					JsonValueListInitIterator(&list, &it);
+
+					while ((v = JsonValueListNext(&list, &it)))
+					{
+						res = executeItem(cxt, &next, v, found);
+
+						if (jperIsError(res) || (!found && res == jperOk))
+						{
+							i = jsp->content.sequence.nelems;
+							break;
+						}
+					}
+				}
+
+				break;
+			}
+
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", jsp->type);
 	}
diff --git a/src/backend/utils/adt/jsonpath_gram.y b/src/backend/utils/adt/jsonpath_gram.y
index 426dbb1..1de6f5a 100644
--- a/src/backend/utils/adt/jsonpath_gram.y
+++ b/src/backend/utils/adt/jsonpath_gram.y
@@ -56,6 +56,7 @@ static JsonPathParseItem *makeAny(int first, int last);
 static JsonPathParseItem *makeItemLikeRegex(JsonPathParseItem *expr,
 											JsonPathString *pattern,
 											JsonPathString *flags);
+static JsonPathParseItem *makeItemSequence(List *elems);
 
 /*
  * Bison doesn't allocate anything that needs to live across parser calls,
@@ -101,9 +102,9 @@ static JsonPathParseItem *makeItemLikeRegex(JsonPathParseItem *expr,
 %type	<value>		scalar_value path_primary expr array_accessor
 					any_path accessor_op key predicate delimited_predicate
 					index_elem starts_with_initial expr_or_predicate
-					datetime_template opt_datetime_template
+					datetime_template opt_datetime_template expr_seq expr_or_seq
 
-%type	<elems>		accessor_expr
+%type	<elems>		accessor_expr expr_list
 
 %type	<indexs>	index_list
 
@@ -127,7 +128,7 @@ static JsonPathParseItem *makeItemLikeRegex(JsonPathParseItem *expr,
 %%
 
 result:
-	pg_opt mode expr_or_predicate	{
+	pg_opt mode expr_or_seq	{
 										*result = palloc(sizeof(JsonPathParseResult));
 										(*result)->expr = $3;
 										(*result)->lax = $2;
@@ -146,6 +147,20 @@ pg_opt:
 	| /* EMPTY */					{ $$ = false; }
 	;
 
+expr_or_seq:
+	expr_or_predicate				{ $$ = $1; }
+	| expr_seq						{ $$ = $1; }
+	;
+
+expr_seq:
+	expr_list						{ $$ = makeItemSequence($1); }
+	;
+
+expr_list:
+	expr_or_predicate ',' expr_or_predicate	{ $$ = list_make2($1, $3); }
+	| expr_list ',' expr_or_predicate		{ $$ = lappend($1, $3); }
+	;
+
 mode:
 	STRICT_P						{ $$ = false; }
 	| LAX_P							{ $$ = true; }
@@ -201,6 +216,7 @@ path_primary:
 	| '$'							{ $$ = makeItemType(jpiRoot); }
 	| '@'							{ $$ = makeItemType(jpiCurrent); }
 	| LAST_P						{ $$ = makeItemType(jpiLast); }
+	| '(' expr_seq ')'				{ $$ = $2; }
 	;
 
 accessor_expr:
@@ -550,6 +566,16 @@ makeItemLikeRegex(JsonPathParseItem *expr, JsonPathString *pattern,
 	return v;
 }
 
+static JsonPathParseItem *
+makeItemSequence(List *elems)
+{
+	JsonPathParseItem  *v = makeItemType(jpiSequence);
+
+	v->value.sequence.elems = elems;
+
+	return v;
+}
+
 /*
  * Convert from XQuery regex flags to those recognized by our regex library.
  */
diff --git a/src/include/utils/jsonpath.h b/src/include/utils/jsonpath.h
index d56175f..4cbdce3 100644
--- a/src/include/utils/jsonpath.h
+++ b/src/include/utils/jsonpath.h
@@ -87,6 +87,7 @@ typedef enum JsonPathItemType
 	jpiLast,					/* LAST array subscript */
 	jpiStartsWith,				/* STARTS WITH predicate */
 	jpiLikeRegex,				/* LIKE_REGEX predicate */
+	jpiSequence,				/* sequence constructor: 'expr, ...' */
 } JsonPathItemType;
 
 /* XQuery regex mode flags for LIKE_REGEX predicate */
@@ -149,6 +150,12 @@ typedef struct JsonPathItem
 
 		struct
 		{
+			int32		nelems;
+			int32	   *elems;
+		}			sequence;
+
+		struct
+		{
 			char	   *data;	/* for bool, numeric and string/key */
 			int32		datalen;	/* filled only for string/key */
 		}			value;
@@ -176,6 +183,7 @@ extern bool jspGetBool(JsonPathItem *v);
 extern char *jspGetString(JsonPathItem *v, int32 *len);
 extern bool jspGetArraySubscript(JsonPathItem *v, JsonPathItem *from,
 								 JsonPathItem *to, int i);
+extern void jspGetSequenceElement(JsonPathItem *v, int i, JsonPathItem *elem);
 
 extern const char *jspOperationName(JsonPathItemType type);
 
@@ -229,6 +237,10 @@ struct JsonPathParseItem
 			uint32		flags;
 		}			like_regex;
 
+		struct {
+			List   *elems;
+		} sequence;
+
 		/* scalars */
 		Numeric numeric;
 		bool		boolean;
diff --git a/src/test/regress/expected/jsonb_jsonpath.out b/src/test/regress/expected/jsonb_jsonpath.out
index ef8db2d..c5252ad 100644
--- a/src/test/regress/expected/jsonb_jsonpath.out
+++ b/src/test/regress/expected/jsonb_jsonpath.out
@@ -2524,3 +2524,57 @@ ORDER BY s1.num, s2.num;
  {"s": "B"}    | {"s": "B"}    | false | true  | true  | true  | false
 (144 rows)
 
+-- extension: path sequences
+select jsonb_path_query('[1,2,3,4,5]', 'pg 10, 20, $[*], 30');
+ jsonb_path_query 
+------------------
+ 10
+ 20
+ 1
+ 2
+ 3
+ 4
+ 5
+ 30
+(8 rows)
+
+select jsonb_path_query('[1,2,3,4,5]', 'pg lax    10, 20, $[*].a, 30');
+ jsonb_path_query 
+------------------
+ 10
+ 20
+ 30
+(3 rows)
+
+select jsonb_path_query('[1,2,3,4,5]', 'pg strict 10, 20, $[*].a, 30');
+ERROR:  jsonpath member accessor can only be applied to an object
+select jsonb_path_query('[1,2,3,4,5]', 'pg -(10, 20, $[1 to 3], 30)');
+ jsonb_path_query 
+------------------
+ -10
+ -20
+ -2
+ -3
+ -4
+ -30
+(6 rows)
+
+select jsonb_path_query('[1,2,3,4,5]', 'pg lax (10, 20.5, $[1 to 3], "30").double()');
+ jsonb_path_query 
+------------------
+ 10
+ 20.5
+ 2
+ 3
+ 4
+ 30
+(6 rows)
+
+select jsonb_path_query('[1,2,3,4,5]', 'pg $[(0, $[*], 5) ? (@ == 3)]');
+ jsonb_path_query 
+------------------
+ 4
+(1 row)
+
+select jsonb_path_query('[1,2,3,4,5]', 'pg $[(0, $[*], 3) ? (@ == 3)]');
+ERROR:  jsonpath array subscript is not a single numeric value
diff --git a/src/test/regress/expected/jsonpath.out b/src/test/regress/expected/jsonpath.out
index 52b36a8..f6a132c 100644
--- a/src/test/regress/expected/jsonpath.out
+++ b/src/test/regress/expected/jsonpath.out
@@ -568,6 +568,42 @@ select '((($ + 1)).a + ((2)).b ? ((((@ > 1)) || (exists(@.c)))))'::jsonpath;
  (($ + 1)."a" + 2."b"?(@ > 1 || exists (@."c")))
 (1 row)
 
+select 'pg 1, 2 + 3, $.a[*] + 5'::jsonpath;
+         jsonpath          
+---------------------------
+ pg 1, 2 + 3, $."a"[*] + 5
+(1 row)
+
+select 'pg (1, 2, $.a)'::jsonpath;
+    jsonpath    
+----------------
+ pg 1, 2, $."a"
+(1 row)
+
+select 'pg (1, 2, $.a).a[*]'::jsonpath;
+        jsonpath         
+-------------------------
+ pg (1, 2, $."a")."a"[*]
+(1 row)
+
+select 'pg (1, 2, $.a) == 5'::jsonpath;
+        jsonpath         
+-------------------------
+ pg ((1, 2, $."a") == 5)
+(1 row)
+
+select 'pg $[(1, 2, $.a) to (3, 4)]'::jsonpath;
+           jsonpath            
+-------------------------------
+ pg $[(1, 2, $."a") to (3, 4)]
+(1 row)
+
+select 'pg $[(1, (2, $.a)), 3, (4, 5)]'::jsonpath;
+            jsonpath            
+--------------------------------
+ pg $[(1, (2, $."a")),3,(4, 5)]
+(1 row)
+
 select '$ ? (@.a < 1)'::jsonpath;
    jsonpath    
 ---------------
diff --git a/src/test/regress/sql/jsonb_jsonpath.sql b/src/test/regress/sql/jsonb_jsonpath.sql
index 591be00..0619ade 100644
--- a/src/test/regress/sql/jsonb_jsonpath.sql
+++ b/src/test/regress/sql/jsonb_jsonpath.sql
@@ -578,3 +578,12 @@ SELECT
 	jsonb_path_query_first(s1.j, '$.s > $s', vars => s2.j) gt
 FROM str s1, str s2
 ORDER BY s1.num, s2.num;
+
+-- extension: path sequences
+select jsonb_path_query('[1,2,3,4,5]', 'pg 10, 20, $[*], 30');
+select jsonb_path_query('[1,2,3,4,5]', 'pg lax    10, 20, $[*].a, 30');
+select jsonb_path_query('[1,2,3,4,5]', 'pg strict 10, 20, $[*].a, 30');
+select jsonb_path_query('[1,2,3,4,5]', 'pg -(10, 20, $[1 to 3], 30)');
+select jsonb_path_query('[1,2,3,4,5]', 'pg lax (10, 20.5, $[1 to 3], "30").double()');
+select jsonb_path_query('[1,2,3,4,5]', 'pg $[(0, $[*], 5) ? (@ == 3)]');
+select jsonb_path_query('[1,2,3,4,5]', 'pg $[(0, $[*], 3) ? (@ == 3)]');
diff --git a/src/test/regress/sql/jsonpath.sql b/src/test/regress/sql/jsonpath.sql
index 315e42f..2a53168 100644
--- a/src/test/regress/sql/jsonpath.sql
+++ b/src/test/regress/sql/jsonpath.sql
@@ -106,6 +106,13 @@ select '($)'::jsonpath;
 select '(($))'::jsonpath;
 select '((($ + 1)).a + ((2)).b ? ((((@ > 1)) || (exists(@.c)))))'::jsonpath;
 
+select 'pg 1, 2 + 3, $.a[*] + 5'::jsonpath;
+select 'pg (1, 2, $.a)'::jsonpath;
+select 'pg (1, 2, $.a).a[*]'::jsonpath;
+select 'pg (1, 2, $.a) == 5'::jsonpath;
+select 'pg $[(1, 2, $.a) to (3, 4)]'::jsonpath;
+select 'pg $[(1, (2, $.a)), 3, (4, 5)]'::jsonpath;
+
 select '$ ? (@.a < 1)'::jsonpath;
 select '$ ? (@.a < -1)'::jsonpath;
 select '$ ? (@.a < +1)'::jsonpath;
-- 
2.7.4

