Fix for bug 3151626: Support for COPY BINARY

This implements support for BINARY format for both COPY TO and COPY FROM.
For COPY TO, PG_HEADER part is only generated by Coordinator,
and Datanodes only generate the data in itself.

Datanodes were generating the whole data, including header and send it back
to client directly.

Patch has been written by xabc1000, with some editorialization by me.

3 files changed, 144 insertions, 0 deletions

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 061d05ed55..3740fc8fa3 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -609,6 +609,13 @@ CopyGetData(CopyState cstate, void *databuf, int minread, int maxread)
 							break;
 					}
 				}
+#ifdef PGXC
+				/* A PGXC Datanode does not need to read the header data received from Coordinator */
+				if (IS_PGXC_DATANODE &&
+					cstate->binary &&
+					cstate->fe_msgbuf->data[cstate->fe_msgbuf->len-1] == '\n')
+					cstate->fe_msgbuf->len--;
+#endif
 				avail = cstate->fe_msgbuf->len - cstate->fe_msgbuf->cursor;
 				if (avail > maxread)
 					avail = maxread;
@@ -1551,6 +1558,10 @@ CopyTo(CopyState cstate)
 
 	if (cstate->binary)
 	{
+#ifdef PGXC
+	if (IS_PGXC_COORDINATOR)
+	{
+#endif
 		/* Generate header for a binary copy */
 		int32		tmp;
 
@@ -1564,6 +1575,12 @@ CopyTo(CopyState cstate)
 		/* No header extension */
 		tmp = 0;
 		CopySendInt32(cstate, tmp);
+
+#ifdef PGXC
+		/* Need to flush out the trailer */
+		CopySendEndOfRow(cstate);
+	}
+#endif
 	}
 	else
 	{
@@ -1646,7 +1663,15 @@ CopyTo(CopyState cstate)
 	}
 #endif
 
+#ifdef PGXC
+	/*
+	 * In PGXC, it is not necessary for a datanode to generate
+	 * the trailer as Coordinator is in charge of it
+	 */
+	if (cstate->binary && IS_PGXC_COORDINATOR)
+#else
 	if (cstate->binary)
+#endif
 	{
 		/* Generate trailer for a binary copy */
 		CopySendInt16(cstate, -1);
@@ -2099,6 +2124,31 @@ CopyFrom(CopyState cstate)
 						(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
 						 errmsg("invalid COPY file header (wrong length)")));
 		}
+#ifdef PGXC
+		if (IS_PGXC_COORDINATOR)
+		{
+			/* Empty buffer info and send header to all the backends involved in COPY */
+			resetStringInfo(&cstate->line_buf);
+
+			enlargeStringInfo(&cstate->line_buf, 19);
+			appendBinaryStringInfo(&cstate->line_buf, BinarySignature, 11);
+			tmp = 0;
+
+			if (cstate->oids)
+				tmp |= (1 << 16);
+			tmp = htonl(tmp);
+
+			appendBinaryStringInfo(&cstate->line_buf, &tmp, 4);
+			tmp = 0;
+			tmp = htonl(tmp);
+			appendBinaryStringInfo(&cstate->line_buf, &tmp, 4);
+
+			if(DataNodeCopyInBinaryForAll(cstate->line_buf.data, 19, cstate->connections))
+					ereport(ERROR,
+							(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
+							 errmsg("invalid COPY file header (COPY SEND)")));
+		}
+#endif
 	}
 
 	if (file_has_oids && cstate->binary)
@@ -2254,6 +2304,18 @@ CopyFrom(CopyState cstate)
 				fld_count == -1)
 			{
 				done = true;
+#ifdef PGXC
+				if (IS_PGXC_COORDINATOR)
+				{
+					/* Empty buffer */
+					resetStringInfo(&cstate->line_buf);
+
+					enlargeStringInfo(&cstate->line_buf, sizeof(uint16));
+					/* Receive field count directly from datanodes */
+					fld_count = htons(fld_count);
+					appendBinaryStringInfo(&cstate->line_buf, &fld_count, sizeof(uint16));
+				}
+#endif
 				break;
 			}
 
@@ -2262,7 +2324,17 @@ CopyFrom(CopyState cstate)
 						(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
 						 errmsg("row field count is %d, expected %d",
 								(int) fld_count, attr_count)));
