v8-0006-Add-support-for-adding-custom-COPY-FROM-format.patch

From f48e7b629a8d15fc70cd4cc4737dd2ad61910cc9 Mon Sep 17 00:00:00 2001
From: Sutou Kouhei <kou@clear-code.com>
Date: Wed, 24 Jan 2024 11:07:14 +0900
Subject: [PATCH v8 06/10] Add support for adding custom COPY FROM format

We use the same approach as we used for custom COPY TO format. Now,
custom COPY format handler can return COPY TO format routines or COPY
FROM format routines based on the "is_from" argument:

    copy_handler(true) returns CopyToRoutine
    copy_handler(false) returns CopyFromRoutine
---
 src/backend/commands/copy.c                   | 53 +++++++++++++------
 src/include/commands/copyapi.h                |  2 +
 .../expected/test_copy_format.out             | 12 +++++
 .../test_copy_format/sql/test_copy_format.sql |  6 +++
 .../test_copy_format/test_copy_format.c       | 50 +++++++++++++++--
 5 files changed, 105 insertions(+), 18 deletions(-)

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index ec6dfff8ab..479f36868c 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -472,12 +472,9 @@ ProcessCopyOptionCustomFormat(ParseState *pstate,
 	}
 
 	/* custom format */
-	if (!is_from)
-	{
-		funcargtypes[0] = INTERNALOID;
-		handlerOid = LookupFuncName(list_make1(makeString(format)), 1,
-									funcargtypes, true);
-	}
+	funcargtypes[0] = INTERNALOID;
+	handlerOid = LookupFuncName(list_make1(makeString(format)), 1,
+								funcargtypes, true);
 	if (!OidIsValid(handlerOid))
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
@@ -486,14 +483,36 @@ ProcessCopyOptionCustomFormat(ParseState *pstate,
 
 	datum = OidFunctionCall1(handlerOid, BoolGetDatum(is_from));
 	routine = DatumGetPointer(datum);
-	if (routine == NULL || !IsA(routine, CopyToRoutine))
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("COPY handler function %s(%u) did not return a CopyToRoutine struct",
-						format, handlerOid),
-				 parser_errposition(pstate, defel->location)));
-
-	opts_out->to_routine = routine;
+	if (is_from)
+	{
+		if (routine == NULL || !IsA(routine, CopyFromRoutine))
+			ereport(
+					ERROR,
+					(errcode(
+							 ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("COPY handler function "
+							"%s(%u) did not return a "
+							"CopyFromRoutine struct",
+							format, handlerOid),
+					 parser_errposition(
+										pstate, defel->location)));
+		opts_out->from_routine = routine;
+	}
+	else
+	{
+		if (routine == NULL || !IsA(routine, CopyToRoutine))
+			ereport(
+					ERROR,
+					(errcode(
+							 ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("COPY handler function "
+							"%s(%u) did not return a "
+							"CopyToRoutine struct",
+							format, handlerOid),
+					 parser_errposition(
+										pstate, defel->location)));
+		opts_out->to_routine = routine;
+	}
 }
 
 /*
@@ -692,7 +711,11 @@ ProcessCopyOptions(ParseState *pstate,
 		{
 			bool		processed = false;
 
-			if (!is_from)
+			if (is_from)
+				processed =
+					opts_out->from_routine->CopyFromProcessOption(
+																  cstate, defel);
+			else
 				processed = opts_out->to_routine->CopyToProcessOption(cstate, defel);
 			if (!processed)
 				ereport(ERROR,
diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h
index 323e4705d2..ef1bb201c2 100644
--- a/src/include/commands/copyapi.h
+++ b/src/include/commands/copyapi.h
@@ -30,6 +30,8 @@ typedef void (*CopyFromEnd_function) (CopyFromState cstate);
 /* Routines for a COPY FROM format implementation. */
 typedef struct CopyFromRoutine
 {
+	NodeTag		type;
+
 	/*
 	 * Called for processing one COPY FROM option. This will return false when
 	 * the given option is invalid.
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 3a24ae7b97..6af69f0eb7 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,18 @@
 CREATE EXTENSION test_copy_format;
 CREATE TABLE public.test (a INT, b INT, c INT);
 INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789);
+COPY public.test FROM stdin WITH (
+	option_before 'before',
+	format 'test_copy_format',
+	option_after 'after'
+);
+NOTICE:  test_copy_format: is_from=true
+NOTICE:  CopyFromProcessOption: "option_before"="before"
+NOTICE:  CopyFromProcessOption: "option_after"="after"
+NOTICE:  CopyFromGetFormat
+NOTICE:  CopyFromStart: natts=3
+NOTICE:  CopyFromOneRow
+NOTICE:  CopyFromEnd
 COPY public.test TO stdout WITH (
 	option_before 'before',
 	format 'test_copy_format',
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 0eb7ed2e11..94d3c789a0 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,12 @@
 CREATE EXTENSION test_copy_format;
 CREATE TABLE public.test (a INT, b INT, c INT);
 INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789);
+COPY public.test FROM stdin WITH (
+	option_before 'before',
+	format 'test_copy_format',
+	option_after 'after'
+);
+\.
 COPY public.test TO stdout WITH (
 	option_before 'before',
 	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 a2219afcde..5e1b40e881 100644
--- a/src/test/modules/test_copy_format/test_copy_format.c
+++ b/src/test/modules/test_copy_format/test_copy_format.c
@@ -18,6 +18,50 @@
 
 PG_MODULE_MAGIC;
 
+static bool
+CopyFromProcessOption(CopyFromState cstate, DefElem *defel)
+{
+	ereport(NOTICE,
+			(errmsg("CopyFromProcessOption: \"%s\"=\"%s\"",
+					defel->defname, defGetString(defel))));
+	return true;
+}
+
+static int16
+CopyFromGetFormat(CopyFromState cstate)
+{
+	ereport(NOTICE, (errmsg("CopyFromGetFormat")));
+	return 0;
+}
+
+static void
+CopyFromStart(CopyFromState cstate, TupleDesc tupDesc)
+{
+	ereport(NOTICE, (errmsg("CopyFromStart: natts=%d", tupDesc->natts)));
+}
+
+static bool
+CopyFromOneRow(CopyFromState cstate, ExprContext *econtext, Datum *values, bool *nulls)
+{
+	ereport(NOTICE, (errmsg("CopyFromOneRow")));
+	return false;
+}
+
+static void
+CopyFromEnd(CopyFromState cstate)
+{
+	ereport(NOTICE, (errmsg("CopyFromEnd")));
+}
+
+static const CopyFromRoutine CopyFromRoutineTestCopyFormat = {
+	.type = T_CopyFromRoutine,
+	.CopyFromProcessOption = CopyFromProcessOption,
+	.CopyFromGetFormat = CopyFromGetFormat,
+	.CopyFromStart = CopyFromStart,
+	.CopyFromOneRow = CopyFromOneRow,
+	.CopyFromEnd = CopyFromEnd,
+};
+
 static bool
 CopyToProcessOption(CopyToState cstate, DefElem *defel)
 {
@@ -71,7 +115,7 @@ test_copy_format(PG_FUNCTION_ARGS)
 			(errmsg("test_copy_format: is_from=%s", is_from ? "true" : "false")));
 
 	if (is_from)
-		elog(ERROR, "COPY FROM isn't supported yet");
-
-	PG_RETURN_POINTER(&CopyToRoutineTestCopyFormat);
+		PG_RETURN_POINTER(&CopyFromRoutineTestCopyFormat);
+	else
+		PG_RETURN_POINTER(&CopyToRoutineTestCopyFormat);
 }
-- 
2.41.0