From 2889fd23be565155c8cd96a8c5a07ac971261515 Mon Sep 17 00:00:00 2001 From: Michael Paquier Date: Mon, 5 Feb 2024 11:12:37 +0900 Subject: [PATCH] Refactor CopyAttributeOut{CSV,Text}() to use a callback in COPY TO These routines are used by the text and CSV formats to send an output representation of a string, applying quotes if required. This is similar to 95fb5b49024a, reducing the number of "if" branches that need to be checked on a per-row basis when sending representation of fields in text or CSV mode. While on it, this simplifies the signature of CopyAttributeOutCSV() as it is possible to know that an attribute is alone on a line thanks to CopyToState. Headers should not use quotes, even if forced at query level. Extracted from a larger patch by the same author. Author: Sutou Kouhei Discussion: https://postgr.es/m/20231204.153548.2126325458835528809.kou@clear-code.com --- src/backend/commands/copyto.c | 47 ++++++++++++++++++++++++----------- 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c index d3dc3fc854..52b42f8a52 100644 --- a/src/backend/commands/copyto.c +++ b/src/backend/commands/copyto.c @@ -54,6 +54,14 @@ typedef enum CopyDest COPY_CALLBACK, /* to callback function */ } CopyDest; +/* + * Per-format callback to send output representation of one attribute for + * a `string`. `use_quote` tracks if quotes are required in the output + * representation. + */ +typedef void (*CopyAttributeOut) (CopyToState cstate, const char *string, + bool use_quote); + /* * This struct contains all the state variables used throughout a COPY TO * operation. @@ -97,6 +105,7 @@ typedef struct CopyToStateData MemoryContext copycontext; /* per-copy execution context */ FmgrInfo *out_functions; /* lookup info for output functions */ + CopyAttributeOut copy_attribute_out; /* output representation callback */ MemoryContext rowcontext; /* per-row evaluation context */ uint64 bytes_processed; /* number of bytes processed so far */ } CopyToStateData; @@ -117,9 +126,12 @@ static const char BinarySignature[11] = "PGCOPY\n\377\r\n\0"; static void EndCopy(CopyToState cstate); static void ClosePipeToProgram(CopyToState cstate); static void CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot); -static void CopyAttributeOutText(CopyToState cstate, const char *string); + +/* Callbacks for copy_attribute_out */ +static void CopyAttributeOutText(CopyToState cstate, const char *string, + bool use_quote); static void CopyAttributeOutCSV(CopyToState cstate, const char *string, - bool use_quote, bool single_attr); + bool use_quote); /* Low-level communications functions */ static void SendCopyBegin(CopyToState cstate); @@ -433,6 +445,15 @@ BeginCopyTo(ParseState *pstate, /* Extract options from the statement node tree */ ProcessCopyOptions(pstate, &cstate->opts, false /* is_from */ , options); + /* Set output representation callback */ + if (!cstate->opts.binary) + { + if (cstate->opts.csv_mode) + cstate->copy_attribute_out = CopyAttributeOutCSV; + else + cstate->copy_attribute_out = CopyAttributeOutText; + } + /* Process the source/target relation or query */ if (rel) { @@ -836,11 +857,8 @@ DoCopyTo(CopyToState cstate) colname = NameStr(TupleDescAttr(tupDesc, attnum - 1)->attname); - if (cstate->opts.csv_mode) - CopyAttributeOutCSV(cstate, colname, false, - list_length(cstate->attnumlist) == 1); - else - CopyAttributeOutText(cstate, colname); + /* Ignore quotes */ + cstate->copy_attribute_out(cstate, colname, false); } CopySendEndOfRow(cstate); @@ -950,12 +968,9 @@ CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot) { string = OutputFunctionCall(&out_functions[attnum - 1], value); - if (cstate->opts.csv_mode) - CopyAttributeOutCSV(cstate, string, - cstate->opts.force_quote_flags[attnum - 1], - list_length(cstate->attnumlist) == 1); - else - CopyAttributeOutText(cstate, string); + + cstate->copy_attribute_out(cstate, string, + cstate->opts.force_quote_flags[attnum - 1]); } else { @@ -985,7 +1000,8 @@ CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot) } while (0) static void -CopyAttributeOutText(CopyToState cstate, const char *string) +CopyAttributeOutText(CopyToState cstate, const char *string, + bool use_quote) { const char *ptr; const char *start; @@ -1139,7 +1155,7 @@ CopyAttributeOutText(CopyToState cstate, const char *string) */ static void CopyAttributeOutCSV(CopyToState cstate, const char *string, - bool use_quote, bool single_attr) + bool use_quote) { const char *ptr; const char *start; @@ -1147,6 +1163,7 @@ CopyAttributeOutCSV(CopyToState cstate, const char *string, char delimc = cstate->opts.delim[0]; char quotec = cstate->opts.quote[0]; char escapec = cstate->opts.escape[0]; + bool single_attr = (list_length(cstate->attnumlist) == 1); /* force quoting if it matches null_print (before conversion!) */ if (!use_quote && strcmp(string, cstate->opts.null_print) == 0)