v10-0002-Add-option-force_array-for-COPY-TO-JSON-fomrat.patch

From c939262f6266cea38761aa89cc6bb495b8d38ccb Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Mon, 1 Apr 2024 19:53:10 +0800
Subject: [PATCH v10 2/2] Add option force_array for COPY TO JSON fomrat

force_array option only apply to COPY TO operation with JSON format.
make add opening brackets and close brackets for the whole json output.
also, separate each json record with a comma.
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        | 17 +++++++++++++++++
 src/backend/commands/copyto.c      | 29 +++++++++++++++++++++++++++++
 src/include/commands/copy.h        |  1 +
 src/test/regress/expected/copy.out | 24 ++++++++++++++++++++++++
 src/test/regress/sql/copy.sql      |  9 +++++++++
 6 files changed, 94 insertions(+)

diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml
index add84dbb..f4100cba 100644
--- a/doc/src/sgml/ref/copy.sgml
+++ b/doc/src/sgml/ref/copy.sgml
@@ -43,6 +43,7 @@ COPY { <replaceable class="parameter">table_name</replaceable> [ ( <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> [, ...] ) | * }
+    FORCE_ARRAY [ <replaceable class="parameter">boolean</replaceable> ]
     ON_ERROR '<replaceable class="parameter">error_action</replaceable>'
     ENCODING '<replaceable class="parameter">encoding_name</replaceable>'
     LOG_VERBOSITY <replaceable class="parameter">mode</replaceable>
@@ -390,6 +391,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>ON_ERROR</literal></term>
     <listitem>
diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 02f16d9e..2ffd6978 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -473,6 +473,7 @@ ProcessCopyOptions(ParseState *pstate,
 	bool		header_specified = false;
 	bool		on_error_specified = false;
 	bool		log_verbosity_specified = false;
+	bool		force_array_specified = false;
 	ListCell   *option;
 
 	/* Support external use for option sanity checking */
@@ -627,6 +628,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)
@@ -830,6 +838,15 @@ ProcessCopyOptions(ParseState *pstate,
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("cannot use JSON mode in COPY FROM")));
 
+	if (!opts_out->json_mode && opts_out->force_array)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("COPY FORCE_ARRAY requires JSON mode")));
+	if (!opts_out->json_mode && force_array_specified)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("COPY FORCE_ARRAY only available in JSON mode")));
+
 	if (opts_out->default_print)
 	{
 		if (!is_from)
diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c
index fe2eb244..8e4edb6e 100644
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -81,6 +81,7 @@ 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? */
+	bool		json_row_delim_needed;	/* need delimiter to start next json array element */
 	copy_data_dest_cb data_dest_cb; /* function for writing data */
 
 	CopyFormatOptions opts;
@@ -853,6 +854,16 @@ DoCopyTo(CopyToState cstate)
 
 			CopySendEndOfRow(cstate);
 		}
+		/*
+		 * If JSON has been requested, and FORCE_ARRAY has been specified send
+		 * the opening bracket.
+		*/
+		if (cstate->opts.json_mode && cstate->opts.force_array)
+		{
+			CopySendChar(cstate, '[');
+			CopySendEndOfRow(cstate);
+		}
+
 	}
 
 	if (cstate->rel)
@@ -900,6 +911,15 @@ DoCopyTo(CopyToState cstate)
 		CopySendEndOfRow(cstate);
 	}
 
+	/*
+	 * If JSON has been requested, and FORCE_ARRAY has been specified send the
+	 * closing bracket.
+	*/
+	if (cstate->opts.json_mode && cstate->opts.force_array)
+	{
+		CopySendChar(cstate, ']');
+		CopySendEndOfRow(cstate);
+	}
 	MemoryContextDelete(cstate->rowcontext);
 
 	if (fe_copy)
@@ -1006,6 +1026,15 @@ CopyOneRowTo(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);
 	}
 
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index ff6ecc7a..f76fe8fa 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -81,6 +81,7 @@ typedef struct CopyFormatOptions
 	List	   *force_null;		/* list of column names */
 	bool		force_null_all; /* FORCE_NULL *? */
 	bool	   *force_null_flags;	/* per-column CSV FN flags */
+	bool		force_array;	/* add JSON array decorations */
 	bool		convert_selectively;	/* do selective binary conversion? */
 	CopyOnErrorChoice on_error; /* what to do when error happened */
 	CopyLogVerbosityChoice log_verbosity;	/* verbosity of logged messages */
diff --git a/src/test/regress/expected/copy.out b/src/test/regress/expected/copy.out
index 0c5ade47..7812768c 100644
--- a/src/test/regress/expected/copy.out
+++ b/src/test/regress/expected/copy.out
@@ -59,6 +59,30 @@ ERROR:  cannot specify HEADER in JSON mode
 -- Error
 copy copytest from stdout (format json);
 ERROR:  cannot use JSON mode in COPY FROM
+--Error
+copy copytest to stdout (format csv, force_array false);
+ERROR:  COPY FORCE_ARRAY only available in JSON mode
+copy copytest from stdin (format json, force_array true);
+ERROR:  cannot use JSON mode in COPY FROM
+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 da6b0b0a..f685193b 100644
--- a/src/test/regress/sql/copy.sql
+++ b/src/test/regress/sql/copy.sql
@@ -64,6 +64,15 @@ copy copytest to stdout (format json, header);
 -- Error
 copy copytest from stdout (format json);
 
+--Error
+copy copytest to stdout (format csv, force_array false);
+copy copytest from stdin (format json, force_array true);
+
+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