From 97ae5e4220aa6475e4711efe71016e268da071b8 Mon Sep 17 00:00:00 2001
From: Andrew Dunstan <andrew@dunslane.net>
Date: Thu, 16 Mar 2023 20:42:13 -0400
Subject: [PATCH 7/9] 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
Discussion: https://postgr.es/m/20220616233130.rparivafipt6doj3@alap3.anarazel.de
Discussion: https://postgr.es/m/abd9b83b-aa66-f230-3d6d-734817f0995d%40postgresql.org
---
 doc/src/sgml/func.sgml                | 10 ++++-
 src/backend/nodes/nodeFuncs.c         | 20 +++++++++-
 src/backend/parser/gram.y             |  7 +++-
 src/backend/parser/parse_expr.c       | 46 ++++++++++++++++-----
 src/backend/utils/adt/ruleutils.c     |  5 ++-
 src/include/nodes/parsenodes.h        |  8 +---
 src/test/regress/expected/sqljson.out | 57 +++++++++++++++++++++++++++
 src/test/regress/sql/sqljson.sql      | 10 +++++
 8 files changed, 139 insertions(+), 24 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 000a86c294..3f39b58802 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -17737,7 +17737,8 @@ $.* ? (@ like_regex "^\\d+$")
          <function>json</function> (
          <replaceable>expression</replaceable>
          <optional> <literal>FORMAT JSON</literal> <optional> <literal>ENCODING UTF8</literal> </optional></optional>
-         <optional> { <literal>WITH</literal> | <literal>WITHOUT</literal> } <literal>UNIQUE</literal> <optional> <literal>KEYS</literal> </optional></optional>)
+         <optional> { <literal>WITH</literal> | <literal>WITHOUT</literal> } <literal>UNIQUE</literal> <optional> <literal>KEYS</literal> </optional></optional>
+         <optional> <literal>RETURNING</literal> <replaceable>json_data_type</replaceable> </optional>)
         </para>
         <para>
          The <replaceable>expression</replaceable> can be any text type or a
@@ -17752,13 +17753,18 @@ $.* ? (@ like_regex "^\\d+$")
          <literal>json('{"a":123, "b":[true,"foo"], "a":"bar"}')</literal>
          <returnvalue>{"a":123, "b":[true,"foo"], "a":"bar"}</returnvalue>
         </para>
+        <para>
+         <literal>json('{"a":123,"b":[true,"foo"],"a":"bar"}' returning jsonb)</literal>
+         <returnvalue>{"a": "bar", "b": [true, "foo"]}</returnvalue>
+        </para>
        </entry>
       </row>
       <row>
        <entry role="func_table_entry">
         <para role="func_signature">
         <indexterm><primary>json_scalar</primary></indexterm>
-        <function>json_scalar</function> (<replaceable>expression</replaceable>)
+        <function>json_scalar</function> (<replaceable>expression</replaceable>
+        <optional> <literal>RETURNING</literal> <replaceable>json_data_type</replaceable> </optional>)
        </para>
        <para>
         Returns a JSON scalar value representing
diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index b7a101cfcc..c79bd03509 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -4333,9 +4333,25 @@ raw_expression_tree_walker_impl(Node *node,
 			}
 			break;
 		case T_JsonParseExpr:
-			return WALK(((JsonParseExpr *) node)->expr);
+			{
+				JsonParseExpr *jpe = (JsonParseExpr *) node;
+
+				if (WALK(jpe->expr))
+					return true;
+				if (WALK(jpe->output))
+					return true;
+			}
+			break;
 		case T_JsonScalarExpr:
-			return WALK(((JsonScalarExpr *) node)->expr);
+			{
+				JsonScalarExpr *jse = (JsonScalarExpr *) node;
+
+				if (WALK(jse->expr))
+					return true;
+				if (WALK(jse->output))
+					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 d517b48e22..89eeda46ce 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -15781,20 +15781,23 @@ func_expr_common_subexpr:
 					n->location = @1;
 					$$ = (Node *) n;
 				} 
-			| 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 '(' 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 2a7a0e28f9..5a58a6d77f 100644
--- a/src/backend/parser/parse_expr.c
+++ b/src/backend/parser/parse_expr.c
@@ -4396,19 +4396,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)
 	{
 		/*
@@ -4448,12 +4477,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 a02edbdd17..9338be0571 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -10067,8 +10067,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);
 }
 
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 5aa0c3920f..6ffa450ab1 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1726,12 +1726,6 @@ typedef enum JsonQuotes
 	JS_QUOTES_OMIT				/* OMIT QUOTES */
 } JsonQuotes;
 
