From 8066477384ad660bfa5438f0b42ae5a24c3f17fb Mon Sep 17 00:00:00 2001
From: Sutou Kouhei <kou@clear-code.com>
Date: Wed, 27 Nov 2024 16:23:55 +0900
Subject: [PATCH v35 7/7] Add CopyFromSkipErrorRow() for custom COPY format
 extension

Extensions must call CopyFromSkipErrorRow() when CopyFromOneRow
callback reports an error by errsave(). CopyFromSkipErrorRow() handles
"ON_ERROR stop" and "LOG_VERBOSITY verbose" cases.
---
 src/backend/commands/copyfromparse.c          | 82 +++++++++++--------
 src/include/commands/copyapi.h                |  2 +
 .../expected/test_copy_format.out             | 47 +++++++++++
 .../test_copy_format/sql/test_copy_format.sql | 24 ++++++
 .../test_copy_format/test_copy_format.c       | 80 +++++++++++++++++-
 5 files changed, 198 insertions(+), 37 deletions(-)

diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c
index d8fd238e72b..2070f51a963 100644
--- a/src/backend/commands/copyfromparse.c
+++ b/src/backend/commands/copyfromparse.c
@@ -938,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().
  *
@@ -1044,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;
 		}
 
diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h
index 2044d8b8c4c..500ece7d5bb 100644
--- a/src/include/commands/copyapi.h
+++ b/src/include/commands/copyapi.h
@@ -110,4 +110,6 @@ typedef struct CopyFromRoutine
 
 extern int	CopyFromStateGetData(CopyFromState cstate, void *dest, int minread, int maxread);
 
+extern void CopyFromSkipErrorRow(CopyFromState cstate);
+
 #endif							/* COPYAPI_H */
diff --git a/src/test/modules/test_copy_format/expected/test_copy_format.out b/src/test/modules/test_copy_format/expected/test_copy_format.out
index 016893e7026..b9a6baa85c0 100644
--- a/src/test/modules/test_copy_format/expected/test_copy_format.out
+++ b/src/test/modules/test_copy_format/expected/test_copy_format.out
@@ -1,6 +1,8 @@
 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: atttypid=21
@@ -8,7 +10,50 @@ NOTICE:  CopyFromInFunc: atttypid=23
 NOTICE:  CopyFromInFunc: atttypid=20
 NOTICE:  CopyFromStart: natts=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: atttypid=21
+NOTICE:  CopyFromInFunc: atttypid=23
+NOTICE:  CopyFromInFunc: atttypid=20
+NOTICE:  CopyFromStart: natts=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: atttypid=21
+NOTICE:  CopyFromInFunc: atttypid=23
+NOTICE:  CopyFromInFunc: atttypid=20
+NOTICE:  CopyFromStart: natts=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: atttypid=21
+NOTICE:  CopyFromInFunc: atttypid=23
+NOTICE:  CopyFromInFunc: atttypid=20
+NOTICE:  CopyFromStart: natts=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: atttypid=21
@@ -18,4 +63,6 @@ NOTICE:  CopyToStart: natts=3
 NOTICE:  CopyToOneRow: tts_nvalid=3
 NOTICE:  CopyToOneRow: tts_nvalid=3
 NOTICE:  CopyToOneRow: tts_nvalid=3
+NOTICE:  CopyToOneRow: tts_nvalid=3
+NOTICE:  CopyToOneRow: tts_nvalid=3
 NOTICE:  CopyToEnd
diff --git a/src/test/modules/test_copy_format/sql/test_copy_format.sql b/src/test/modules/test_copy_format/sql/test_copy_format.sql
index 0dfdfa00080..86db71bce7f 100644
--- a/src/test/modules/test_copy_format/sql/test_copy_format.sql
+++ b/src/test/modules/test_copy_format/sql/test_copy_format.sql
@@ -1,6 +1,30 @@
 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');
diff --git a/src/test/modules/test_copy_format/test_copy_format.c b/src/test/modules/test_copy_format/test_copy_format.c
index abafc668463..96a54dab7ec 100644
--- a/src/test/modules/test_copy_format/test_copy_format.c
+++ b/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"
 
 PG_MODULE_MAGIC;
@@ -34,8 +35,85 @@ CopyFromStart(CopyFromState cstate, TupleDesc tupDesc)
 static bool
 CopyFromOneRow(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
-- 
2.47.2

