Refactor CopyAttributeOut{CSV,Text}() to use a callback in COPY TO
authorMichael Paquier <michael@paquier.xyz>
Mon, 5 Feb 2024 02:12:37 +0000 (11:12 +0900)
committerMichael Paquier <michael@paquier.xyz>
Mon, 5 Feb 2024 02:12:37 +0000 (11:12 +0900)
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

index d3dc3fc854f1437d222e581176e20aaed1dd456e..52b42f8a522ef902349ad04c518215d23bc4c69c 100644 (file)
@@ -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)