Add JSON_QUERY support for row, array and domain types

Nikita Glukhov · Nikita Glukhov · commit f51fb8f6b951 · 2018-06-14T17:55:44.000+03:00
diff --git a/src/backend/executor/execExpr.c b/src/backend/executor/execExpr.c
@@ -2187,6 +2187,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"
@@ -4232,6 +4233,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
@@ -3051,6 +3051,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
@@ -662,6 +662,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
@@ -206,6 +206,10 @@ extern text *transform_json_string_values(text *json, void *action_state,
 
 extern char *JsonEncodeDateTime(char *buf, Datum value, Oid typid, int *tz);
 
+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
@@ -751,6 +751,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
`@@ -2187,6 +2187,8 @@ ExecInitExprRec(Expr node, ExprState state,`
`2187`	`2187`	`lappend(scratch.d.jsonexpr.args, var);`
`2188`	`2188`	`}`
`2189`	`2189`
	`2190`	`+ scratch.d.jsonexpr.cache = NULL;`
	`2191`	`+`
`2190`	`2192`	`if (jexpr->coercions)`
`2191`	`2193`	`{`
`2192`	`2194`	`JsonCoercion **coercion;`
Original file line number	Diff line number	Diff line change
`@@ -662,6 +662,8 @@ typedef struct ExprEvalStep`
`662`	`662`	`ExprState default_on_error; / ON ERROR DEFAULT expression */`
`663`	`663`	`List args; / passing arguments */`
`664`	`664`
	`665`	`+ void cache; / cache for json_populate_type() */`
	`666`	`+`
`665`	`667`	`struct JsonCoercionsState`
`666`	`668`	`{`
`667`	`669`	`struct JsonCoercionState`