Use nodes JsonFomat and JsonReturning in JsonCtorExpr transformation

Nikita Glukhov · Nikita Glukhov · commit fe6eb36fa44e · 2020-07-29T15:47:15.000+03:00
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
@@ -2304,11 +2304,7 @@ _copyJsonCtorExpr(const JsonCtorExpr *from)
 
 	COPY_NODE_FIELD(func);
 	COPY_SCALAR_FIELD(type);
-	COPY_SCALAR_FIELD(returning.format.type);
-	COPY_SCALAR_FIELD(returning.format.encoding);
-	COPY_LOCATION_FIELD(returning.format.location);
-	COPY_SCALAR_FIELD(returning.typid);
-	COPY_SCALAR_FIELD(returning.typmod);
+	COPY_NODE_FIELD(returning);
 	COPY_SCALAR_FIELD(absent_on_null);
 	COPY_SCALAR_FIELD(unique);
 
@@ -2355,7 +2351,7 @@ _copyJsonOutput(const JsonOutput *from)
 	JsonOutput	   *newnode = makeNode(JsonOutput);
 
 	COPY_NODE_FIELD(typeName);
-	COPY_SCALAR_FIELD(returning);
+	COPY_NODE_FIELD(returning);
 
 	return newnode;
 }
diff --git a/src/backend/nodes/equalfuncs.c b/src/backend/nodes/equalfuncs.c
@@ -853,11 +853,7 @@ _equalJsonCtorExpr(const JsonCtorExpr *a, const JsonCtorExpr *b)
 {
 	COMPARE_NODE_FIELD(func);
 	COMPARE_SCALAR_FIELD(type);
-	COMPARE_SCALAR_FIELD(returning.format.type);
-	COMPARE_SCALAR_FIELD(returning.format.encoding);
-	COMPARE_LOCATION_FIELD(returning.format.location);
-	COMPARE_SCALAR_FIELD(returning.typid);
-	COMPARE_SCALAR_FIELD(returning.typmod);
+	COMPARE_NODE_FIELD(returning);
 	COMPARE_SCALAR_FIELD(absent_on_null);
 	COMPARE_SCALAR_FIELD(unique);
 
diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
@@ -2291,7 +2291,13 @@ expression_tree_walker(Node *node,
 			}
 			break;
 		case T_JsonCtorExpr:
-			return walker(((JsonCtorExpr *) node)->func, context);
+			{
+				JsonCtorExpr *ctor = (JsonCtorExpr *) node;
+				
+				if (walker(ctor->func, context))
+					return true;
+			}
+			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d",
 				 (int) nodeTag(node));
@@ -3245,6 +3251,7 @@ expression_tree_mutator(Node *node,
 
 				FLATCOPY(newnode, jve, JsonCtorExpr);
 				MUTATE(newnode->func, jve->func, FuncExpr *);
+				MUTATE(newnode->returning, jve->returning, JsonReturning *);
 
 				return (Node *) newnode;
 			}
@@ -3958,9 +3965,25 @@ raw_expression_tree_walker(Node *node,
 			}
 			break;
 		case T_JsonCtorExpr:
-			return walker(((JsonCtorExpr *) node)->func, context);
+			{
+				JsonCtorExpr *ctor = (JsonCtorExpr *) node;
+
+				if (walker(ctor->func, context))
+					return true;
+				if (walker(ctor->returning, context))
+					return true;
+			}
+			break;
 		case T_JsonOutput:
-			return walker(((JsonOutput *) node)->typeName, context);
+			{
+				JsonOutput *out = (JsonOutput *) node;
+
+				if (walker(out->typeName, context))
+					return true;
+				if (walker(out->returning, context))
+					return true;
+			}
+			break;
 		case T_JsonKeyValue:
 			{
 				JsonKeyValue *jkv = (JsonKeyValue *) node;
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
@@ -1746,11 +1746,7 @@ _outJsonCtorExpr(StringInfo str, const JsonCtorExpr *node)
 
 	WRITE_NODE_FIELD(func);
 	WRITE_INT_FIELD(type);
-	WRITE_ENUM_FIELD(returning.format.type, JsonFormatType);
-	WRITE_ENUM_FIELD(returning.format.encoding, JsonEncoding);
-	WRITE_LOCATION_FIELD(returning.format.location);
-	WRITE_OID_FIELD(returning.typid);
-	WRITE_INT_FIELD(returning.typmod);
+	WRITE_NODE_FIELD(returning);
 	WRITE_BOOL_FIELD(unique);
 	WRITE_BOOL_FIELD(absent_on_null);
 }
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
@@ -1398,11 +1398,7 @@ _readJsonCtorExpr(void)
 
 	READ_NODE_FIELD(func);
 	READ_INT_FIELD(type);
