Add JSON_QUERY support for row, array and domain types

Author: Nikita Glukhov

src/backend/executor/execExpr.c

@@ -2272,6 +2272,8 @@ ExecInitExprRec(Expr *node, ExprState *state,
 						lappend(scratch.d.jsonexpr.args, var);
 				}
 
+				scratch.d.jsonexpr.cache = NULL;
+
 				if (jexpr->coercions)
 				{
 					JsonCoercion **coercion;

src/backend/executor/execExprInterp.c

@@ -4563,6 +4563,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;

src/backend/utils/adt/jsonfuncs.c

@@ -2628,11 +2628,11 @@ populate_array_dim_jsonb(PopulateArrayContext *ctx, /* context */
 
 	check_stack_depth();
 
-	if (jbv->type != jbvBinary || !JsonContainerIsArray(jbc))
+	if (jbv->type != jbvBinary ||
+		!JsonContainerIsArray(jbc) ||
+		JsonContainerIsScalar(jbc))
 		populate_array_report_expected_array(ctx, ndim - 1);
 
-	Assert(!JsonContainerIsScalar(jbc));
-
 	it = JsonbIteratorInit(jbc);
 
 	tok = JsonbIteratorNext(&it, &val, true);
@@ -3104,6 +3104,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

@@ -684,6 +684,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

src/include/utils/jsonfuncs.h

@@ -58,4 +58,8 @@ 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);
+
 #endif

src/test/regress/expected/jsonb_sqljson.out

@@ -758,6 +758,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,

src/test/regress/sql/jsonb_sqljson.sql

@@ -214,6 +214,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 (