v8-0002-Add-force_array-for-COPY-TO-json-fomrat.patch

From 214ad534d13730cba13008798c3d70f8b363436f Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Tue, 23 Jan 2024 12:26:43 +0800
Subject: [PATCH v8 2/2] Add force_array for COPY TO json fomrat.

make add open brackets and close for the whole output.
separate each json row with comma after the first row.
---
 doc/src/sgml/ref/copy.sgml         | 14 ++++++++++++++
 src/backend/commands/copy.c        | 17 +++++++++++++++++
 src/backend/commands/copyto.c      | 30 ++++++++++++++++++++++++++++++
 src/backend/parser/gram.y          |  4 ++++
 src/include/commands/copy.h        |  1 +
 src/test/regress/expected/copy.out | 24 ++++++++++++++++++++++++
 src/test/regress/sql/copy.sql      | 10 ++++++++++
 7 files changed, 100 insertions(+)

diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml
index ccd90b61..d19332ac 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>'
 </synopsis>
@@ -379,6 +380,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 5d5b733d..e15056e1 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -456,6 +456,7 @@ ProcessCopyOptions(ParseState *pstate,
 	bool		freeze_specified = false;
 	bool		header_specified = false;
 	bool		on_error_specified = false;
+	bool		force_array_specified = false;
 	ListCell   *option;
 
 	/* Support external use for option sanity checking */
@@ -610,6 +611,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)
@@ -806,6 +814,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("specify COPY FORCE_ARRAY is only allowed 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 4f55d6d5..d9245df0 100644
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -88,6 +88,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;
@@ -858,6 +859,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)
@@ -905,6 +916,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,16 @@ 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/backend/parser/gram.y b/src/backend/parser/gram.y
index 702f04c3..4e13a0ab 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -3468,6 +3468,10 @@ copy_opt_item:
 				{
 					$$ = makeDefElem("encoding", (Node *) makeString($2), @1);
 				}
+			| FORCE ARRAY
+				{
+					$$ = makeDefElem("force_array", (Node *) makeBoolean(true), @1);
+				}
 		;
 
 /* The following exist for backward compatibility with very old versions */
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index f591b613..51656eec 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -72,6 +72,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 */
 	List	   *convert_select; /* list of column names (can be NIL) */
diff --git a/src/test/regress/expected/copy.out b/src/test/regress/expected/copy.out
index 0c5ade47..1b200b0d 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:  specify COPY FORCE_ARRAY is only allowed 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 96e4f0b6..a07d27af 100644
--- a/src/test/regress/sql/copy.sql
+++ b/src/test/regress/sql/copy.sql
@@ -64,6 +64,16 @@ 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 (
-- 
2.34.1