0006-RETURNING-clause-for-JSON-and-JSON_SCALAR-v64.patch

From 31ead5d65852012876af2194f2ff8daf6fceeaf9 Mon Sep 17 00:00:00 2001
From: Andrew Dunstan <andrew@dunslane.net>
Date: Sat, 5 Mar 2022 08:07:15 -0500
Subject: [PATCH 6/6] RETURNING clause for JSON() and JSON_SCALAR()

---
 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        |  2 +
 src/test/regress/expected/sqljson.out | 57 +++++++++++++++++++++++++++
 src/test/regress/sql/sqljson.sql      | 10 +++++
 6 files changed, 113 insertions(+), 14 deletions(-)

diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 175e14e51d..70f541ece0 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -15401,21 +15401,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 bbd80e5c6e..18622fd013 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 c4cac590af..bc5b545fd1 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -10052,8 +10052,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 13a18e02a2..cf19a225db 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1643,6 +1643,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;
@@ -1655,6 +1656,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 11f5eb2d2c..6cadd87868 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 98bd93c110..51fc659b58 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.25.1