+#ifdef PGXC
+			if (IS_PGXC_COORDINATOR)
+			{
+				/* Empty buffer */
+				resetStringInfo(&cstate->line_buf);
 
+				enlargeStringInfo(&cstate->line_buf, sizeof(uint16));
+				fld_count = htons(fld_count);
+				appendBinaryStringInfo(&cstate->line_buf, &fld_count, sizeof(uint16));
+			}
+#endif
 			if (file_has_oids)
 			{
 				cstate->cur_attname = "oid";
@@ -3302,6 +3374,7 @@ CopyReadBinaryAttribute(CopyState cstate,
 						bool *isnull)
 {
 	int32		fld_size;
+	int32 		nSize;
 	Datum		result;
 
 	if (!CopyGetInt32(cstate, &fld_size))
@@ -3317,6 +3390,15 @@ CopyReadBinaryAttribute(CopyState cstate,
 		ereport(ERROR,
 				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
 				 errmsg("invalid field size")));
+#ifdef PGXC
+	if (IS_PGXC_COORDINATOR)
+	{
+		/* Get the field size from Datanode */
+		enlargeStringInfo(&cstate->line_buf, sizeof(int32));
+		nSize = htonl(fld_size);
+		appendBinaryStringInfo(&cstate->line_buf, &nSize, sizeof(int32));
+	}
+#endif
 
 	/* reset attribute_buf to empty, and load raw data in it */
 	resetStringInfo(&cstate->attribute_buf);
@@ -3330,6 +3412,14 @@ CopyReadBinaryAttribute(CopyState cstate,
 
 	cstate->attribute_buf.len = fld_size;
 	cstate->attribute_buf.data[fld_size] = '\0';
+#ifdef PGXC
+	if (IS_PGXC_COORDINATOR)
+	{
+		/* Get binary message from Datanode */
+		enlargeStringInfo(&cstate->line_buf, fld_size);
+		appendBinaryStringInfo(&cstate->line_buf, cstate->attribute_buf.data, fld_size);
+	}
+#endif
 
 	/* Call the column type's binary input converter */
 	result = ReceiveFunctionCall(flinfo, &cstate->attribute_buf,
diff --git a/src/backend/pgxc/pool/execRemote.c b/src/backend/pgxc/pool/execRemote.c
index 1d41a5cfbc..dc1d68ce5a 100644
--- a/src/backend/pgxc/pool/execRemote.c
+++ b/src/backend/pgxc/pool/execRemote.c
@@ -4570,3 +4570,56 @@ PGXCNodeGetNodeList(PGXC_NodeId **datanodes,
 	if (!PersistentConnections)
 		release_handles();
 }
+
+/*
+ * DataNodeCopyInBinaryForAll
+ *
+ * In a COPY TO, send to all datanodes PG_HEADER for a COPY TO in binary mode.
+ */
+int DataNodeCopyInBinaryForAll(char *msg_buf, int len, PGXCNodeHandle** copy_connections)
+{
+	int 		i;
+	int 		conn_count = 0;
+	PGXCNodeHandle *connections[NumDataNodes];
+	int msgLen = 4 + len + 1;
+	int nLen = htonl(msgLen);
+
+	for (i = 0; i < NumDataNodes; i++)
+	{
+		PGXCNodeHandle *handle = copy_connections[i];
+
+		if (!handle)
+			continue;
+
+		connections[conn_count++] = handle;
+	}
+
+	for (i = 0; i < conn_count; i++)
+	{
+		PGXCNodeHandle *handle = connections[i];
+		if (handle->state == DN_CONNECTION_STATE_COPY_IN)
+		{
+			/* msgType + msgLen */
+			if (ensure_out_buffer_capacity(handle->outEnd + 1 + msgLen, handle) != 0)
+			{
+				ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					errmsg("out of memory")));
+			}
+
+			handle->outBuffer[handle->outEnd++] = 'd';
+			memcpy(handle->outBuffer + handle->outEnd, &nLen, 4);
+			handle->outEnd += 4;
+			memcpy(handle->outBuffer + handle->outEnd, msg_buf, len);
+			handle->outEnd += len;
+			handle->outBuffer[handle->outEnd++] = '\n';
+		}
+		else
+		{
+			add_error_message(handle, "Invalid data node connection");
+			return EOF;
+		}
+	}
+
+	return 0;
+}
diff --git a/src/include/pgxc/execRemote.h b/src/include/pgxc/execRemote.h
index aaa12a19f4..23ea4a816e 100644
--- a/src/include/pgxc/execRemote.h
+++ b/src/include/pgxc/execRemote.h
@@ -139,6 +139,7 @@ extern PGXCNodeHandle** DataNodeCopyBegin(const char *query, List *nodelist, Sna
 extern int DataNodeCopyIn(char *data_row, int len, ExecNodes *exec_nodes, PGXCNodeHandle** copy_connections);
 extern uint64 DataNodeCopyOut(ExecNodes *exec_nodes, PGXCNodeHandle** copy_connections, FILE* copy_file);
 extern void DataNodeCopyFinish(PGXCNodeHandle** copy_connections, int primary_data_node, CombineType combine_type);
+extern int DataNodeCopyInBinaryForAll(char *msg_buf, int len, PGXCNodeHandle** copy_connections);
 
 extern int ExecCountSlotsRemoteQuery(RemoteQuery *node);
 extern RemoteQueryState *ExecInitRemoteQuery(RemoteQuery *node, EState *estate, int eflags);