Add support for implementing custom COPY handler as extension

* TO: Add CopyToStateData::opaque that can be used to keep
  data for custom COPY TO handler implementation
* TO: Export CopySendEndOfRow() to send end of row data as
  CopyToStateFlush()
* FROM: Add CopyFromStateData::opaque that can be used to
  keep data for custom COPY FROM handler implementation
* FROM: Export CopyGetData() to get the next data as
  CopyFromStateGetData()
* FROM: Add CopyFromSkipErrorRow() for "ON_ERROR stop" and
  "LOG_VERBOSITY verbose"

COPY FROM extensions must call CopyFromSkipErrorRow() when
CopyFromOneRow callback reports an error by
errsave(). CopyFromSkipErrorRow() handles "ON_ERROR stop" and
"LOG_VERBOSITY verbose" cases.

Author: kou

src/backend/commands/copyfromparse.c

@@ -739,6 +739,17 @@ CopyReadBinaryData(CopyFromState cstate, char *dest, int nbytes)
 	return copied_bytes;
 }
 
+/*
+ * Export CopyGetData() for extensions. We want to keep CopyGetData() as a
+ * static function for optimization. CopyGetData() calls in this file may be
+ * optimized by a compiler.
+ */
+int
+CopyFromStateGetData(CopyFromState cstate, void *dest, int minread, int maxread)
+{
+	return CopyGetData(cstate, dest, minread, maxread);
+}
+
 /*
  * This function is exposed for use by extensions that read raw fields in the
  * next line. See NextCopyFromRawFieldsInternal() for details.
@@ -927,6 +938,51 @@ CopyFromCSVOneRow(CopyFromState cstate, ExprContext *econtext, Datum *values,
 	return CopyFromTextLikeOneRow(cstate, econtext, values, nulls, true);
 }
 
+/*
+ * Call this when you report an error by errsave() in your CopyFromOneRow
+ * callback. This handles "ON_ERROR stop" and "LOG_VERBOSITY verbose" cases
+ * for you.
+ */
+void
+CopyFromSkipErrorRow(CopyFromState cstate)
+{
+	Assert(cstate->opts.on_error != COPY_ON_ERROR_STOP);
+
+	cstate->num_errors++;
+
+	if (cstate->opts.log_verbosity == COPY_LOG_VERBOSITY_VERBOSE)
+	{
+		/*
+		 * Since we emit line number and column info in the below notice
+		 * message, we suppress error context information other than the
+		 * relation name.
+		 */
+		Assert(!cstate->relname_only);
+		cstate->relname_only = true;
+
+		if (cstate->cur_attval)
+		{
+			char	   *attval;
+
+			attval = CopyLimitPrintoutLength(cstate->cur_attval);
+			ereport(NOTICE,
+					errmsg("skipping row due to data type incompatibility at line %llu for column \"%s\": \"%s\"",
+						   (unsigned long long) cstate->cur_lineno,
+						   cstate->cur_attname,
+						   attval));
+			pfree(attval);
+		}
+		else
+			ereport(NOTICE,
+					errmsg("skipping row due to data type incompatibility at line %llu for column \"%s\": null input",
+						   (unsigned long long) cstate->cur_lineno,
+						   cstate->cur_attname));
+
+		/* reset relname_only */
+		cstate->relname_only = false;
+	}
+}
+
 /*
  * Workhorse for CopyFromTextOneRow() and CopyFromCSVOneRow().
  *
@@ -1033,42 +1089,7 @@ CopyFromTextLikeOneRow(CopyFromState cstate, ExprContext *econtext,
 										(Node *) cstate->escontext,
 										&values[m]))
 		{
-			Assert(cstate->opts.on_error != COPY_ON_ERROR_STOP);
-
-			cstate->num_errors++;
-
-			if (cstate->opts.log_verbosity == COPY_LOG_VERBOSITY_VERBOSE)
-			{
-				/*
-				 * Since we emit line number and column info in the below
-				 * notice message, we suppress error context information other
-				 * than the relation name.
-				 */
-				Assert(!cstate->relname_only);
-				cstate->relname_only = true;
-
-				if (cstate->cur_attval)
-				{
-					char	   *attval;
-
-					attval = CopyLimitPrintoutLength(cstate->cur_attval);
-					ereport(NOTICE,
-							errmsg("skipping row due to data type incompatibility at line %llu for column \"%s\": \"%s\"",
-								   (unsigned long long) cstate->cur_lineno,
-								   cstate->cur_attname,
-								   attval));
-					pfree(attval);
-				}
-				else
-					ereport(NOTICE,
-							errmsg("skipping row due to data type incompatibility at line %llu for column \"%s\": null input",
-								   (unsigned long long) cstate->cur_lineno,
-								   cstate->cur_attname));
-
-				/* reset relname_only */
-				cstate->relname_only = false;
-			}
-
+			CopyFromSkipErrorRow(cstate);
 			return true;
 		}
 

src/backend/commands/copyto.c

@@ -456,6 +456,18 @@ CopySendEndOfRow(CopyToState cstate)
 	resetStringInfo(fe_msgbuf);
 }
 
