COPY TO FORMAT JSON: use the query's column names as keys

CopyToJsonOneRow() stamped the row datum with the query's descriptor
only when the slot's descriptor was of type RECORD.  A scan that does
not project returns the table's row type, so the keys came from the
table, not the query, losing column aliases and mislabelling rows
from UNION ALL branches with different column orders.

Fix by always using the query's descriptor.  For non-virtual slots,
heap_copy_tuple_as_datum() can just be passed that descriptor, so
there is no extra cost.  This is safe because a scan skips projection
only when tlist_matches_tupdesc() holds, so the layouts match.

Bug introduced in 7dadd38cda.

Reported-by: Shihao Zhong <zhong950419@gmail.com>
Reviewed-by: Shihao Zhong <zhong950419@gmail.com>
Reviewed-by: João Marcelo Detomini <joao.detomini@enterprisedb.com>
Discussion: https://postgr.es/m/CAGRkXqQHjxgufHPz86+dxGaQZt37hikduUArYtfbn2_3mryZAw@mail.gmail.com
Backpatch-through: 19

src/backend/commands/copyto.c

@@ -395,22 +395,38 @@ CopyToJsonOneRow(CopyToState cstate, TupleTableSlot *slot)
 	else
 	{
 		/*
-		 * Full table or query without column list.  For queries, the slot's
-		 * TupleDesc may carry RECORDOID, which is not registered in the type
-		 * cache and would cause composite_to_json's lookup_rowtype_tupdesc
-		 * call to fail.  Build a HeapTuple stamped with the blessed
-		 * descriptor so the type can be looked up correctly.
+		 * Full table or query without column list.  For a query, the slot's
+		 * descriptor is either an unregistered RECORD type, which
+		 * composite_to_json's lookup_rowtype_tupdesc() cannot look up, or,
+		 * when the top plan node does not project, the row type of a scanned
+		 * table, whose column names need not match the query's.  Either way,
+		 * the datum must be stamped with the query's blessed descriptor.
+		 *
+		 * A virtual slot has no physical tuple, so form one directly.
+		 * Otherwise copy the slot's tuple and stamp the copy with the query's
+		 * descriptor.  That is safe because the tuple's physical layout
+		 * matches the query's result descriptor: a scan returns its scan
+		 * tuple unprojected only if tlist_matches_tupdesc() holds, which
+		 * rules out dropped columns and columns with missing values.
 		 */
-		if (!cstate->rel && slot->tts_tupleDescriptor->tdtypeid == RECORDOID)
-		{
-			HeapTuple	tup = heap_form_tuple(cstate->tupDesc,
-											  slot->tts_values,
-											  slot->tts_isnull);
-
-			rowdata = HeapTupleGetDatum(tup);
-		}
-		else
+		if (cstate->rel)
 			rowdata = ExecFetchSlotHeapTupleDatum(slot);
+		else if (TTS_IS_VIRTUAL(slot))
+			rowdata = HeapTupleGetDatum(heap_form_tuple(cstate->tupDesc,
+														slot->tts_values,
+														slot->tts_isnull));
+		else
+		{
+			HeapTuple	tup;
+			bool		shouldFree;
+
+			Assert(slot->tts_tupleDescriptor->natts == cstate->tupDesc->natts);
+
+			tup = ExecFetchSlotHeapTuple(slot, false, &shouldFree);
+			rowdata = heap_copy_tuple_as_datum(tup, cstate->tupDesc);
+			if (shouldFree)
+				heap_freetuple(tup);
+		}
 	}
 
 	composite_to_json(rowdata, cstate->json_buf, false);

src/test/regress/expected/copy.out

@@ -117,6 +117,19 @@ copy (select * from copytest) to stdout (format json);
 {"style":"Unix","test":"abc\ndef","filler":2}
 {"style":"Mac","test":"abc\rdef","filler":3}
 {"style":"esc\\ape","test":"a\\r\\\r\\\n\\nb","filler":4}
+-- keys are the query's column names, also when a scan does not project
+copy (select a as x, b as y, c as z from copytest_order) to stdout (format json);
+{"x":1,"y":2,"z":3}
+create temp table copytest_order2 (c int, b int, a int);
+insert into copytest_order2 values (4, 5, 6);
+copy (select * from copytest_order union all select * from copytest_order2)
+  to stdout (format json);
+{"a":1,"b":2,"c":3}
+{"a":4,"b":5,"c":6}
+-- columns with missing values force a projection
+alter table copytest_order2 add column d int default 7;
+copy (select * from copytest_order2) to stdout (format json);
+{"c":4,"b":5,"a":6,"d":7}
 -- all of the following should yield error
 copy copytest to stdout (format json, delimiter '|');
 ERROR:  cannot specify DELIMITER in JSON mode

src/test/regress/sql/copy.sql

@@ -101,6 +101,16 @@ copy copytest to stdout json;
 copy copytest to stdout (format json);
 copy (select * from copytest) to stdout (format json);
 
+-- keys are the query's column names, also when a scan does not project
+copy (select a as x, b as y, c as z from copytest_order) to stdout (format json);
+create temp table copytest_order2 (c int, b int, a int);
+insert into copytest_order2 values (4, 5, 6);
+copy (select * from copytest_order union all select * from copytest_order2)
+  to stdout (format json);
+-- columns with missing values force a projection
+alter table copytest_order2 add column d int default 7;
+copy (select * from copytest_order2) to stdout (format json);
+
 -- all of the following should yield error
 copy copytest to stdout (format json, delimiter '|');
 copy copytest to stdout (format json, null '\N');