Refactor output coercion in SQL/JSON constructors

Nikita Glukhov · Nikita Glukhov · commit 39d0a434e3cb · 2020-07-29T15:47:17.000+03:00
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -3117,6 +3117,7 @@ JumbleExpr(pgssJumbleState *jstate, Node *node)
 				JsonCtorExpr *ctor = (JsonCtorExpr *) node;
 
 				JumbleExpr(jstate, (Node *) ctor->func);
+				JumbleExpr(jstate, (Node *) ctor->coercion);
 				JumbleExpr(jstate, (Node *) ctor->returning);
 				APP_JUMB(ctor->type);
 				APP_JUMB(ctor->unique);
diff --git a/src/backend/executor/execExpr.c b/src/backend/executor/execExpr.c
@@ -2141,40 +2141,55 @@ ExecInitExprRec(Expr *node, ExprState *state,
 				if (ctor->func)
 				{
 					ExecInitExprRec(ctor->func, state, resv, resnull);
-					break;
 				}
-
-				scratch.opcode = EEOP_JSON_CTOR;
-				scratch.d.json_ctor.ctor = ctor;
-				scratch.d.json_ctor.arg_values = palloc(sizeof(Datum) * nargs);
-				scratch.d.json_ctor.arg_nulls = palloc(sizeof(bool) * nargs);
-				scratch.d.json_ctor.arg_types = palloc(sizeof(Oid) * nargs);
-				scratch.d.json_ctor.nargs = nargs;
-
-				foreach(lc, args)
+				else
 				{
-					Expr	   *arg = (Expr *) lfirst(lc);
+					scratch.opcode = EEOP_JSON_CTOR;
+					scratch.d.json_ctor.ctor = ctor;
+					scratch.d.json_ctor.arg_values = palloc(sizeof(Datum) * nargs);
+					scratch.d.json_ctor.arg_nulls = palloc(sizeof(bool) * nargs);
+					scratch.d.json_ctor.arg_types = palloc(sizeof(Oid) * nargs);
+					scratch.d.json_ctor.nargs = nargs;
+
+					foreach(lc, args)
+					{
+						Expr	   *arg = (Expr *) lfirst(lc);
 
-					scratch.d.json_ctor.arg_types[argno] = exprType((Node *) arg);
+						scratch.d.json_ctor.arg_types[argno] = exprType((Node *) arg);
 
-					if (IsA(arg, Const))
-					{
-						/* Don't evaluate const arguments every round */
-						Const	   *con = (Const *) arg;
+						if (IsA(arg, Const))
+						{
+							/* Don't evaluate const arguments every round */
+							Const	   *con = (Const *) arg;
 
-						scratch.d.json_ctor.arg_values[argno] = con->constvalue;
-						scratch.d.json_ctor.arg_nulls[argno] = con->constisnull;
-					}
-					else
-					{
-						ExecInitExprRec(arg, state,
-										&scratch.d.json_ctor.arg_values[argno],
-										&scratch.d.json_ctor.arg_nulls[argno]);
+							scratch.d.json_ctor.arg_values[argno] = con->constvalue;
+							scratch.d.json_ctor.arg_nulls[argno] = con->constisnull;
+						}
+						else
+						{
+							ExecInitExprRec(arg, state,
+											&scratch.d.json_ctor.arg_values[argno],
+											&scratch.d.json_ctor.arg_nulls[argno]);
+						}
+						argno++;
 					}
-					argno++;
+
+					ExprEvalPushStep(state, &scratch);
 				}
 
-				ExprEvalPushStep(state, &scratch);
+				if (ctor->coercion)
+				{
+					Datum	   *innermost_caseval = state->innermost_caseval;
+					bool	   *innermost_isnull = state->innermost_casenull;
+
+					state->innermost_caseval = resv;
+					state->innermost_casenull = resnull;
+
+					ExecInitExprRec(ctor->coercion, state, resv, resnull);
+
+					state->innermost_caseval = innermost_caseval;
+					state->innermost_casenull = innermost_isnull;
+				}
 			}
 			break;
 
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
@@ -2302,9 +2302,10 @@ _copyJsonCtorExpr(const JsonCtorExpr *from)
 {
 	JsonCtorExpr *newnode = makeNode(JsonCtorExpr);
 
+	COPY_SCALAR_FIELD(type);
 	COPY_NODE_FIELD(args);
 	COPY_NODE_FIELD(func);
-	COPY_SCALAR_FIELD(type);
+	COPY_NODE_FIELD(coercion);
 	COPY_NODE_FIELD(returning);
 	COPY_SCALAR_FIELD(absent_on_null);
 	COPY_SCALAR_FIELD(unique);
diff --git a/src/backend/nodes/equalfuncs.c b/src/backend/nodes/equalfuncs.c
@@ -851,9 +851,10 @@ _equalJsonValueExpr(const JsonValueExpr *a, const JsonValueExpr *b)
 static bool
 _equalJsonCtorExpr(const JsonCtorExpr *a, const JsonCtorExpr *b)
 {
+	COMPARE_SCALAR_FIELD(type);
 	COMPARE_NODE_FIELD(args);
 	COMPARE_NODE_FIELD(func);
-	COMPARE_SCALAR_FIELD(type);
+	COMPARE_NODE_FIELD(coercion);
 	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
@@ -266,7 +266,7 @@ exprType(const Node *expr)
 			}
 			break;
 		case T_JsonCtorExpr:
-			type = ((const JsonCtorExpr *) expr)->returning->format->format == JS_FORMAT_JSONB ? JSONBOID : JSONOID;
+			type = ((const JsonCtorExpr *) expr)->returning->typid;
 			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
@@ -924,8 +924,14 @@ exprCollation(const Node *expr)
 			coll = exprCollation((Node *) ((const JsonValueExpr *) expr)->formatted_expr);
 			break;
 		case T_JsonCtorExpr:
-			/* coll = exprCollation((Node *) ((const JsonCtorExpr *) expr)->func); */
-			coll = InvalidOid;	/* keep compiler quiet */
+			{
+				const JsonCtorExpr *ctor = (const JsonCtorExpr *) expr;
+
+				if (ctor->coercion)
+					coll = exprCollation((Node *) ctor->coercion);
+				else
+					coll = InvalidOid;
+			}
 			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
@@ -1135,8 +1141,14 @@ exprSetCollation(Node *expr, Oid collation)
 							 collation);
 			break;
 		case T_JsonCtorExpr:
-			/* exprSetCollation((Node *) ((const JsonCtorExpr *) expr)->func, collation); */
-			Assert(!OidIsValid(collation)); /* result is always an json[b] type */
+			{
+				JsonCtorExpr *ctor = (JsonCtorExpr *) expr;
+
+				if (ctor->coercion)
+					exprSetCollation((Node *) ctor->coercion, collation);
+				else
+					Assert(!OidIsValid(collation)); /* result is always an json[b] type */
+			}
 			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
@@ -2299,6 +2311,8 @@ expression_tree_walker(Node *node,
 					return true;
 				if (walker(ctor->func, context))
 					return true;
+				if (walker(ctor->coercion, context))
+					return true;
 			}
 			break;
 		default:
@@ -3255,6 +3269,7 @@ expression_tree_mutator(Node *node,
 				FLATCOPY(newnode, jve, JsonCtorExpr);
 				MUTATE(newnode->args, jve->args, List *);
 				MUTATE(newnode->func, jve->func, Expr *);
+				MUTATE(newnode->coercion, jve->coercion, Expr *);
 				MUTATE(newnode->returning, jve->returning, JsonReturning *);
 
 				return (Node *) newnode;
@@ -3976,6 +3991,8 @@ raw_expression_tree_walker(Node *node,
 					return true;
 				if (walker(ctor->func, context))
 					return true;
+				if (walker(ctor->coercion, context))
+					return true;
 				if (walker(ctor->returning, context))
 					return true;
 			}
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
@@ -1746,6 +1746,7 @@ _outJsonCtorExpr(StringInfo str, const JsonCtorExpr *node)
 
 	WRITE_NODE_FIELD(args);
 	WRITE_NODE_FIELD(func);
+	WRITE_NODE_FIELD(coercion);
 	WRITE_INT_FIELD(type);
 	WRITE_NODE_FIELD(returning);
 	WRITE_BOOL_FIELD(unique);
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
@@ -1398,6 +1398,7 @@ _readJsonCtorExpr(void)
 
 	READ_NODE_FIELD(args);
 	READ_NODE_FIELD(func);
+	READ_NODE_FIELD(coercion);
 	READ_INT_FIELD(type);
 	READ_NODE_FIELD(returning);
 	READ_BOOL_FIELD(unique);
diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
@@ -3974,6 +3974,46 @@ coerceJsonFuncExpr(ParseState *pstate, Node *expr,
 	return res;
 }
 
+static Node *
+makeJsonCtorExpr(ParseState *pstate, JsonCtorType type, List *args, Expr *fexpr,
+				 JsonReturning *returning, bool unique, bool absent_on_null,
+				 int location)
+{
+	Node	   *placeholder;
+	Node	   *coercion;
+	JsonCtorExpr *jsctor = makeNode(JsonCtorExpr);
+	Oid			intermediate_typid =
+		returning->format->format == JS_FORMAT_JSONB ? JSONBOID : JSONOID;
+
+	jsctor->args = args;
+	jsctor->func = fexpr;
+	jsctor->type = type;
+	jsctor->returning = returning;
+	jsctor->unique = unique;
+	jsctor->absent_on_null = absent_on_null;
+	jsctor->location = location;
+
+	if (fexpr)
+		placeholder = makeCaseTestExpr((Node *) fexpr);
+	else
+	{
+		CaseTestExpr *cte = makeNode(CaseTestExpr);
+
+		cte->typeId = intermediate_typid;
+		cte->typeMod = -1;
+		cte->collation = InvalidOid;
+
+		placeholder = (Node *) cte;
+	}
+
+	coercion = coerceJsonFuncExpr(pstate, placeholder, returning, true);
+
+	if (coercion != placeholder)
+		jsctor->coercion = (Expr *) coercion;
+
+	return (Node *) jsctor;
+}
+
 /*
  * Transform JSON_OBJECT() constructor.
  *
@@ -3986,7 +4026,7 @@ coerceJsonFuncExpr(ParseState *pstate, Node *expr,
 static Node *
 transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 {
-	JsonCtorExpr *jsctor;
+	JsonReturning *returning;
 	List	   *args = NIL;
 
 	/* transform key-value pairs, if any */
@@ -4007,15 +4047,11 @@ transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 		}
 	}
 
