v17-0002-Add-option-force_array-for-COPY-JSON-FORMAT.patch

From 44c494fd8d7fdb9d8fd5d2d2a48f49b779d1bcb9 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Fri, 4 Jul 2025 13:25:00 +0800
Subject: [PATCH v17 2/2] Add option force_array for COPY JSON FORMAT

force_array option can only be used in COPY TO with JSON format.  it make the
output json output behave like json array type.  refactored by Junwang Zhao to
adapt the newly introduced CopyToRoutine struct(2e4127b6d2).

Author: Joe Conway <mail@joeconway.com>
discussion: https://postgr.es/m/CALvfUkBxTYy5uWPFVwpk_7ii2zgT07t3d-yR_cy4sfrrLU%3Dkcg%40mail.gmail.com
discussion: https://postgr.es/m/6a04628d-0d53-41d9-9e35-5a8dc302c34c@joeconway.com
---
 doc/src/sgml/ref/copy.sgml         | 14 +++++++++++
 src/backend/commands/copy.c        | 13 +++++++++++
 src/backend/commands/copyto.c      | 37 +++++++++++++++++++++++++++++-
 src/bin/psql/tab-complete.in.c     |  2 +-
 src/include/commands/copy.h        |  1 +
 src/test/regress/expected/copy.out | 23 +++++++++++++++++++
 src/test/regress/sql/copy.sql      |  8 +++++++
 7 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml
