Add JSON_QUERY support for row, array and domain types

Nikita Glukhov · Nikita Glukhov · commit b321bdfb20ef · 2019-06-17T18:39:54.000+03:00
diff --git a/src/backend/executor/execExpr.c b/src/backend/executor/execExpr.c
@@ -2177,6 +2177,8 @@ ExecInitExprRec(Expr *node, ExprState *state,
 						lappend(scratch.d.jsonexpr.args, var);
 				}
 
+				scratch.d.jsonexpr.cache = NULL;
+
 				if (jexpr->coercions)
 				{
 					JsonCoercion **coercion;
diff --git a/src/backend/executor/execExprInterp.c b/src/backend/executor/execExprInterp.c
@@ -73,6 +73,7 @@
 #include "utils/date.h"
 #include "utils/datum.h"
 #include "utils/expandedrecord.h"
+#include "utils/jsonapi.h"
 #include "utils/jsonb.h"
 #include "utils/jsonpath.h"
 #include "utils/lsyscache.h"
@@ -4249,6 +4250,13 @@ ExecEvalJsonExprCoercion(ExprEvalStep *op, ExprContext *econtext,
 	else if (op->d.jsonexpr.result_expr)
 		res = ExecEvalExprPassingCaseValue(op->d.jsonexpr.result_expr, econtext,
 										   isNull, res, *isNull);
+	else if (coercion && coercion->via_populate)
+		res = json_populate_type(res, JSONBOID,
+								 jexpr->returning.typid,
+								 jexpr->returning.typmod,
+								 &op->d.jsonexpr.cache,
+								 econtext->ecxt_per_query_memory,
+								 isNull);
 	/* else no coercion, simply return item */
 
 	return res;
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
@@ -3050,6 +3050,50 @@ populate_record_field(ColumnIOData *col,
 	}
 }
 
+/* recursively populate specified type from a json/jsonb value */
+Datum
+json_populate_type(Datum json_val, Oid json_type, Oid typid, int32 typmod,
+				   void **cache, MemoryContext mcxt, bool *isnull)
+{
+	JsValue		jsv = { 0 };
+	JsonbValue	jbv;
+
+	jsv.is_json = json_type == JSONOID;
+
+	if (*isnull)
+	{
+		if (jsv.is_json)
+			jsv.val.json.str = NULL;
+		else
+			jsv.val.jsonb = NULL;
+	}
+	else if (jsv.is_json)
+	{
+		text	   *json = DatumGetTextPP(json_val);
+
+		jsv.val.json.str = VARDATA_ANY(json);
+		jsv.val.json.len = VARSIZE_ANY_EXHDR(json);
+		jsv.val.json.type = JSON_TOKEN_INVALID; /* not used in populate_composite() */
+	}
+	else
+	{
+		Jsonb	   *jsonb = DatumGetJsonbP(json_val);
+
+		jsv.val.jsonb = &jbv;
+
+		/* fill binary jsonb value pointing to jb */
+		jbv.type = jbvBinary;
+		jbv.val.binary.data = &jsonb->root;
+		jbv.val.binary.len = VARSIZE(jsonb) - VARHDRSZ;
+	}
+
+	if (!*cache)
+		*cache = MemoryContextAllocZero(mcxt, sizeof(ColumnIOData));
+
+	return populate_record_field(*cache , typid, typmod, NULL, mcxt,
+								 PointerGetDatum(NULL), &jsv, isnull);
+}
+
 static RecordIOData *
 allocate_record_info(MemoryContext mcxt, int ncolumns)
 {
diff --git a/src/include/executor/execExpr.h b/src/include/executor/execExpr.h
@@ -680,6 +680,8 @@ typedef struct ExprEvalStep
 			ExprState  *default_on_error;	/* ON ERROR DEFAULT expression */
 			List	   *args;				/* passing arguments */
 
+			void	   *cache;				/* cache for json_populate_type() */
+
 			struct JsonCoercionsState
 			{
 				struct JsonCoercionState
diff --git a/src/include/utils/jsonapi.h b/src/include/utils/jsonapi.h
@@ -217,6 +217,10 @@ extern text *transform_json_string_values(text *json, void *action_state,
 extern char *JsonEncodeDateTime(char *buf, Datum value, Oid typid,
 				   const int *tzp);
 
+extern Datum json_populate_type(Datum json_val, Oid json_type,
+								Oid typid, int32 typmod,
+								void **cache, MemoryContext mcxt, bool *isnull);
+
 extern Json *JsonCreate(text *json);
 extern JsonbIteratorToken JsonIteratorNext(JsonIterator **pit, JsonbValue *val,
 				 bool skipNested);
diff --git a/src/test/regress/expected/jsonb_sqljson.out b/src/test/regress/expected/jsonb_sqljson.out
@@ -752,6 +752,53 @@ FROM
  4 | 4 | [4]
 (25 rows)
 
+-- Extension: record types returning
+CREATE TYPE sqljsonb_rec AS (a int, t text, js json, jb jsonb, jsa json[]);
+CREATE TYPE sqljsonb_reca AS (reca sqljsonb_rec[]);
+SELECT JSON_QUERY(jsonb '[{"a": 1, "b": "foo", "t": "aaa", "js": [1, "2", {}], "jb": {"x": [1, "2", {}]}},  {"a": 2}]', '$[0]' RETURNING sqljsonb_rec);
+                     json_query                      
+-----------------------------------------------------
+ (1,aaa,"[1, ""2"", {}]","{""x"": [1, ""2"", {}]}",)
+(1 row)
+
+SELECT * FROM unnest((JSON_QUERY(jsonb '{"jsa":  [{"a": 1, "b": ["foo"]}, {"a": 2, "c": {}}, 123]}', '$' RETURNING sqljsonb_rec)).jsa);
+         unnest         
+------------------------
+ {"a": 1, "b": ["foo"]}
+ {"a": 2, "c": {}}
+ 123
+(3 rows)
+
+SELECT * FROM unnest((JSON_QUERY(jsonb '{"reca": [{"a": 1, "t": ["foo", []]}, {"a": 2, "jb": [{}, true]}]}', '$' RETURNING sqljsonb_reca)).reca);
+ a |      t      | js |     jb     | jsa 
+---+-------------+----+------------+-----
+ 1 | ["foo", []] |    |            | 
+ 2 |             |    | [{}, true] | 
+(2 rows)
+
+-- Extension: array types returning
+SELECT JSON_QUERY(jsonb '[1,2,null,"3"]', '$[*]' RETURNING int[] WITH WRAPPER);
+  json_query  
+--------------
+ {1,2,NULL,3}
+(1 row)
+
+SELECT * FROM unnest(JSON_QUERY(jsonb '[{"a": 1, "t": ["foo", []]}, {"a": 2, "jb": [{}, true]}]', '$' RETURNING sqljsonb_rec[]));
+ a |      t      | js |     jb     | jsa 
+---+-------------+----+------------+-----
+ 1 | ["foo", []] |    |            | 
+ 2 |             |    | [{}, true] | 
+(2 rows)
+
+-- Extension: domain types returning
+SELECT JSON_QUERY(jsonb '{"a": 1}', '$.a' RETURNING sqljsonb_int_not_null);
+ json_query 
+------------
+          1
+(1 row)
+
+SELECT JSON_QUERY(jsonb '{"a": 1}', '$.b' RETURNING sqljsonb_int_not_null);
+ERROR:  domain sqljsonb_int_not_null does not allow null values
 -- Test constraints
 CREATE TABLE test_jsonb_constraints (
 	js text,
diff --git a/src/test/regress/sql/jsonb_sqljson.sql b/src/test/regress/sql/jsonb_sqljson.sql
@@ -213,6 +213,22 @@ FROM
 	generate_series(0, 4) x,
 	generate_series(0, 4) y;
 
+-- Extension: record types returning
+CREATE TYPE sqljsonb_rec AS (a int, t text, js json, jb jsonb, jsa json[]);
+CREATE TYPE sqljsonb_reca AS (reca sqljsonb_rec[]);
+
+SELECT JSON_QUERY(jsonb '[{"a": 1, "b": "foo", "t": "aaa", "js": [1, "2", {}], "jb": {"x": [1, "2", {}]}},  {"a": 2}]', '$[0]' RETURNING sqljsonb_rec);
+SELECT * FROM unnest((JSON_QUERY(jsonb '{"jsa":  [{"a": 1, "b": ["foo"]}, {"a": 2, "c": {}}, 123]}', '$' RETURNING sqljsonb_rec)).jsa);
+SELECT * FROM unnest((JSON_QUERY(jsonb '{"reca": [{"a": 1, "t": ["foo", []]}, {"a": 2, "jb": [{}, true]}]}', '$' RETURNING sqljsonb_reca)).reca);
+
+-- Extension: array types returning
+SELECT JSON_QUERY(jsonb '[1,2,null,"3"]', '$[*]' RETURNING int[] WITH WRAPPER);
+SELECT * FROM unnest(JSON_QUERY(jsonb '[{"a": 1, "t": ["foo", []]}, {"a": 2, "jb": [{}, true]}]', '$' RETURNING sqljsonb_rec[]));
+
+-- Extension: domain types returning
+SELECT JSON_QUERY(jsonb '{"a": 1}', '$.a' RETURNING sqljsonb_int_not_null);
+SELECT JSON_QUERY(jsonb '{"a": 1}', '$.b' RETURNING sqljsonb_int_not_null);
+
 -- Test constraints
 
 CREATE TABLE test_jsonb_constraints (

Original file line number	Diff line number	Diff line change
`@@ -2177,6 +2177,8 @@ ExecInitExprRec(Expr node, ExprState state,`
`2177`	`2177`	`lappend(scratch.d.jsonexpr.args, var);`
`2178`	`2178`	`}`
`2179`	`2179`
	`2180`	`+ scratch.d.jsonexpr.cache = NULL;`
	`2181`	`+`
`2180`	`2182`	`if (jexpr->coercions)`
`2181`	`2183`	`{`
`2182`	`2184`	`JsonCoercion **coercion;`
Original file line number	Diff line number	Diff line change
`@@ -680,6 +680,8 @@ typedef struct ExprEvalStep`
`680`	`680`	`ExprState default_on_error; / ON ERROR DEFAULT expression */`
`681`	`681`	`List args; / passing arguments */`
`682`	`682`
	`683`	`+ void cache; / cache for json_populate_type() */`
	`684`	`+`
`683`	`685`	`struct JsonCoercionsState`
`684`	`686`	`{`
`685`	`687`	`struct JsonCoercionState`