-	jsctor = makeNode(JsonCtorExpr);
-	jsctor->args = args;
-	jsctor->type = JSCTOR_JSON_OBJECT;
-	jsctor->returning = transformJsonCtorOutput(pstate, ctor->output, args);
-	jsctor->unique = ctor->unique;
-	jsctor->absent_on_null = ctor->absent_on_null;
-	jsctor->location = ctor->location;
+	returning = transformJsonCtorOutput(pstate, ctor->output, args);
 
-	return coerceJsonFuncExpr(pstate, (Node *) jsctor, jsctor->returning, true);
+	return makeJsonCtorExpr(pstate, JSCTOR_JSON_OBJECT, args, NULL,
+							returning, ctor->unique, ctor->absent_on_null,
+							ctor->location);
 }
 
 /*
@@ -4093,7 +4129,6 @@ transformJsonAggCtor(ParseState *pstate, JsonAggCtor *agg_ctor,
 					 bool unique, bool absent_on_null)
 {
 	Oid			aggfnoid;
-	JsonCtorExpr *jsctor;
 	Node	   *node;
 	Expr	   *aggfilter = agg_ctor->agg_filter ? (Expr *)
 		transformWhereClause(pstate, agg_ctor->agg_filter,
@@ -4156,15 +4191,8 @@ transformJsonAggCtor(ParseState *pstate, JsonAggCtor *agg_ctor,
 		node = (Node *) aggref;
 	}
 
-	jsctor = makeNode(JsonCtorExpr);
-	jsctor->func = (Expr *) node;
-	jsctor->type = ctor_type;
-	jsctor->returning = returning;
-	jsctor->unique = unique;
-	jsctor->absent_on_null = absent_on_null;
-	jsctor->location = agg_ctor->location;
-
-	return coerceJsonFuncExpr(pstate, (Node *) jsctor, returning, true);
+	return makeJsonCtorExpr(pstate, ctor_type, NIL, (Expr *) node, returning,
+							unique, absent_on_null, agg_ctor->location);
 }
 
 /*
@@ -4277,7 +4305,7 @@ transformJsonArrayAgg(ParseState *pstate, JsonArrayAgg *agg)
 static Node *
 transformJsonArrayCtor(ParseState *pstate, JsonArrayCtor *ctor)
 {
-	JsonCtorExpr *jsctor;
+	JsonReturning *returning;
 	List	   *args = NIL;
 
 	/* transform element expressions, if any */
