0005-Introduce-NULL_ON_ERROR-option-for-COPY-FROM.patch

From 1fc1b4231b079078c303be75864c4ef71d0742a1 Mon Sep 17 00:00:00 2001
From: Nikita Glukhov <n.gluhov@postgrespro.ru>
Date: Wed, 31 Aug 2022 22:29:47 +0300
Subject: [PATCH 5/5] Introduce NULL_ON_ERROR option for COPY FROM

---
 src/backend/commands/copy.c          | 20 +++++++++++++++++
 src/backend/commands/copyfrom.c      | 33 ++++++++++++++++++++++++++++
 src/backend/commands/copyfromparse.c | 21 +++++++++++++-----
 src/include/commands/copy.h          |  2 ++
 src/test/regress/expected/copy.out   | 15 +++++++++++++
 src/test/regress/sql/copy.sql        | 13 +++++++++++
 6 files changed, 99 insertions(+), 5 deletions(-)

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 49924e476af..ce7e875f655 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -551,6 +551,19 @@ ProcessCopyOptions(ParseState *pstate,
 								defel->defname),
 						 parser_errposition(pstate, defel->location)));
 		}
+		else if (strcmp(defel->defname, "null_on_error") == 0)
+		{
+			if (opts_out->null_on_error)
+				errorConflictingDefElem(defel, pstate);
+			if (defel->arg && IsA(defel->arg, List))
+				opts_out->null_on_error = castNode(List, defel->arg);
+			else
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						 errmsg("argument to option \"%s\" must be a list of column names",
+								defel->defname),
+						 parser_errposition(pstate, defel->location)));
+		}
 		else
 			ereport(ERROR,
 					(errcode(ERRCODE_SYNTAX_ERROR),
@@ -701,6 +714,13 @@ ProcessCopyOptions(ParseState *pstate,
 		ereport(ERROR,
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("CSV quote character must not appear in the NULL specification")));
+
+	/* Check null_on_error */
+	if (opts_out->null_on_error != NIL && !is_from)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("COPY null on error only available using COPY FROM")));
+
 }
 
 /*
diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c
index e8bb168aea8..99b39d8d7c2 100644
--- a/src/backend/commands/copyfrom.c
+++ b/src/backend/commands/copyfrom.c
@@ -47,6 +47,7 @@
 #include "rewrite/rewriteHandler.h"
 #include "storage/fd.h"
 #include "tcop/tcopprot.h"
+#include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/portal.h"
@@ -1316,6 +1317,29 @@ BeginCopyFrom(ParseState *pstate,
 		}
 	}
 
+	/* Convert NULL_ON_ERROR name list to per-column flags, check validity */
+	cstate->opts.null_on_error_flags = (bool *) palloc0(num_phys_attrs * sizeof(bool));
+	if (cstate->opts.null_on_error)
+	{
+		List	   *attnums;
+		ListCell   *cur;
+
+		attnums = CopyGetAttnums(tupDesc, cstate->rel, cstate->opts.null_on_error);
+
+		foreach(cur, attnums)
+		{
+			int			attnum = lfirst_int(cur);
+			Form_pg_attribute attr = TupleDescAttr(tupDesc, attnum - 1);
+
+			if (!list_member_int(cstate->attnumlist, attnum))
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_COLUMN_REFERENCE),
+						 errmsg("NULL_ON_ERROR column \"%s\" not referenced by COPY",
+								NameStr(attr->attname))));
+			cstate->opts.null_on_error_flags[attnum - 1] = true;
+		}
+	}
+
 	/* Use client encoding when ENCODING option is not specified. */
 	if (cstate->opts.file_encoding < 0)
 		cstate->file_encoding = pg_get_client_encoding();
@@ -1417,6 +1441,15 @@ BeginCopyFrom(ParseState *pstate,
 							 &in_func_oid, &typioparams[attnum - 1]);
 		fmgr_info(in_func_oid, &in_functions[attnum - 1]);
 
