RETURNING clause for JSON() and JSON_SCALAR()

This patch is extracted from a larger patch that allowed setting the
default returned value from these functions to json or jsonb. That had
problems, but this piece of it is fine. For these functions only json or
jsonb can be specified in the RETURNING clause.

Extracted from an original patch from Nikita Glukhov

Reviewers have included (in no particular order) Andres Freund, Alexander
Korotkov, Pavel Stehule, Andrew Alsup, Erik Rijkers, Zihong Yu,
Himanshu Upadhyaya, Daniel Gustafsson, Justin Pryzby.

Discussion: https://postgr.es/m/cd0bb935-0158-78a7-08b5-904886deac4b@postgrespro.ru

6 files changed, 66 insertions, 16 deletions

diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 56505557bff..11c016495e3 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -2354,6 +2354,7 @@ _copyJsonParseExpr(const JsonParseExpr *from)
 	JsonParseExpr  *newnode = makeNode(JsonParseExpr);
 
 	COPY_NODE_FIELD(expr);
+	COPY_NODE_FIELD(output);
 	COPY_SCALAR_FIELD(unique_keys);
 	COPY_LOCATION_FIELD(location);
 
@@ -2369,6 +2370,7 @@ _copyJsonScalarExpr(const JsonScalarExpr *from)
 	JsonScalarExpr *newnode = makeNode(JsonScalarExpr);
 
 	COPY_NODE_FIELD(expr);
+	COPY_NODE_FIELD(output);
 	COPY_LOCATION_FIELD(location);
 
 	return newnode;