-/*
- * JsonPathSpec -
- *		representation of JSON path constant
- */
-typedef char *JsonPathSpec;
-
 /*
  * JsonOutput -
  *		representation of JSON output clause (RETURNING type [FORMAT format])
@@ -1805,6 +1799,7 @@ typedef struct JsonParseExpr
 {
 	NodeTag		type;
 	JsonValueExpr *expr;		/* string expression */
+	JsonOutput *output;			/* RETURNING clause, if specified */
 	bool		unique_keys;	/* WITH UNIQUE KEYS? */
 	int			location;		/* token location, or -1 if unknown */
 } JsonParseExpr;
@@ -1817,6 +1812,7 @@ typedef struct JsonScalarExpr
 {
 	NodeTag		type;
 	Expr	   *expr;			/* scalar expression */
+	JsonOutput *output;			/* RETURNING clause, if specified */
 	int			location;		/* token location, or -1 if unknown */
 } JsonScalarExpr;
 
diff --git a/src/test/regress/expected/sqljson.out b/src/test/regress/expected/sqljson.out
index 615af42b8a..5866a0ad14 100644
--- a/src/test/regress/expected/sqljson.out
+++ b/src/test/regress/expected/sqljson.out
@@ -113,6 +113,49 @@ EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' WITHOUT UNIQUE KEYS);
    Output: JSON('123'::json)
 (2 rows)
 
+SELECT JSON('123' RETURNING text);
+ERROR:  cannot use RETURNING type text in JSON()
+LINE 1: SELECT JSON('123' RETURNING text);
+                                    ^
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123');
+         QUERY PLAN          
+-----------------------------
+ Result
+   Output: JSON('123'::json)
+(2 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' RETURNING json);
+         QUERY PLAN          
+-----------------------------
+ Result
+   Output: JSON('123'::json)
+(2 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' RETURNING jsonb);
+                  QUERY PLAN                  
+----------------------------------------------
+ Result
+   Output: JSON('123'::jsonb RETURNING jsonb)
+(2 rows)
+
+SELECT pg_typeof(JSON('123'));
+ pg_typeof 
+-----------
+ json
+(1 row)
+
+SELECT pg_typeof(JSON('123' RETURNING json));
+ pg_typeof 
+-----------
+ json
+(1 row)
+
+SELECT pg_typeof(JSON('123' RETURNING jsonb));
+ pg_typeof 
+-----------
+ jsonb
+(1 row)
+
 -- JSON_SCALAR()
 SELECT JSON_SCALAR();
 ERROR:  syntax error at or near ")"
@@ -204,6 +247,20 @@ EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR('123');
    Output: JSON_SCALAR('123'::text)
 (2 rows)
 
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR(123 RETURNING json);
+         QUERY PLAN         
+----------------------------
+ Result
+   Output: JSON_SCALAR(123)
+(2 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR(123 RETURNING jsonb);
+                 QUERY PLAN                 
+--------------------------------------------
+ Result
+   Output: JSON_SCALAR(123 RETURNING jsonb)
+(2 rows)
+
 -- JSON_SERIALIZE()
 SELECT JSON_SERIALIZE();
 ERROR:  syntax error at or near ")"
diff --git a/src/test/regress/sql/sqljson.sql b/src/test/regress/sql/sqljson.sql
index c8d3b80c9e..c2742b40f1 100644
--- a/src/test/regress/sql/sqljson.sql
+++ b/src/test/regress/sql/sqljson.sql
@@ -23,6 +23,14 @@ EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123'::bytea FORMAT JSON ENCODING UTF8)
 EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' WITH UNIQUE KEYS);
 EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' WITHOUT UNIQUE KEYS);
 
+SELECT JSON('123' RETURNING text);
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123');
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' RETURNING json);
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON('123' RETURNING jsonb);
+SELECT pg_typeof(JSON('123'));
+SELECT pg_typeof(JSON('123' RETURNING json));
+SELECT pg_typeof(JSON('123' RETURNING jsonb));
 
 -- JSON_SCALAR()
 SELECT JSON_SCALAR();
@@ -41,6 +49,8 @@ SELECT JSON_SCALAR('{}'::jsonb);
 
 EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR(123);
 EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR('123');
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR(123 RETURNING json);
+EXPLAIN (VERBOSE, COSTS OFF) SELECT JSON_SCALAR(123 RETURNING jsonb);
 
 -- JSON_SERIALIZE()
 SELECT JSON_SERIALIZE();
-- 
2.34.1