+		/* check whether input function supports returning errors */
+		if (cstate->opts.null_on_error_flags[attnum - 1] &&
+			!func_is_safe(in_func_oid))
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("input function for datatype \"%s\" does not support error handling",
+							format_type_be(att->atttypid))));
+
+
 		/* Get default info if needed */
 		if (!list_member_int(cstate->attnumlist, attnum) && !att->attgenerated)
 		{
diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c
index 097414ef12d..04a36de444c 100644
--- a/src/backend/commands/copyfromparse.c
+++ b/src/backend/commands/copyfromparse.c
@@ -897,6 +897,8 @@ NextCopyFrom(CopyFromState cstate, ExprContext *econtext,
 			int			attnum = lfirst_int(cur);
 			int			m = attnum - 1;
 			Form_pg_attribute att = TupleDescAttr(tupDesc, m);
+			ErrorData  *edata = NULL;
+			ErrorData **p_edata;
 
 			if (fieldno >= fldct)
 				ereport(ERROR,
@@ -938,11 +940,20 @@ NextCopyFrom(CopyFromState cstate, ExprContext *econtext,
 
 			cstate->cur_attname = NameStr(att->attname);
 			cstate->cur_attval = string;
-			values[m] = InputFunctionCall(&in_functions[m],
-										  string,
-										  typioparams[m],
-										  att->atttypmod);
-			if (string != NULL)
+
+			p_edata = cstate->opts.null_on_error_flags[m] ? &edata : NULL;
+
+			values[m] = InputFunctionCallOptError(&in_functions[m],
+												  string,
+												  typioparams[m],
+												  att->atttypmod,
+												  p_edata);
+			if (edata)
+			{
+				FreeErrorData(edata);
+				nulls[m] = true;
+			}
+			else if (string != NULL)
 				nulls[m] = false;
 			cstate->cur_attname = NULL;
 			cstate->cur_attval = NULL;
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index cb0096aeb67..359fb4cc9e0 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -59,6 +59,8 @@ typedef struct CopyFormatOptions
 	bool	   *force_null_flags;	/* per-column CSV FN flags */
 	bool		convert_selectively;	/* do selective binary conversion? */
 	List	   *convert_select; /* list of column names (can be NIL) */
+	List	   *null_on_error;	/* list of column names */
+	bool	   *null_on_error_flags;	/* per-column NOE flags */
 } CopyFormatOptions;
 
 /* These are private in commands/copy[from|to].c */
diff --git a/src/test/regress/expected/copy.out b/src/test/regress/expected/copy.out
index 3fad1c52d1f..82839d2f67a 100644
--- a/src/test/regress/expected/copy.out
+++ b/src/test/regress/expected/copy.out
@@ -240,3 +240,18 @@ SELECT * FROM header_copytest ORDER BY a;
 (5 rows)
 
 drop table header_copytest;
+create table null_on_error_copytest(i int, f float4);
+copy null_on_error_copytest from stdin with (null_on_error (i, f));
+ERROR:  input function for datatype "integer" does not support error handling
+copy null_on_error_copytest from stdin with (null_on_error (f));
+copy null_on_error_copytest from stdin with (null_on_error (f));
+ERROR:  invalid input syntax for type integer: "err"
+CONTEXT:  COPY null_on_error_copytest, line 2, column i: "err"
+select * from null_on_error_copytest;
+ i |  f  
+---+-----
+ 1 | 2.3
+ 2 |    
+(2 rows)
+
+drop table null_on_error_copytest;
diff --git a/src/test/regress/sql/copy.sql b/src/test/regress/sql/copy.sql
index 285022e07c6..914e29d5ba1 100644
--- a/src/test/regress/sql/copy.sql
+++ b/src/test/regress/sql/copy.sql
@@ -268,3 +268,16 @@ a	c	b
 
 SELECT * FROM header_copytest ORDER BY a;
 drop table header_copytest;
+
+create table null_on_error_copytest(i int, f float4);
+copy null_on_error_copytest from stdin with (null_on_error (i, f));
+copy null_on_error_copytest from stdin with (null_on_error (f));
+1	2.3
+2	err
+\.
+copy null_on_error_copytest from stdin with (null_on_error (f));
+1	2.3
+err	4.5
+\.
+select * from null_on_error_copytest;
+drop table null_on_error_copytest;
-- 
2.25.1