index 219604ad306..c01927864bd 100644
--- a/doc/src/sgml/ref/copy.sgml
+++ b/doc/src/sgml/ref/copy.sgml
@@ -40,6 +40,7 @@ COPY { <replaceable class="parameter">table_name</replaceable> [ ( <replaceable
     HEADER [ <replaceable class="parameter">boolean</replaceable> | <replaceable class="parameter">integer</replaceable> | MATCH ]
     QUOTE '<replaceable class="parameter">quote_character</replaceable>'
     ESCAPE '<replaceable class="parameter">escape_character</replaceable>'
+    FORCE_ARRAY [ <replaceable class="parameter">boolean</replaceable> ]
     FORCE_QUOTE { ( <replaceable class="parameter">column_name</replaceable> [, ...] ) | * }
     FORCE_NOT_NULL { ( <replaceable class="parameter">column_name</replaceable> [, ...] ) | * }
     FORCE_NULL { ( <replaceable class="parameter">column_name</replaceable> [, ...] ) | * }
@@ -366,6 +367,19 @@ COPY { <replaceable class="parameter">table_name</replaceable> [ ( <replaceable
     </listitem>
    </varlistentry>
 
+   <varlistentry>
+    <term><literal>FORCE_ARRAY</literal></term>
+    <listitem>
+     <para>
+      Force output of square brackets as array decorations at the beginning
+      and end of output, and commas between the rows. It is allowed only in
+      <command>COPY TO</command>, and only when using
+      <literal>json</literal> format. The default is
+      <literal>false</literal>.
+     </para>
+    </listitem>
+   </varlistentry>
+
    <varlistentry>
     <term><literal>FORCE_QUOTE</literal></term>
     <listitem>
diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 213e59cc435..d23ef99e395 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -514,6 +514,7 @@ ProcessCopyOptions(ParseState *pstate,
 	bool		on_error_specified = false;
 	bool		log_verbosity_specified = false;
 	bool		reject_limit_specified = false;
+	bool		force_array_specified = false;
 	ListCell   *option;
 
 	/* Support external use for option sanity checking */
@@ -670,6 +671,13 @@ ProcessCopyOptions(ParseState *pstate,
 								defel->defname),
 						 parser_errposition(pstate, defel->location)));
 		}
+		else if (strcmp(defel->defname, "force_array") == 0)
+		{
+			if (force_array_specified)
+				errorConflictingDefElem(defel, pstate);
+			force_array_specified = true;
+			opts_out->force_array = defGetBoolean(defel);
+		}
 		else if (strcmp(defel->defname, "on_error") == 0)
 		{
 			if (on_error_specified)
@@ -918,6 +926,11 @@ ProcessCopyOptions(ParseState *pstate,
 				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 				errmsg("COPY %s mode cannot be used with %s", "json", "COPY FROM"));
 
+	if (opts_out->format != COPY_FORMAT_JSON && opts_out->force_array)
+		ereport(ERROR,
+				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("COPY %s can only used with JSON mode", "FORCE_ARRAY"));
+
 	if (opts_out->default_print)
 	{
 		if (!is_from)
diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c
index 34b72936bca..18061661fb1 100644
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -84,6 +84,10 @@ typedef struct CopyToStateData
 	List	   *attnumlist;		/* integer list of attnums to copy */
 	char	   *filename;		/* filename, or NULL for STDOUT */
 	bool		is_program;		/* is 'filename' a program to popen? */
+
+	/* need delimiter to start next json array element */
+	bool		json_row_delim_needed;
+
 	copy_data_dest_cb data_dest_cb; /* function for writing data */
 
 	CopyFormatOptions opts;
@@ -128,6 +132,7 @@ static void CopyToTextLikeOneRow(CopyToState cstate, TupleTableSlot *slot,
 								 bool is_csv);
 static void CopyToTextLikeEnd(CopyToState cstate);
 static void CopyToJsonOneRow(CopyToState cstate, TupleTableSlot *slot);
+static void CopyToJsonEnd(CopyToState cstate);
 static void CopyToBinaryStart(CopyToState cstate, TupleDesc tupDesc);
 static void CopyToBinaryOutFunc(CopyToState cstate, Oid atttypid, FmgrInfo *finfo);
 static void CopyToBinaryOneRow(CopyToState cstate, TupleTableSlot *slot);
@@ -172,7 +177,7 @@ static const CopyToRoutine CopyToRoutineJson = {
 	.CopyToStart = CopyToTextLikeStart,
 	.CopyToOutFunc = CopyToTextLikeOutFunc,
 	.CopyToOneRow = CopyToJsonOneRow,
-	.CopyToEnd = CopyToTextLikeEnd,
+	.CopyToEnd = CopyToJsonEnd,
 };
 
 /* binary format */
@@ -238,6 +243,16 @@ CopyToTextLikeStart(CopyToState cstate, TupleDesc tupDesc)
 
 		CopySendTextLikeEndOfRow(cstate);
 	}
+
+	/*
+	 * If JSON has been requested, and FORCE_ARRAY has been specified send the
+	 * opening bracket.
+	 */
+	if (cstate->opts.format == COPY_FORMAT_JSON && cstate->opts.force_array)
+	{
+		CopySendChar(cstate, '[');
+		CopySendTextLikeEndOfRow(cstate);
+	}
 }
 
 /*
@@ -349,11 +364,31 @@ CopyToJsonOneRow(CopyToState cstate, TupleTableSlot *slot)
 	result = makeStringInfo();
 	composite_to_json(rowdata, result, false);
 
+	if (cstate->json_row_delim_needed && cstate->opts.force_array)
+		CopySendChar(cstate, ',');
+	else if (cstate->opts.force_array)
+	{
+		/* first row needs no delimiter */
+		CopySendChar(cstate, ' ');
+		cstate->json_row_delim_needed = true;
+	}
+
 	CopySendData(cstate, result->data, result->len);
 
 	CopySendTextLikeEndOfRow(cstate);
 }
 
+/* Implementation of the end callback for json format */
+static void
+CopyToJsonEnd(CopyToState cstate)
+{
+	if (cstate->opts.force_array)
+	{
+		CopySendChar(cstate, ']');
+		CopySendTextLikeEndOfRow(cstate);
+	}
+}
+
 /*
  * Implementation of the start callback for binary format. Send a header
  * for a binary copy.
diff --git a/src/bin/psql/tab-complete.in.c b/src/bin/psql/tab-complete.in.c
index bd4c03be050..55313612398 100644
--- a/src/bin/psql/tab-complete.in.c
+++ b/src/bin/psql/tab-complete.in.c
@@ -3303,7 +3303,7 @@ match_previous_words(int pattern_id,
 	else if (Matches("COPY|\\copy", MatchAny, "FROM|TO", MatchAny, "WITH", "("))
 		COMPLETE_WITH("FORMAT", "FREEZE", "DELIMITER", "NULL",
 					  "HEADER", "QUOTE", "ESCAPE", "FORCE_QUOTE",
-					  "FORCE_NOT_NULL", "FORCE_NULL", "ENCODING", "DEFAULT",
+					  "FORCE_NOT_NULL", "FORCE_NULL", "FORCE_ARRAY", "ENCODING", "DEFAULT",
 					  "ON_ERROR", "LOG_VERBOSITY", "REJECT_LIMIT");
 
 	/* Complete COPY <sth> FROM|TO filename WITH (FORMAT */
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index 85aedc267d6..7274b0d3ca5 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -87,6 +87,7 @@ typedef struct CopyFormatOptions
 	List	   *force_notnull;	/* list of column names */
 	bool		force_notnull_all;	/* FORCE_NOT_NULL *? */
 	bool	   *force_notnull_flags;	/* per-column CSV FNN flags */
+	bool		force_array;	/* add JSON array decorations */
 	List	   *force_null;		/* list of column names */
 	bool		force_null_all; /* FORCE_NULL *? */
 	bool	   *force_null_flags;	/* per-column CSV FN flags */
diff --git a/src/test/regress/expected/copy.out b/src/test/regress/expected/copy.out
index fcb8823e101..f3196fd5609 100644
--- a/src/test/regress/expected/copy.out
+++ b/src/test/regress/expected/copy.out
@@ -110,6 +110,29 @@ LINE 1: copy copytest to stdout (format json, on_error ignore);
 copy copytest from stdin(format json);
 ERROR:  COPY json mode cannot be used with COPY FROM
 -- all of the above should yield error
+--Error
+copy copytest to stdout (format csv, force_array true);
+ERROR:  COPY FORCE_ARRAY can only used with JSON mode
+--ok
+copy copytest to stdout (format json, force_array);
+[
+ {"style":"DOS","test":"abc\r\ndef","filler":1}
+,{"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}
+]
+copy copytest to stdout (format json, force_array true);
+[
+ {"style":"DOS","test":"abc\r\ndef","filler":1}
+,{"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}
+]
+copy copytest to stdout (format json, force_array false);
+{"style":"DOS","test":"abc\r\ndef","filler":1}
+{"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}
 -- embedded escaped characters
 create temp table copyjsontest (
     id bigserial,
diff --git a/src/test/regress/sql/copy.sql b/src/test/regress/sql/copy.sql
index 80bd4a59239..c55aa08e99d 100644
--- a/src/test/regress/sql/copy.sql
+++ b/src/test/regress/sql/copy.sql
@@ -100,6 +100,14 @@ copy copytest to stdout (format json, on_error ignore);
 copy copytest from stdin(format json);
 -- all of the above should yield error
 
+--Error
+copy copytest to stdout (format csv, force_array true);
+
+--ok
+copy copytest to stdout (format json, force_array);
+copy copytest to stdout (format json, force_array true);
+copy copytest to stdout (format json, force_array false);
+
 -- embedded escaped characters
 create temp table copyjsontest (
     id bigserial,
-- 
2.34.1