diff --git a/src/backend/nodes/equalfuncs.c b/src/backend/nodes/equalfuncs.c
index 9ea3c5abf23..722dbe6a0d8 100644
--- a/src/backend/nodes/equalfuncs.c
+++ b/src/backend/nodes/equalfuncs.c
@@ -875,6 +875,7 @@ static bool
 _equalJsonParseExpr(const JsonParseExpr *a, const JsonParseExpr *b)
 {
 	COMPARE_NODE_FIELD(expr);
+	COMPARE_NODE_FIELD(output);
 	COMPARE_SCALAR_FIELD(unique_keys);
 	COMPARE_LOCATION_FIELD(location);
 
@@ -885,6 +886,7 @@ static bool
 _equalJsonScalarExpr(const JsonScalarExpr *a, const JsonScalarExpr *b)
 {
 	COMPARE_NODE_FIELD(expr);
+	COMPARE_NODE_FIELD(output);
 	COMPARE_LOCATION_FIELD(location);
 
 	return true;
diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index 4789ba69113..a094317bfc1 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -4364,9 +4364,25 @@ raw_expression_tree_walker(Node *node,
 			}
 			break;
 		case T_JsonParseExpr:
-			return walker(((JsonParseExpr *) node)->expr, context);
+			{
+				JsonParseExpr *jpe = (JsonParseExpr *) node;
+
+				if (walker(jpe->expr, context))
+					return true;
+				if (walker(jpe->output, context))
+					return true;
+			}
+			break;
 		case T_JsonScalarExpr:
-			return walker(((JsonScalarExpr *) node)->expr, context);
+			{
+				JsonScalarExpr *jse = (JsonScalarExpr *) node;
+
+				if (walker(jse->expr, context))
+					return true;
+				if (walker(jse->output, context))
+					return true;
+			}
+			break;
 		case T_JsonSerializeExpr:
 			{
 				JsonSerializeExpr *jse = (JsonSerializeExpr *) node;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index eefcf901879..e5a3c528aad 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -15614,21 +15614,24 @@ json_func_expr:
 		;
 
 json_parse_expr:
-			JSON '(' json_value_expr json_key_uniqueness_constraint_opt ')'
+			JSON '(' json_value_expr json_key_uniqueness_constraint_opt
+					 json_returning_clause_opt ')'
 				{
 					JsonParseExpr *n = makeNode(JsonParseExpr);
 					n->expr = (JsonValueExpr *) $3;
 					n->unique_keys = $4;
+					n->output = (JsonOutput *) $5;
 					n->location = @1;
 					$$ = (Node *) n;
 				}
 		;
 
 json_scalar_expr:
-			JSON_SCALAR '(' a_expr ')'
+			JSON_SCALAR '(' a_expr json_returning_clause_opt ')'
 				{
 					JsonScalarExpr *n = makeNode(JsonScalarExpr);
 					n->expr = (Expr *) $3;
+					n->output = (JsonOutput *) $4;
 					n->location = @1;
 					$$ = (Node *) n;
 				}
diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
index e1400072501..31f0c9f693d 100644
--- a/src/backend/parser/parse_expr.c
+++ b/src/backend/parser/parse_expr.c
@@ -4450,19 +4450,48 @@ transformJsonFuncExpr(ParseState *pstate, JsonFuncExpr *func)
 	return (Node *) jsexpr;
 }
 
+static JsonReturning *
+transformJsonConstructorRet(ParseState *pstate, JsonOutput *output, const char *fname)
+{
+	JsonReturning *returning;
+
+	if (output)
+	{
+		returning = transformJsonOutput(pstate, output, false);
+
+		Assert(OidIsValid(returning->typid));
+
+		if (returning->typid != JSONOID && returning->typid != JSONBOID)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("cannot use RETURNING type %s in %s",
+							format_type_be(returning->typid), fname),
+					 parser_errposition(pstate, output->typeName->location)));
+	}
+	else
+	{
+		Oid			targettype = JSONOID;
+		JsonFormatType format = JS_FORMAT_JSON;
+
+		returning = makeNode(JsonReturning);
+		returning->format = makeJsonFormat(format, JS_ENC_DEFAULT, -1);
+		returning->typid = targettype;
+		returning->typmod = -1;
+	}
+
+	return returning;
+}
+
 /*
  * Transform a JSON() expression.
  */
 static Node *
 transformJsonParseExpr(ParseState *pstate, JsonParseExpr *jsexpr)
 {
-	JsonReturning *returning = makeNode(JsonReturning);
+	JsonReturning *returning = transformJsonConstructorRet(pstate, jsexpr->output,
+													"JSON()");
 	Node	   *arg;
 
-	returning->format = makeJsonFormat(JS_FORMAT_JSON, JS_ENC_DEFAULT, -1);
-	returning->typid = JSONOID;
-	returning->typmod = -1;
-
 	if (jsexpr->unique_keys)
 	{
 		/*
@@ -4502,12 +4531,9 @@ transformJsonParseExpr(ParseState *pstate, JsonParseExpr *jsexpr)
 static Node *
 transformJsonScalarExpr(ParseState *pstate, JsonScalarExpr *jsexpr)
 {
-	JsonReturning *returning = makeNode(JsonReturning);
 	Node	   *arg = transformExprRecurse(pstate, (Node *) jsexpr->expr);
-
-	returning->format = makeJsonFormat(JS_FORMAT_JSON, JS_ENC_DEFAULT, -1);
-	returning->typid = JSONOID;
-	returning->typmod = -1;
+	JsonReturning *returning = transformJsonConstructorRet(pstate, jsexpr->output,
+													"JSON_SCALAR()");
 
 	if (exprType(arg) == UNKNOWNOID)
 		arg = coerce_to_specific_type(pstate, arg, TEXTOID, "JSON_SCALAR");
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 010d5a7a751..4458d2ff90a 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -10092,8 +10092,9 @@ get_json_constructor_options(JsonConstructorExpr *ctor, StringInfo buf)
 	if (ctor->unique)
 		appendStringInfoString(buf, " WITH UNIQUE KEYS");
 
-	if (ctor->type != JSCTOR_JSON_PARSE &&
-		ctor->type != JSCTOR_JSON_SCALAR)
+	if (!((ctor->type == JSCTOR_JSON_PARSE ||
+		   ctor->type == JSCTOR_JSON_SCALAR) &&
+		  ctor->returning->typid == JSONOID))
 		get_json_returning(ctor->returning, buf, true);
 }