[Sanaba, Bilva <bilvas@amazon.com>]

Hi hackers,



Currently, the COPY TO api does not support callback functions, while the COPY FROM api does. The COPY TO code does, however, include placeholders for supporting callbacks in the future.



Rounding out the support of callback functions to both could be very beneficial for extension development. In particular, supporting callbacks for COPY TO will allow developers to utilize the preexisting command in order to create tools that give users more support for moving data for storage, backup, analytics, etc.



We are aiming to get the support in core PostgreSQL and add COPY TO callback support in the next commitfest.  The attached patch contains a change to COPY TO api to support callbacks.



Best,

Bilva



Attachments:

0001-Support-COPY-TO-callback-functions.patch (application/octet-stream, 4.9K) download
From 5c1ec351edb97c42b1161ce9717819281024fddf Mon Sep 17 00:00:00 2001
From: Bilva Sanaba <bilvas@amazon.com>
Date: Tue, 30 Jun 2020 16:08:07 -0700
Subject: [PATCH] Support COPY TO callback functions

---
 src/backend/commands/copy.c | 31 +++++++++++++++++--------------
 src/include/commands/copy.h |  7 +++++++
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 3e199bdfd0..2322b6d558 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -133,6 +133,7 @@ typedef struct CopyStateData
 	char	   *filename;		/* filename, or NULL for STDIN/STDOUT */
 	bool		is_program;		/* is 'filename' a program to popen? */
 	copy_data_source_cb data_source_cb; /* function for reading data */
+	copy_data_destination_cb data_destination_cb; /* function to write data */
 	bool		binary;			/* binary format? */
 	bool		freeze;			/* freeze rows on loading? */
 	bool		csv_mode;		/* Comma Separated Value format? */
@@ -356,11 +357,6 @@ static CopyState BeginCopy(ParseState *pstate, bool is_from, Relation rel,
 						   List *options);
 static void EndCopy(CopyState cstate);
 static void ClosePipeToProgram(CopyState cstate);
-static CopyState BeginCopyTo(ParseState *pstate, Relation rel, RawStmt *query,
-							 Oid queryRelId, const char *filename, bool is_program,
-							 List *attnamelist, List *options);
-static void EndCopyTo(CopyState cstate);
-static uint64 DoCopyTo(CopyState cstate);
 static uint64 CopyTo(CopyState cstate);
 static void CopyOneRowTo(CopyState cstate, TupleTableSlot *slot);
 static bool CopyReadLine(CopyState cstate);
@@ -586,7 +582,7 @@ CopySendEndOfRow(CopyState cstate)
 			(void) pq_putmessage('d', fe_msgbuf->data, fe_msgbuf->len);
 			break;
 		case COPY_CALLBACK:
-			Assert(false);		/* Not yet supported. */
+			cstate->data_destination_cb(fe_msgbuf->data, fe_msgbuf->len);
 			break;
 	}
 
@@ -1075,7 +1071,8 @@ DoCopy(ParseState *pstate, const CopyStmt *stmt,
 	{
 		cstate = BeginCopyTo(pstate, rel, query, relid,
 							 stmt->filename, stmt->is_program,
-							 stmt->attlist, stmt->options);
+							 stmt->attlist, stmt->options,
+							 NULL);
 		*processed = DoCopyTo(cstate);	/* copy from database to file */
 		EndCopyTo(cstate);
 	}
@@ -1817,7 +1814,7 @@ EndCopy(CopyState cstate)
 /*
  * Setup CopyState to read tuples from a table or a query for COPY TO.
  */
-static CopyState
+CopyState
 BeginCopyTo(ParseState *pstate,
 			Relation rel,
 			RawStmt *query,
@@ -1825,10 +1822,11 @@ BeginCopyTo(ParseState *pstate,
 			const char *filename,
 			bool is_program,
 			List *attnamelist,
-			List *options)
+			List *options,
+			copy_data_destination_cb data_destination_cb)
 {
 	CopyState	cstate;
-	bool		pipe = (filename == NULL);
+	bool		pipe = (filename == NULL) && (data_destination_cb == NULL);
 	MemoryContext oldcontext;
 
 	if (rel != NULL && rel->rd_rel->relkind != RELKIND_RELATION)
@@ -1873,7 +1871,12 @@ BeginCopyTo(ParseState *pstate,
 					   options);
 	oldcontext = MemoryContextSwitchTo(cstate->copycontext);
 
-	if (pipe)
+	if (data_destination_cb)
+	{
+	    cstate->copy_dest = COPY_CALLBACK;
+	    cstate->data_destination_cb = data_destination_cb;
+	}
+	else if (pipe)
 	{
 		Assert(!is_program);	/* the grammar does not allow this */
 		if (whereToSendOutput != DestRemote)
@@ -1953,10 +1956,10 @@ BeginCopyTo(ParseState *pstate,
  * This intermediate routine exists mainly to localize the effects of setjmp
  * so we don't need to plaster a lot of variables with "volatile".
  */
-static uint64
+uint64
 DoCopyTo(CopyState cstate)
 {
-	bool		pipe = (cstate->filename == NULL);
+	bool		pipe = (cstate->filename == NULL) && (cstate->data_destination_cb == NULL);
 	bool		fe_copy = (pipe && whereToSendOutput == DestRemote);
 	uint64		processed;
 
@@ -1988,7 +1991,7 @@ DoCopyTo(CopyState cstate)
 /*
  * Clean up storage and release resources for COPY TO.
  */
-static void
+void
 EndCopyTo(CopyState cstate)
 {
 	if (cstate->queryDesc != NULL)
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index c639833565..6617bacd87 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -22,6 +22,7 @@
 /* CopyStateData is private in commands/copy.c */
 typedef struct CopyStateData *CopyState;
 typedef int (*copy_data_source_cb) (void *outbuf, int minread, int maxread);
+typedef void (*copy_data_destination_cb) (void *data, int len);
 
 extern void DoCopy(ParseState *state, const CopyStmt *stmt,
 				   int stmt_location, int stmt_len,
@@ -41,4 +42,10 @@ extern uint64 CopyFrom(CopyState cstate);
 
 extern DestReceiver *CreateCopyDestReceiver(void);
 
+extern CopyState BeginCopyTo(ParseState *pstate, Relation rel, RawStmt *query,
+	Oid queryRelId, const char *filename, bool is_program,
+	List *attnamelist, List *options, copy_data_destination_cb data_destination_cb);
+extern uint64 DoCopyTo(CopyState cstate);
+extern void EndCopyTo(CopyState cstate);
+
 #endif							/* COPY_H */
-- 
2.22.0