Add JSON_QUERY support for row, array and domain types

Author: Nikita Glukhov

src/backend/executor/execExpr.c

@@ -2101,6 +2101,8 @@ ExecInitExprRec(Expr *node, PlanState *parent, ExprState *state,
 						lappend(scratch.d.jsonexpr.args, var);
 				}
 
+				scratch.d.jsonexpr.cache = NULL;
+
 				memset(&scratch.d.jsonexpr.scalar, 0,
 					   sizeof(scratch.d.jsonexpr.scalar));
 

src/backend/executor/execExprInterp.c

@@ -70,6 +70,7 @@
 #include "pgstat.h"
 #include "utils/builtins.h"
 #include "utils/date.h"
+#include "utils/jsonapi.h"
 #include "utils/jsonb.h"
 #include "utils/jsonpath.h"
 #include "utils/lsyscache.h"
@@ -3739,6 +3740,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 (jexpr->coerce_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;

src/backend/parser/parse_expr.c

@@ -4483,11 +4483,10 @@ coerceJsonExpr(ParseState *pstate, Node *expr, JsonReturning *returning,
 
 	typtype = get_typtype(returning->typid);
 
-	if (coerce_via_populate &&
-		(returning->typid == RECORDOID ||
-		 typtype == TYPTYPE_COMPOSITE ||
-		 typtype == TYPTYPE_DOMAIN ||
-		 type_is_array(returning->typid)))
+	if (returning->typid == RECORDOID ||
+		typtype == TYPTYPE_COMPOSITE ||
+		typtype == TYPTYPE_DOMAIN ||
+		type_is_array(returning->typid))
 		*coerce_via_populate = true;
 	else
 		*coerce_via_io = true;
@@ -4542,7 +4541,7 @@ transformJsonFuncExprOutput(ParseState *pstate,	JsonFuncExpr *func,
 			jsexpr->result_expr = coerceJsonExpr(pstate, placeholder,
 												 &jsexpr->returning,
 												 &jsexpr->coerce_via_io,
-												 NULL);
+												 &jsexpr->coerce_via_populate);
 		}
 	}
 	else

src/backend/utils/adt/jsonfuncs.c

@@ -3047,6 +3047,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)
 {

src/include/executor/execExpr.h

@@ -584,6 +584,8 @@ typedef struct ExprEvalStep
 			ExprState  *default_on_error;	/* ON ERROR DEFAULT expression */
 			List	   *args;				/* passing arguments */
 
+			void	   *cache;				/* cache for json_populate_type() */
+
 			struct JsonScalarCoercions
 			{
 				struct JsonScalarCoercionExprState

src/include/nodes/primnodes.h

@@ -1291,6 +1291,7 @@ typedef struct JsonExpr
 	Node	   *raw_expr;		/* raw context item expression */
 	Node	   *formatted_expr;	/* formatted context item expression */
 	Node	   *result_expr;	/* resulting expression (coerced to RETURNING type) */
+	bool		coerce_via_populate; /* coerce result using json_populate_type() */
 	bool		coerce_via_io;	/* coerce result using type input function */
 	Oid			coerce_via_io_collation; /* collation for conversion through I/O */
 	JsonFormat	format;			/* context item format (JSON/JSONB) */

src/include/utils/jsonapi.h

@@ -190,6 +190,10 @@ extern Jsonb *transform_jsonb_string_values(Jsonb *jsonb, void *action_state,
 extern text *transform_json_string_values(text *json, void *action_state,
 							 JsonTransformStringValuesAction transform_action);
 
+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);

src/test/regress/expected/jsonb_sqljson.out

@@ -751,6 +751,53 @@ FROM
  4 | 4 | [4]
 (25 rows)
 
+-- Conversion to record types
+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);
+                      ?column?                       
+-----------------------------------------------------
+ (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)
+
+-- Conversion to array types
+SELECT JSON_QUERY(jsonb '[1,2,null,"3"]', '$[*]' RETURNING int[] WITH WRAPPER);
+   ?column?   
+--------------
+ {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)
+
+-- Conversion to domain types
+SELECT JSON_QUERY(jsonb '{"a": 1}', '$.a' RETURNING sqljsonb_int_not_null);
+ ?column? 
+----------
+        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,

src/test/regress/sql/jsonb_sqljson.sql

@@ -213,6 +213,22 @@ FROM
 	generate_series(0, 4) x,
 	generate_series(0, 4) y;
 
+-- Conversion to record types
+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);
+
+-- Conversion to array types
+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[]));
+
+-- Conversion to domain types
+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 (