+/*
+ * Export CopySendEndOfRow() for extensions. We want to keep
+ * CopySendEndOfRow() as a static function for
+ * optimization. CopySendEndOfRow() calls in this file may be optimized by a
+ * compiler.
+ */
+void
+CopyToStateFlush(CopyToState cstate)
+{
+	CopySendEndOfRow(cstate);
+}
+
 /*
  * Wrapper function of CopySendEndOfRow for text and CSV formats. Sends the
  * line termination and do common appropriate things for the end of row.

src/include/commands/copyapi.h

@@ -56,6 +56,8 @@ typedef struct CopyToRoutine
 	void		(*CopyToEnd) (CopyToState cstate);
 } CopyToRoutine;
 
+extern void CopyToStateFlush(CopyToState cstate);
+
 /*
  * API structure for a COPY FROM format implementation. Note this must be
  * allocated in a server-lifetime manner, typically as a static const struct.
@@ -106,4 +108,8 @@ typedef struct CopyFromRoutine
 	void		(*CopyFromEnd) (CopyFromState cstate);
 } CopyFromRoutine;
 
+extern int	CopyFromStateGetData(CopyFromState cstate, void *dest, int minread, int maxread);
+
+extern void CopyFromSkipErrorRow(CopyFromState cstate);
+
 #endif							/* COPYAPI_H */

src/include/commands/copyfrom_internal.h

@@ -181,6 +181,9 @@ typedef struct CopyFromStateData
 #define RAW_BUF_BYTES(cstate) ((cstate)->raw_buf_len - (cstate)->raw_buf_index)
 
 	uint64		bytes_processed;	/* number of bytes processed so far */
+
+	/* For custom format implementation */
+	void	   *opaque;			/* private space */
 } CopyFromStateData;
 
 extern void ReceiveCopyBegin(CopyFromState cstate);

src/include/commands/copyto_internal.h

@@ -78,6 +78,9 @@ typedef struct CopyToStateData
 	FmgrInfo   *out_functions;	/* lookup info for output functions */
 	MemoryContext rowcontext;	/* per-row evaluation context */
 	uint64		bytes_processed;	/* number of bytes processed so far */
+
+	/* For custom format implementation */
+	void	   *opaque;			/* private space */
 } CopyToStateData;
 
 #endif							/* COPYTO_INTERNAL_H */

src/test/modules/test_copy_format/expected/no_schema.out

@@ -1,14 +1,59 @@
 CREATE EXTENSION test_copy_format;
 CREATE TABLE public.test (a smallint, b integer, c bigint);
 INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789);
+-- 987 is accepted.
+-- 654 is a hard error because ON_ERROR is stop by default.
 COPY public.test FROM stdin WITH (FORMAT 'test_copy_format');
 NOTICE:  test_copy_format: is_from=true
 NOTICE:  CopyFromInFunc: attribute: smallint
 NOTICE:  CopyFromInFunc: attribute: integer
 NOTICE:  CopyFromInFunc: attribute: bigint
 NOTICE:  CopyFromStart: the number of attributes: 3
 NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromOneRow
+ERROR:  invalid value: "6"
+CONTEXT:  COPY test, line 2, column a: "6"
+-- 987 is accepted.
+-- 654 is a soft error because ON_ERROR is ignore.
+COPY public.test FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore);
+NOTICE:  test_copy_format: is_from=true
+NOTICE:  CopyFromInFunc: attribute: smallint
+NOTICE:  CopyFromInFunc: attribute: integer
+NOTICE:  CopyFromInFunc: attribute: bigint
+NOTICE:  CopyFromStart: the number of attributes: 3
+NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromOneRow
+NOTICE:  1 row was skipped due to data type incompatibility
+NOTICE:  CopyFromEnd
+-- 987 is accepted.
+-- 654 is a soft error because ON_ERROR is ignore.
+COPY public.test FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore, LOG_VERBOSITY verbose);
+NOTICE:  test_copy_format: is_from=true
+NOTICE:  CopyFromInFunc: attribute: smallint
+NOTICE:  CopyFromInFunc: attribute: integer
+NOTICE:  CopyFromInFunc: attribute: bigint
+NOTICE:  CopyFromStart: the number of attributes: 3
+NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromOneRow
+NOTICE:  skipping row due to data type incompatibility at line 2 for column "a": "6"
+NOTICE:  CopyFromOneRow
+NOTICE:  1 row was skipped due to data type incompatibility
 NOTICE:  CopyFromEnd
+-- 987 is accepted.
+-- 654 is a soft error because ON_ERROR is ignore.
+-- 321 is a hard error.
+COPY public.test FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore);
+NOTICE:  test_copy_format: is_from=true
+NOTICE:  CopyFromInFunc: attribute: smallint
+NOTICE:  CopyFromInFunc: attribute: integer
+NOTICE:  CopyFromInFunc: attribute: bigint
+NOTICE:  CopyFromStart: the number of attributes: 3
+NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromOneRow
+ERROR:  too much lines: 3
+CONTEXT:  COPY test, line 3
 COPY public.test TO stdout WITH (FORMAT 'test_copy_format');
 NOTICE:  test_copy_format: is_from=false
 NOTICE:  CopyToOutFunc: attribute: smallint