-	READ_ENUM_FIELD(returning.format.type, JsonFormatType);
-	READ_ENUM_FIELD(returning.format.encoding, JsonEncoding);
-	READ_LOCATION_FIELD(returning.format.location);
-	READ_OID_FIELD(returning.typid);
-	READ_INT_FIELD(returning.typmod);
+	READ_NODE_FIELD(returning);
 	READ_BOOL_FIELD(unique);
 	READ_BOOL_FIELD(absent_on_null);
 
diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
@@ -3762,7 +3762,7 @@ checkJsonOutputFormat(ParseState *pstate, const JsonFormat *format,
 					  Oid targettype, bool allow_format_for_non_strings)
 {
 	if (!allow_format_for_non_strings &&
-		format->type != JS_FORMAT_DEFAULT &&
+		format->format != JS_FORMAT_DEFAULT &&
 		(targettype != BYTEAOID &&
 		 targettype != JSONOID &&
 		 targettype != JSONBOID))
@@ -3779,7 +3779,7 @@ checkJsonOutputFormat(ParseState *pstate, const JsonFormat *format,
 					 errmsg("cannot use JSON format with non-string output types")));
 	}
 
-	if (format->type == JS_FORMAT_JSON)
+	if (format->format == JS_FORMAT_JSON)
 	{
 		JsonEncoding enc = format->encoding != JS_ENC_DEFAULT ?
 						   format->encoding : JS_ENC_UTF8;
@@ -3807,23 +3807,25 @@ checkJsonOutputFormat(ParseState *pstate, const JsonFormat *format,
  * Assigns default format or checks specified format for its applicability to
  * the target type.
  */
-static void
+static JsonReturning *
 transformJsonOutput(ParseState *pstate, const JsonOutput *output,
-					bool allow_format, JsonReturning *ret)
+					bool allow_format)
 {
+	JsonReturning *ret;
+
 	/* if output clause is not specified, make default clause value */
 	if (!output)
 	{
-		ret->format.type = JS_FORMAT_DEFAULT;
-		ret->format.encoding = JS_ENC_DEFAULT;
-		ret->format.location = -1;
+		ret = makeNode(JsonReturning);
+
+		ret->format = makeJsonFormat(JS_FORMAT_DEFAULT, JS_ENC_DEFAULT, -1);
 		ret->typid = InvalidOid;
 		ret->typmod = -1;
 
-		return;
+		return ret;
 	}
 
-	*ret = output->returning;
+	ret = copyObject(output->returning);
 
 	typenameTypeIdAndMod(pstate, output->typeName, &ret->typid, &ret->typmod);
 
@@ -3832,20 +3834,73 @@ transformJsonOutput(ParseState *pstate, const JsonOutput *output,
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("returning SETOF types is not supported in SQL/JSON functions")));
 
-	if (ret->format.type == JS_FORMAT_DEFAULT)
+	if (ret->format->format == JS_FORMAT_DEFAULT)
 		/* assign JSONB format when returning jsonb, or JSON format otherwise */
-		ret->format.type =
+		ret->format->format =
 			ret->typid == JSONBOID ? JS_FORMAT_JSONB : JS_FORMAT_JSON;
 	else
-		checkJsonOutputFormat(pstate, &ret->format, ret->typid, allow_format);
+		checkJsonOutputFormat(pstate, ret->format, ret->typid, allow_format);
+
+	return ret;
+}
+
+/*
+ * Transform JSON output clause of JSON contructor functions.
+ *
+ * Derive RETURNING type, if not specified, from argument types.
+ */
+static JsonReturning *
+transformJsonCtorOutput(ParseState *pstate, JsonOutput *output, List *args)
+{
+	JsonReturning *returning = transformJsonOutput(pstate, output, true);
+
+	if (!OidIsValid(returning->typid))
+	{
+		ListCell   *lc;
+		bool		have_json = false;
+		bool		have_jsonb = false;
+
+		foreach(lc, args)
+		{
+			Node	   *expr = lfirst(lc);
+			Oid			typid = exprType(expr);
+
+			have_json |= typid == JSONOID;
+			have_jsonb |= typid == JSONBOID;
+
+			if (have_jsonb)
+				break;
+		}
+
+		if (have_jsonb)
+		{
+			returning->typid = JSONBOID;
+			returning->format->format = JS_FORMAT_JSONB;
+		}
+		else if (have_json)
+		{
+			returning->typid = JSONOID;
+			returning->format->format = JS_FORMAT_JSON;
+		}
+		else
+		{
+			/* XXX TEXTOID is default by standard */
+			returning->typid = JSONOID;
+			returning->format->format = JS_FORMAT_JSON;
+		}
+
+		returning->typmod = -1;
+	}
+
+	return returning;
 }
 
 /*
  * Coerce json[b]-valued function expression to the output type.
  */
 static Node *
