From aee1dea1a0ab5d4d058a1cc8766667c78f77d175 Mon Sep 17 00:00:00 2001 From: Junwang Zhao Date: Sun, 2 Mar 2025 04:41:48 +0000 Subject: [PATCH v15 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 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 | 38 +++++++++++++++++++++++++++--- 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, 95 insertions(+), 4 deletions(-) diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml index 9c519d8a9e2..7b3c913d4ee 100644 --- a/doc/src/sgml/ref/copy.sgml +++ b/doc/src/sgml/ref/copy.sgml @@ -43,6 +43,7 @@ COPY { table_name [ ( column_name [, ...] ) | * } FORCE_NOT_NULL { ( column_name [, ...] ) | * } FORCE_NULL { ( column_name [, ...] ) | * } + FORCE_ARRAY [ boolean ] ON_ERROR error_action REJECT_LIMIT maxerror ENCODING 'encoding_name' @@ -392,6 +393,19 @@ COPY { table_name [ ( + + FORCE_ARRAY + + + 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 + COPY TO, and only when using + JSON format. The default is + false. + + + + ON_ERROR diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c index b6f74c798d0..7b4c64ea97e 100644 --- a/src/backend/commands/copy.c +++ b/src/backend/commands/copy.c @@ -504,6 +504,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 */ @@ -658,6 +659,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) @@ -906,6 +914,11 @@ ProcessCopyOptions(ParseState *pstate, errcode(ERRCODE_INVALID_PARAMETER_VALUE), errmsg("COPY json mode cannot be used with %s", "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 69b7fd11e03..86842b764e8 100644 --- a/src/backend/commands/copyto.c +++ b/src/backend/commands/copyto.c @@ -84,6 +84,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; @@ -128,6 +129,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); @@ -171,7 +173,7 @@ static const CopyToRoutine CopyToRoutineJson = { .CopyToStart = CopyToTextLikeStart, .CopyToOutFunc = CopyToTextLikeOutFunc, .CopyToOneRow = CopyToJsonOneRow, - .CopyToEnd = CopyToTextLikeEnd, + .CopyToEnd = CopyToJsonEnd, }; /* binary format */ @@ -235,6 +237,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); + } } /* @@ -345,11 +357,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. @@ -554,8 +586,8 @@ CopySendEndOfRow(CopyToState cstate) } /* - * Wrapper function of CopySendEndOfRow for text and CSV formats. Sends the - * line termination and do common appropriate things for the end of row. + * Wrapper function of CopySendEndOfRow for text, CSV and json formats. Sends + * the line termination and do common appropriate things for the end of row. */ static inline void CopySendTextLikeEndOfRow(CopyToState cstate) diff --git a/src/bin/psql/tab-complete.in.c b/src/bin/psql/tab-complete.in.c index 6069b118834..e0c7412ec0a 100644 --- a/src/bin/psql/tab-complete.in.c +++ b/src/bin/psql/tab-complete.in.c @@ -3278,7 +3278,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"); /* Complete COPY FROM|TO filename WITH (FORMAT */ diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h index 07fcc2bc9ac..fa8a8ab7e31 100644 --- a/src/include/commands/copy.h +++ b/src/include/commands/copy.h @@ -89,6 +89,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 d9e30937feb..7aa0005fe7a 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 5e4d1f0781a..3e062275d85 100644 --- a/src/test/regress/sql/copy.sql +++ b/src/test/regress/sql/copy.sql @@ -101,6 +101,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.39.5