@@ -4296,13 +4324,8 @@ transformJsonArrayCtor(ParseState *pstate, JsonArrayCtor *ctor)
 		}
 	}
 
-	jsctor = makeNode(JsonCtorExpr);
-	jsctor->args = args;
-	jsctor->type = JSCTOR_JSON_ARRAY;
-	jsctor->returning = transformJsonCtorOutput(pstate, ctor->output, args);
-	jsctor->unique = false;
-	jsctor->absent_on_null = ctor->absent_on_null;
-	jsctor->location = ctor->location;
+	returning = transformJsonCtorOutput(pstate, ctor->output, args);
 
-	return coerceJsonFuncExpr(pstate, (Node *) jsctor, jsctor->returning, true);
+	return makeJsonCtorExpr(pstate, JSCTOR_JSON_ARRAY, args, NULL, returning,
+							false, ctor->absent_on_null, ctor->location);
 }
diff --git a/src/include/nodes/primnodes.h b/src/include/nodes/primnodes.h
@@ -1266,6 +1266,7 @@ typedef struct JsonCtorExpr
 	JsonCtorType type;			/* constructor type */
 	List	   *args;
 	Expr	   *func;			/* underlying json[b]_xxx() function call */
+	Expr	   *coercion;		/* coercion to RETURNING type */
 	JsonReturning *returning;	/* RETURNING clause */
 	bool		absent_on_null;	/* ABSENT ON NULL? */
 	bool		unique;			/* WITH UNIQUE KEYS? (JSON_OBJECT[AGG] only) */
