Add option force_array for COPY JSON FORMAT

This adds the force_array option, which is available exclusively
when using COPY TO with the JSON format.

When enabled, this option wraps the output in a top-level JSON array
(enclosed in square brackets with comma-separated elements), making the
entire result a valid single JSON value.  Without this option, the
default behavior is to output a stream of independent JSON objects.

Attempting to use this option with COPY FROM or with formats other than
JSON will raise an error.

Author: Joe Conway <mail@joeconway.com>
Author: jian he <jian.universality@gmail.com>
Reviewed-by: Junwang Zhao <zhjwpku@gmail.com>
Reviewed-by: Masahiko Sawada <sawada.mshk@gmail.com>
Reviewed-by: Florents Tselai <florents.tselai@gmail.com>
Reviewed-by: Andrew Dunstan <andrew@dunslane.net>
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

diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml
index f9c2717339b47004969ee1990c44f9c653436c4a..4706c9a44100c14090ba601c3ec5e77bc88d69f7 100644 (file)
--- 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> [, ...] ) | * }
@@ -382,6 +383,19 @@ COPY (SELECT j FROM (VALUES ('null'::json), (NULL::json)) v(j))
     </listitem>
    </varlistentry>
 
+   <varlistentry id="sql-copy-params-force-array">
+    <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 id="sql-copy-params-force-quote">
     <term><literal>FORCE_QUOTE</literal></term>
     <listitem>
@@ -1119,6 +1133,22 @@ COPY country TO STDOUT (DELIMITER '|');
 </programlisting>
   </para>
 
+<para>
+   When the <literal>FORCE_ARRAY</literal> option is enabled,
+   the entire output is wrapped in a single JSON array with rows separated by commas:
+<programlisting>
+COPY (SELECT * FROM (VALUES(1),(2)) val(id)) TO STDOUT  (FORMAT JSON, FORCE_ARRAY);
+</programlisting>
+The output is as follows:
+<screen>
+[
+ {"id":1}
+,{"id":2}
+]
+</screen>
+</para>
+
+
   <para>
    To copy data from a file into the <literal>country</literal> table:
 <programlisting>

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 29e22d91ecdc4541f301a050c5e2295012ca9d0b..e837f417d0d194bd05c2380fe021266fc467cc2e 100644 (file)
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -569,6 +569,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 */
@@ -725,6 +726,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)
@@ -967,6 +975,11 @@ ProcessCopyOptions(ParseState *pstate,
                                errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                                errmsg("COPY %s is not supported for %s", "FORMAT 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 be 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 489e73be7b4ac4ec9af552cd7608e347d6ffd000..faf62d959b4a37295983ea1101f85260ab9792ca 100644 (file)
--- 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 before next row */
        StringInfo      json_buf;               /* reusable buffer for JSON output,
                                                                 * initialized in BeginCopyTo */
        TupleDesc       tupDesc;                /* Descriptor for JSON output; for a column
@@ -142,6 +143,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);
@@ -183,7 +185,7 @@ static const CopyToRoutine CopyToRoutineJson = {
        .CopyToStart = CopyToTextLikeStart,
        .CopyToOutFunc = CopyToTextLikeOutFunc,
        .CopyToOneRow = CopyToJsonOneRow,
-       .CopyToEnd = CopyToTextLikeEnd,
+       .CopyToEnd = CopyToJsonEnd,
 };
 
 /* binary format */
@@ -249,6 +251,15 @@ CopyToTextLikeStart(CopyToState cstate, TupleDesc tupDesc)
 
                CopySendTextLikeEndOfRow(cstate);
        }
+
+       /*
+        * If FORCE_ARRAY has been specified, send the opening bracket.
+        */
+       if (cstate->opts.format == COPY_FORMAT_JSON && cstate->opts.force_array)
+       {
+               CopySendChar(cstate, '[');
+               CopySendTextLikeEndOfRow(cstate);
+       }
 }
 
 /*
@@ -325,13 +336,24 @@ CopyToTextLikeOneRow(CopyToState cstate,
        CopySendTextLikeEndOfRow(cstate);
 }
 
-/* Implementation of the end callback for text, CSV, and json formats */
+/* Implementation of the end callback for text and CSV formats */
 static void
 CopyToTextLikeEnd(CopyToState cstate)
 {
        /* Nothing to do here */
 }
 
+/* Implementation of the end callback for json format */
+static void
+CopyToJsonEnd(CopyToState cstate)
+{
+       if (cstate->opts.force_array)
+       {
+               CopySendChar(cstate, ']');
+               CopySendTextLikeEndOfRow(cstate);
+       }
+}
+
 /* Implementation of per-row callback for json format */
 static void
 CopyToJsonOneRow(CopyToState cstate, TupleTableSlot *slot)