-coerceJsonFuncExpr(ParseState *pstate, Node *expr, JsonReturning *returning,
-				   bool report_error)
+coerceJsonFuncExpr(ParseState *pstate, Node *expr,
+				   const JsonReturning *returning, bool report_error)
 {
 	Node	   *res;
 	int			location;
@@ -3858,16 +3913,16 @@ coerceJsonFuncExpr(ParseState *pstate, Node *expr, JsonReturning *returning,
 	location = exprLocation(expr);
 
 	if (location < 0)
-		location = returning ? returning->format.location : -1;
+		location = returning ? returning->format->location : -1;
 
 	/* special case for RETURNING bytea FORMAT json */
-	if (returning->format.type == JS_FORMAT_JSON &&
+	if (returning->format->format == JS_FORMAT_JSON &&
 		returning->typid == BYTEAOID)
 	{
 		/* encode json text into bytea using pg_convert_to() */
 		Node	   *texpr = coerce_to_specific_type(pstate, expr, TEXTOID,
 													"JSON_FUNCTION");
-		Const	   *enc = getJsonEncodingConst(&returning->format);
+		Const	   *enc = getJsonEncodingConst(returning->format);
 		FuncExpr   *fexpr = makeFuncExpr(F_PG_CONVERT_TO, BYTEAOID,
 										 list_make2(texpr, enc),
 										 InvalidOid, InvalidOid,
@@ -3908,7 +3963,7 @@ coerceJsonFuncExpr(ParseState *pstate, Node *expr, JsonReturning *returning,
 static Node *
 transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 {
-	JsonReturning returning;
+	JsonReturning *returning;
 	JsonCtorExpr *jsctor;
 	FuncExpr   *fexpr;
 	List	   *args = NIL;
@@ -3937,9 +3992,9 @@ transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 		}
 	}
 
-	transformJsonOutput(pstate, ctor->output, true, &returning);
+	returning = transformJsonCtorOutput(pstate, ctor->output, args);
 
-	if (returning.format.type == JS_FORMAT_JSONB)
+	if (returning->format->format == JS_FORMAT_JSONB)
 	{
 		funcid = args ? F_JSONB_BUILD_OBJECT_EXT : F_JSONB_BUILD_OBJECT_NOARGS;
 		funcrettype = JSONBOID;
@@ -3961,5 +4016,5 @@ transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 	jsctor->unique = ctor->unique;
 	jsctor->absent_on_null = ctor->absent_on_null;
 
-	return coerceJsonFuncExpr(pstate, (Node *) jsctor, &returning, true);
+	return coerceJsonFuncExpr(pstate, (Node *) jsctor, returning, true);
 }
diff --git a/src/include/nodes/primnodes.h b/src/include/nodes/primnodes.h
@@ -1265,7 +1265,7 @@ typedef struct JsonCtorExpr
 	Expr		xpr;
 	JsonCtorType type;			/* constructor type */
 	FuncExpr   *func;			/* underlying json[b]_xxx() function call */
-	JsonReturning returning;	/* RETURNING clause */
+	JsonReturning *returning;	/* RETURNING clause */
 	bool		absent_on_null;	/* ABSENT ON NULL? */
 	bool		unique;			/* WITH UNIQUE KEYS? (JSON_OBJECT[AGG] only) */
 } JsonCtorExpr;
diff --git a/src/test/regress/expected/sqljson.out b/src/test/regress/expected/sqljson.out
@@ -100,9 +100,9 @@ SELECT JSON_OBJECT('foo': NULL::jsonb FORMAT JSON);
 WARNING:  FORMAT JSON has no effect for json and jsonb types
 LINE 1: SELECT JSON_OBJECT('foo': NULL::jsonb FORMAT JSON);
                                               ^
-  json_object   
-----------------
- {"foo" : null}
+  json_object  
+---------------
+ {"foo": null}
 (1 row)
 
 SELECT JSON_OBJECT('foo': NULL::jsonb FORMAT JSON ENCODING UTF8);
@@ -176,9 +176,9 @@ SELECT JSON_OBJECT(
 	'c': json '[ 1,true,{ } ]',
 	'd': jsonb '{ "x" : 123.45 }'
 );
-                              json_object                               
-------------------------------------------------------------------------
- {"a" : "123", "1.23" : 123, "c" : [ 1,true,{ } ], "d" : {"x": 123.45}}
+                            json_object                            
+-------------------------------------------------------------------
+ {"a": "123", "c": [1, true, {}], "d": {"x": 123.45}, "1.23": 123}
 (1 row)
 
 SELECT JSON_OBJECT(
@@ -207,9 +207,9 @@ SELECT JSON_OBJECT('a': '123', 'b': JSON_OBJECT('a': 111, 'b': 'aaa'));
 (1 row)
 
 SELECT JSON_OBJECT('a': '123', 'b': JSON_OBJECT('a': 111, 'b': 'aaa' RETURNING jsonb));
-                 json_object                 
----------------------------------------------
- {"a" : "123", "b" : {"a": 111, "b": "aaa"}}
+                json_object                
+-------------------------------------------
+ {"a": "123", "b": {"a": 111, "b": "aaa"}}
 (1 row)
 
 SELECT JSON_OBJECT('a': JSON_OBJECT('b': 1 RETURNING text));
