From c1ab564adbe8db74c97c1d0e4ee6ec552457474c Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Mon, 10 Nov 2025 11:26:10 -0800
Subject: [PATCH 3/6] Support custom COPY format for COPY FROM and COPY TO.

This commit enables extensions to register their own COPY format
implementation for COPY TO and COPY FROM. If extensions register the
custom format during _PG_init_ time, the format would be available for
all backends whereas the backend can LOAD the custom format.

Author:
Reviewed-by:
Discussion: https://postgr.es/m/
Backpatch-through:
---
 src/backend/commands/Makefile             |   1 +
 src/backend/commands/copy_custom_format.c | 142 ++++++++++++++++++++++
 src/backend/commands/copyfrom.c           |   7 ++
 src/backend/commands/copyfromparse.c      |   9 ++
 src/backend/commands/copyto.c             |  17 +++
 src/backend/commands/meson.build          |   1 +
 src/include/commands/copy.h               |   6 +
 src/include/commands/copyapi.h            |   9 ++
 8 files changed, 192 insertions(+)
 create mode 100644 src/backend/commands/copy_custom_format.c

diff --git a/src/backend/commands/Makefile b/src/backend/commands/Makefile
index f99acfd2b4b..74668d4ac9d 100644
--- a/src/backend/commands/Makefile
+++ b/src/backend/commands/Makefile
@@ -27,6 +27,7 @@ OBJS = \
 	copyfrom.o \
 	copyfromparse.o \
 	copyto.o \
+	copy_custom_format.o \
 	createas.o \
 	dbcommands.o \
 	define.o \
diff --git a/src/backend/commands/copy_custom_format.c b/src/backend/commands/copy_custom_format.c
new file mode 100644
index 00000000000..8bef6e779ac
--- /dev/null
+++ b/src/backend/commands/copy_custom_format.c
@@ -0,0 +1,142 @@
+/*-------------------------------------------------------------------------
+ *
+ * copy_custom_format.c
+ *
+ * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ *
+ * IDENTIFICATION
+ *	  src/backend/commands/copy_custom_format.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "commands/copy.h"
+#include "commands/copyapi.h"
+#include "utils/memutils.h"
+
+
+typedef struct CopyCustomFormat
+{
+	const char *fmt_name;
+	const CopyFromRoutine *from_routine;
+	const CopyToRoutine *to_routine;
+}			CopyCustomFormat;
+
+static CopyCustomFormat * CopyCustomFormatArray = NULL;
+static int	CopyCustomFormatAssigned = 0;
+static int	CopyCustomFormatAllocated = 0;
+
+static char *BuiltinFormatNames[] = {
+	"text", "csv", "binary"
+};
+
+void
+RegisterCopyCustomFormat(const char *fmt_name, const CopyFromRoutine *from_routine,
+						 const CopyToRoutine *to_routine)
+{
+	CopyCustomFormat *entry;
+
+	/* Check the duplication with built-in formats */
+	for (int i = 0; i < lengthof(BuiltinFormatNames); i++)
+	{
+		if (strcmp(BuiltinFormatNames[i], fmt_name) == 0)
+			ereport(ERROR,
+					errmsg("failed to register custom COPY format \"%s\"", fmt_name),
+					errdetail("COPY format \"%s\" is a builtin format",
+							  BuiltinFormatNames[i]));
+
+	}
+
+	/* Check the duplication with the registered custom formats */
+	if (FindCustomCopyFormat(fmt_name))
+		ereport(ERROR,
+				errmsg("failed to register custom COPY format \"%s\"", fmt_name),
+				errdetail("Custom COPY format \"%s\" already registered",
+						  fmt_name));
+
+	/* If there is no array yet, create one */
+	if (CopyCustomFormatArray == NULL)
+	{
+		CopyCustomFormatAllocated = 16;
+		CopyCustomFormatArray = (CopyCustomFormat *)
+			MemoryContextAlloc(TopMemoryContext,
+							   CopyCustomFormatAllocated * sizeof(CopyCustomFormat));
+	}
+
+	/* If there's an array but it's current full, expand it */
+	if (CopyCustomFormatAssigned >= CopyCustomFormatAllocated)
+	{
+		int			i = pg_nextpower2_32(CopyCustomFormatAllocated + 1);
+
+		CopyCustomFormatArray = (CopyCustomFormat *)
+			repalloc(CopyCustomFormatArray, i * sizeof(CopyCustomFormat));
+		CopyCustomFormatAllocated = i;
+	}
+
+	Assert(to_routine != NULL || from_routine != NULL);
+	Assert((to_routine == NULL) ||
+		   (to_routine->CopyToEstimateStateSpace != NULL &&
+			to_routine->CopyToOutFunc != NULL &&
+			to_routine->CopyToStart != NULL &&
+			to_routine->CopyToOneRow != NULL &&
+			to_routine->CopyToEnd != NULL));
+	Assert((from_routine == NULL) ||
+		   (from_routine->CopyFromEstimateStateSpace != NULL &&
+			from_routine->CopyFromInFunc != NULL &&
+			from_routine->CopyFromStart != NULL &&
+			from_routine->CopyFromOneRow != NULL &&
+			from_routine->CopyFromEnd != NULL));
+
+	entry = &CopyCustomFormatArray[CopyCustomFormatAssigned++];
+	entry->fmt_name = fmt_name;
+	entry->from_routine = from_routine;
+	entry->to_routine = to_routine;
+}
+
+/*
+ * Returns true if the given custom format name is registered.
+ */
+bool
+FindCustomCopyFormat(const char *fmt_name)
+{
+	for (int i = 0; i < CopyCustomFormatAssigned; i++)
+	{
+		if (strcmp(CopyCustomFormatArray[i].fmt_name, fmt_name) == 0)
+			return true;
+	}
+
+	return false;
+}
+
+bool
+GetCustomCopyToRoutine(const char *fmt_name, const CopyToRoutine **to_routine)
+{
+	for (int i = 0; i < CopyCustomFormatAssigned; i++)
+	{
+		if (strcmp(CopyCustomFormatArray[i].fmt_name, fmt_name) == 0)
+		{
+			*to_routine = CopyCustomFormatArray[i].to_routine;
+			return true;
+		}
+	}
+
+	return false;
+}
+
+bool
+GetCustomCopyFromRoutine(const char *fmt_name, const CopyFromRoutine **from_routine)
+{
+	for (int i = 0; i < CopyCustomFormatAssigned; i++)
+	{
+		if (strcmp(CopyCustomFormatArray[i].fmt_name, fmt_name) == 0)
+		{
+			*from_routine = CopyCustomFormatArray[i].from_routine;
+			return true;
+		}
+	}
+
+	return false;
+}
diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c
index 0c51e5ba5e1..592be4fcb5d 100644
--- a/src/backend/commands/copyfrom.c
+++ b/src/backend/commands/copyfrom.c
@@ -1619,6 +1619,13 @@ create_copyfrom_state(ParseState *pstate, List *options)
 				routine = &CopyFromRoutineCSV;
 			else if (strcmp(fmt, "binary") == 0)
 				routine = &CopyFromRoutineBinary;