@@ -393,6 +415,18 @@ CopyToJsonOneRow(CopyToState cstate, TupleTableSlot *slot)
 
        composite_to_json(rowdata, cstate->json_buf, false);
 
+       if (cstate->opts.force_array)
+       {
+               if (cstate->json_row_delim_needed)
+                       CopySendChar(cstate, ',');
+               else
+               {
+                       /* first row needs no delimiter */
+                       CopySendChar(cstate, ' ');
+                       cstate->json_row_delim_needed = true;
+               }
+       }
+
        CopySendData(cstate, cstate->json_buf->data, cstate->json_buf->len);
 
        CopySendTextLikeEndOfRow(cstate);

diff --git a/src/bin/psql/tab-complete.in.c b/src/bin/psql/tab-complete.in.c
index b6cbb077326de45431fe781212a58779ef472c11..e3bdbf34a5c6c1a6f8c357c1f23449b1d272e033 100644 (file)
--- a/src/bin/psql/tab-complete.in.c
+++ b/src/bin/psql/tab-complete.in.c
@@ -1240,7 +1240,7 @@ Copy_common_options, "DEFAULT", "FORCE_NOT_NULL", "FORCE_NULL", "FREEZE", \
 
 /* COPY TO options */
 #define Copy_to_options \
-Copy_common_options, "FORCE_QUOTE"
+Copy_common_options, "FORCE_QUOTE", "FORCE_ARRAY"
 
 /*
  * These object types were introduced later than our support cutoff of

diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index 2b5bef6738e843a5a302b0fa0f48ff87553d9f6e..abecfe510981ed56c6d1ad868773adb5a3decd13 100644 (file)
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -88,6 +88,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 7f2d2e065f67076947405441c959f6911e9348fb..1714faab39c2fd6a94a0d1e7c4af8247de2ff078 100644 (file)
--- a/src/test/regress/expected/copy.out
+++ b/src/test/regress/expected/copy.out
@@ -83,6 +83,16 @@ copy (select 1 as foo union all select 2) to stdout with (format json);
 copy (values (1), (2)) TO stdout with (format json);
 {"column1":1}
 {"column1":2}
+copy (select 1 union all select 2) to stdout with (format json, force_array true);
+[
+ {"?column?":1}
+,{"?column?":2}
+]
+copy (values (1), (2)) TO stdout with (format json, force_array true);
+[
+ {"column1":1}
+,{"column1":2}
+]
 copy copytest to stdout json;
 {"style":"DOS","test":"abc\r\ndef","filler":1}
 {"style":"Unix","test":"abc\ndef","filler":2}
@@ -134,6 +144,33 @@ copy copytest (style, test, filler) 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}
+-- should fail: force_array requires json format
+copy copytest to stdout (format csv, force_array true);
+ERROR:  COPY FORCE_ARRAY can only be used with JSON mode
+-- force_array variants
+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(style, test) to stdout (format json, force_array true);
+[
+ {"style":"DOS","test":"abc\r\ndef"}
+,{"style":"Unix","test":"abc\ndef"}
+,{"style":"Mac","test":"abc\rdef"}
+,{"style":"esc\\ape","test":"a\\r\\\r\\\n\\nb"}
+]
+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}
+-- force_array with empty result set
+copy (select 1 where false) to stdout (format json, force_array);
+[
+]
 -- column list with diverse data types
 create temp table copyjsontest_types (
     id int,

diff --git a/src/test/regress/sql/copy.sql b/src/test/regress/sql/copy.sql
index 404f432108574356a5d1c4beb8dadae8cf2f6534..eaad290b257bc9a06e202926632e69ac4bdb7030 100644 (file)
--- a/src/test/regress/sql/copy.sql
+++ b/src/test/regress/sql/copy.sql
@@ -86,6 +86,8 @@ copy copytest3 to stdout csv header;
 copy (select 1 union all select 2) to stdout with (format json);
 copy (select 1 as foo union all select 2) to stdout with (format json);
 copy (values (1), (2)) TO stdout with (format json);
+copy (select 1 union all select 2) to stdout with (format json, force_array true);
+copy (values (1), (2)) TO stdout with (format json, force_array true);
 copy copytest to stdout json;
 copy copytest to stdout (format json);
 copy (select * from copytest) to stdout (format json);
@@ -109,6 +111,17 @@ copy copytest from stdin(format json);
 -- column list with json format
 copy copytest (style, test, filler) to stdout (format json);
 
+-- should fail: force_array requires json format
+copy copytest to stdout (format csv, force_array true);
+
+-- force_array variants
+copy copytest to stdout (format json, force_array);
+copy copytest(style, test) to stdout (format json, force_array true);
+copy copytest to stdout (format json, force_array false);
+
+-- force_array with empty result set
+copy (select 1 where false) to stdout (format json, force_array);
+
 -- column list with diverse data types
 create temp table copyjsontest_types (
     id int,