@@ -18,6 +63,8 @@ NOTICE:  CopyToStart: the number of attributes: 3
 NOTICE:  CopyToOneRow: the number of valid values: 3
 NOTICE:  CopyToOneRow: the number of valid values: 3
 NOTICE:  CopyToOneRow: the number of valid values: 3
+NOTICE:  CopyToOneRow: the number of valid values: 3
+NOTICE:  CopyToOneRow: the number of valid values: 3
 NOTICE:  CopyToEnd
 DROP TABLE public.test;
 DROP EXTENSION test_copy_format;

src/test/modules/test_copy_format/sql/no_schema.sql

@@ -1,7 +1,31 @@
 CREATE EXTENSION test_copy_format;
 CREATE TABLE public.test (a smallint, b integer, c bigint);
 INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789);
+-- 987 is accepted.
+-- 654 is a hard error because ON_ERROR is stop by default.
 COPY public.test FROM stdin WITH (FORMAT 'test_copy_format');
+987
+654
+\.
+-- 987 is accepted.
+-- 654 is a soft error because ON_ERROR is ignore.
+COPY public.test FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore);
+987
+654
+\.
+-- 987 is accepted.
+-- 654 is a soft error because ON_ERROR is ignore.
+COPY public.test FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore, LOG_VERBOSITY verbose);
+987
+654
+\.
+-- 987 is accepted.
+-- 654 is a soft error because ON_ERROR is ignore.
+-- 321 is a hard error.
+COPY public.test FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore);
+987
+654
+321
 \.
 COPY public.test TO stdout WITH (FORMAT 'test_copy_format');
 DROP TABLE public.test;

src/test/modules/test_copy_format/test_copy_format.c

@@ -14,6 +14,7 @@
 #include "postgres.h"
 
 #include "commands/copyapi.h"
+#include "commands/copyfrom_internal.h"
 #include "commands/defrem.h"
 #include "utils/builtins.h"
 
@@ -35,8 +36,85 @@ TestCopyFromStart(CopyFromState cstate, TupleDesc tupDesc)
 static bool
 TestCopyFromOneRow(CopyFromState cstate, ExprContext *econtext, Datum *values, bool *nulls)
 {
+	int			n_attributes = list_length(cstate->attnumlist);
+	char	   *line;
+	int			line_size = n_attributes + 1;	/* +1 is for new line */
+	int			read_bytes;
+
 	ereport(NOTICE, (errmsg("CopyFromOneRow")));
-	return false;
+
+	cstate->cur_lineno++;
+	line = palloc(line_size);
+	read_bytes = CopyFromStateGetData(cstate, line, line_size, line_size);
+	if (read_bytes == 0)
+		return false;
+	if (read_bytes != line_size)
+		ereport(ERROR,
+				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
+				 errmsg("one line must be %d bytes: %d",
+						line_size, read_bytes)));
+
+	if (cstate->cur_lineno == 1)
+	{
+		/* Success */
+		TupleDesc	tupDesc = RelationGetDescr(cstate->rel);
+		ListCell   *cur;
+		int			i = 0;
+
+		foreach(cur, cstate->attnumlist)
+		{
+			int			attnum = lfirst_int(cur);
+			int			m = attnum - 1;
+			Form_pg_attribute att = TupleDescAttr(tupDesc, m);
+
+			if (att->atttypid == INT2OID)
+			{
+				values[i] = Int16GetDatum(line[i] - '0');
+			}
+			else if (att->atttypid == INT4OID)
+			{
+				values[i] = Int32GetDatum(line[i] - '0');
+			}
+			else if (att->atttypid == INT8OID)
+			{
+				values[i] = Int64GetDatum(line[i] - '0');
+			}
+			nulls[i] = false;
+			i++;
+		}
+	}
+	else if (cstate->cur_lineno == 2)
+	{
+		/* Soft error */
+		TupleDesc	tupDesc = RelationGetDescr(cstate->rel);
+		int			attnum = lfirst_int(list_head(cstate->attnumlist));
+		int			m = attnum - 1;
+		Form_pg_attribute att = TupleDescAttr(tupDesc, m);
+		char		value[2];
+
+		cstate->cur_attname = NameStr(att->attname);
+		value[0] = line[0];
+		value[1] = '\0';
+		cstate->cur_attval = value;
+		errsave((Node *) cstate->escontext,
+				(
+				 errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
+				 errmsg("invalid value: \"%c\"", line[0])));
+		CopyFromSkipErrorRow(cstate);
+		cstate->cur_attname = NULL;
+		cstate->cur_attval = NULL;
+		return true;
+	}
+	else
+	{
+		/* Hard error */
+		ereport(ERROR,
+				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
+				 errmsg("too much lines: %llu",
+						(unsigned long long) cstate->cur_lineno)));
+	}
+
+	return true;
 }
 
 static void