+			else if (GetCustomCopyFromRoutine(fmt, &routine))
+			{
+				if (routine == NULL)
+					ereport(ERROR,
+							errmsg("COPY format \"%s\" does not support COPY FROM", fmt),
+							parser_errposition(pstate, defel->location));
+			}
 			else
 				ereport(ERROR,
 						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c
index 669dbfb8459..30d8047f21a 100644
--- a/src/backend/commands/copyfromparse.c
+++ b/src/backend/commands/copyfromparse.c
@@ -228,6 +228,15 @@ ReceiveCopyBinaryHeader(CopyFromState ccstate)
 	}
 }
 
+/*
+ * Exporting CopyGetData() function for custom COPY FROM format implementations.
+ */
+int
+CopyFromGetData(CopyFromState cstate, void *databuf, int minread, int maxread)
+{
+	return CopyGetData(cstate, databuf, minread, maxread);
+}
+
 /*
  * CopyGetData reads data from the source (file or frontend)
  *
diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c
index 6a0a66507ba..b1b3ae141eb 100644
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -435,6 +435,16 @@ CopySendChar(CopyToState cstate, char c)
 	appendStringInfoCharMacro(cstate->fe_msgbuf, c);
 }
 
+/*
+ * Exporting CopySendEndOfRow() function for custom COPY TO format
+ * implementations.
+ */
+void
+CopyToFlushData(CopyToState cstate)
+{
+	CopySendEndOfRow(cstate);
+}
+
 static void
 CopySendEndOfRow(CopyToState cstate)
 {
@@ -631,6 +641,13 @@ create_copyto_state(ParseState *pstate, List *options)
 				routine = &CopyToRoutineCSV;
 			else if (strcmp(fmt, "binary") == 0)
 				routine = &CopyToRoutineBinary;
+			else if (GetCustomCopyToRoutine(fmt, &routine))
+			{
+				if (routine == NULL)
+					ereport(ERROR,
+							errmsg("COPY format \"%s\" does not support COPY TO", fmt),
+							parser_errposition(pstate, defel->location));
+			}
 			else
 				ereport(ERROR,
 						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
diff --git a/src/backend/commands/meson.build b/src/backend/commands/meson.build
index 9f640ad4810..827f809eb53 100644
--- a/src/backend/commands/meson.build
+++ b/src/backend/commands/meson.build
@@ -15,6 +15,7 @@ backend_sources += files(
   'copyfrom.c',
   'copyfromparse.c',
   'copyto.c',
+  'copy_custom_format.c',
   'createas.c',
   'dbcommands.c',
   'define.c',
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index 30a1d2bff6e..82f07b05823 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -120,6 +120,12 @@ extern uint64 CopyFrom(CopyFromState cstate);
 
 extern DestReceiver *CreateCopyDestReceiver(void);
 
+extern void ProcessCopyBuiltinOptions(List *options, CopyFormatOptions *opts_out,
+									  bool is_from, List **other_options, ParseState *pstate);
+
+/* defined in copy_custom_format.c */
+extern bool FindCustomCopyFormat(const char *fmt_name);
+
 /*
  * internal prototypes
  */
diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h
index c3d2199a0b6..2090a4e1c61 100644
--- a/src/include/commands/copyapi.h
+++ b/src/include/commands/copyapi.h
@@ -111,4 +111,13 @@ typedef struct CopyFromRoutine
 	 */
 	void		(*CopyFromEnd) (CopyFromState cstate);
 } CopyFromRoutine;
+
+extern int CopyFromGetData(CopyFromState cstate, void *databuf, int minread, int maxread);
+extern void CopyToFlushData(CopyToState cstate);
+
+extern void RegisterCopyCustomFormat(const char *fmt_name, const CopyFromRoutine *from_routine,
+									 const CopyToRoutine *to_routine);
+extern bool GetCustomCopyToRoutine(const char *fmt_name, const CopyToRoutine **to_routine);
+extern bool GetCustomCopyFromRoutine(const char *fmt_name, const CopyFromRoutine **from_routine);
+
 #endif							/* COPYAPI_H */
-- 
2.47.3

