Merged https://github.com/arssher/postgresql/tree/foreign_copy_from.

Author: arssher

contrib/postgres_fdw/connection.c

@@ -19,6 +19,8 @@
 #include "access/xact.h"
 #include "access/xtm.h"
 #include "access/transam.h"
+#include "access/xlog.h"
+#include "libpq-int.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "pgstat.h"
@@ -1231,21 +1233,40 @@ pgfdw_cancel_query(PGconn *conn)
 	endtime = TimestampTzPlusMilliseconds(GetCurrentTimestamp(), 30000);
 
 	/*
-	 * Issue cancel request.  Unfortunately, there's no good way to limit the
-	 * amount of time that we might block inside PQgetCancel().
+	 * If COPY IN in progress, send CopyFail. Otherwise send cancel request.
+	 * TODO: make it less hackish, without libpq-int.h inclusion and handling
+	 * EAGAIN.
 	 */
-	if ((cancel = PQgetCancel(conn)))
+	if (conn->asyncStatus == PGASYNC_COPY_IN)
 	{
-		if (!PQcancel(cancel, errbuf, sizeof(errbuf)))
+		if (PQputCopyEnd(conn, "postgres_fdw: transaction abort on source node") != 1)
 		{
 			ereport(WARNING,
 					(errcode(ERRCODE_CONNECTION_FAILURE),
-					 errmsg("could not send cancel request: %s",
+					 errmsg("could not send abort copy request: %s",
 							errbuf)));
-			PQfreeCancel(cancel);
 			return false;
 		}
-		PQfreeCancel(cancel);
+	}
+	else
+	{
+		/*
+		 * Issue cancel request.  Unfortunately, there's no good way to limit the
+		 * amount of time that we might block inside PQgetCancel().
+		 */
+		if ((cancel = PQgetCancel(conn)))
+		{
+			if (!PQcancel(cancel, errbuf, sizeof(errbuf)))
+			{
+				ereport(WARNING,
+						(errcode(ERRCODE_CONNECTION_FAILURE),
+						 errmsg("could not send cancel request: %s",
+								errbuf)));
+				PQfreeCancel(cancel);
+				return false;
+			}
+			PQfreeCancel(cancel);
+		}
 	}
 
 	/* Get and discard the result of the query. */

contrib/postgres_fdw/deparse.c

@@ -45,6 +45,8 @@
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "commands/defrem.h"
+#include "commands/copy.h"
+#include "mb/pg_wchar.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "nodes/plannodes.h"
@@ -3172,3 +3174,47 @@ get_relation_column_alias_ids(Var *node, RelOptInfo *foreignrel,
 	/* Shouldn't get here */
 	elog(ERROR, "unexpected expression in subquery output");
 }
+
+/*
+ * Deparse COPY FROM
+ */
+void
+deparseCopyFromSql(StringInfo buf, Relation rel, CopyState cstate)
+{
+	appendStringInfoString(buf, "COPY ");
+	deparseRelation(buf, rel);
+	appendStringInfoString(buf, " FROM STDIN WITH (");
+
+	/* TODO: deparse column names */
+	if (cstate->binary)
+	{
+			ereport(ERROR,
+					(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+					 errmsg("cannot copy to postgres_fdw table \"%s\" in binary format ",
+							RelationGetRelationName(rel))));
+	}
+	if (cstate->csv_mode)
+	{
+		appendStringInfoString(buf, " FORMAT csv ");
+		appendStringInfo(buf, ", QUOTE '%c'", *(cstate->quote));
+		appendStringInfo(buf, ", ESCAPE '%c'", *(cstate->escape));
+		/* TODO: force quote, force not null, force null */
+	}
+	else
+	{
+		appendStringInfoString(buf, " FORMAT text ");
+	}
+
+	appendStringInfo(buf, ", OIDS %d", cstate->oids);
+	appendStringInfo(buf, ", FREEZE %d", cstate->freeze);
+	appendStringInfo(buf, ", DELIMITER '%c'", *(cstate->delim));
+	appendStringInfo(buf, ", NULL %s", quote_literal_cstr(cstate->null_print));
+	/*
+	 * cstate->line_buf is passed to us already converted to this server
+	 * encoding.
+	 */
+	appendStringInfo(buf, ", ENCODING %s",
+					 quote_literal_cstr(
+						 pg_encoding_to_char(GetDatabaseEncoding())));
+	appendStringInfoChar(buf, ')');
+}

contrib/postgres_fdw/postgres_fdw.c

@@ -352,6 +352,14 @@ static void postgresGetForeignUpperPaths(PlannerInfo *root,
 							 UpperRelationKind stage,
 							 RelOptInfo *input_rel,
 							 RelOptInfo *output_rel);
+static void postgresBeginForeignCopyFrom(EState *estate,
+										 ResultRelInfo *rinfo,
+										 CopyState cstate);
+static void postgresForeignNextCopyFrom(EState *estate,
+										ResultRelInfo *rinfo,
+										CopyState cstate);
+static void postgresEndForeignCopyFrom(EState *estate,
+									   ResultRelInfo *rinfo);
 
 /*
  * Helper functions
@@ -476,6 +484,11 @@ postgres_fdw_handler(PG_FUNCTION_ARGS)
 	/* Support functions for upper relation push-down */
 	routine->GetForeignUpperPaths = postgresGetForeignUpperPaths;
 
+	/* Functions for COPY FROM */
+	routine->BeginForeignCopyFrom = postgresBeginForeignCopyFrom;
+	routine->ForeignNextCopyFrom = postgresForeignNextCopyFrom;
+	routine->EndForeignCopyFrom = postgresEndForeignCopyFrom;
+
 	PG_RETURN_POINTER(routine);
 }
 
@@ -5208,3 +5221,80 @@ _PG_init(void)
 						  &UseTsDtmTransactions, false, PGC_USERSET, 0, NULL,
 							 NULL, NULL);
 }
+
+/* Begin COPY FROM to foreign table */
+static void
+postgresBeginForeignCopyFrom(EState *estate, ResultRelInfo *rinfo,
+							 CopyState cstate)
+{
+	Relation		rel = rinfo->ri_RelationDesc;
+	RangeTblEntry	*rte;
+	Oid				userid;
+	ForeignTable	*table;
+	UserMapping		*user;
+	StringInfoData 	sql;
+	PGconn	   *conn;
+	PGresult   *res;
+
+	/*
+	 * Identify which user to do the remote access as.  This should match what
+	 * ExecCheckRTEPerms() does.
+	 */
+	rte = rt_fetch(rinfo->ri_RangeTableIndex, estate->es_range_table);
+	userid = rte->checkAsUser ? rte->checkAsUser : GetUserId();
+
+	/* Get info about foreign table. */
+	table = GetForeignTable(RelationGetRelid(rel));
+	user = GetUserMapping(userid, table->serverid);
+
+	/* Open connection */
+	conn = GetConnection(user, false);
+	rinfo->ri_FdwState = conn;
+
+	/* deparse COPY stmt */
+	initStringInfo(&sql);
+	deparseCopyFromSql(&sql, rel, cstate);
+
+	res = PQexec(conn, sql.data);
+	if (PQresultStatus(res) != PGRES_COPY_IN)
+	{
+		pgfdw_report_error(ERROR, res, conn, true, sql.data);
+	}
+	PQclear(res);
+}
+
+/* COPY FROM next row to foreign table */
+static void
+postgresForeignNextCopyFrom(EState *estate, ResultRelInfo *rinfo,
+							CopyState cstate)
+{
+	PGconn	   *conn = (PGconn *) rinfo->ri_FdwState;
+
+	Assert(!cstate->binary);
+	/* TODO: distinuish failure and nonblocking-send EAGAIN */
+	if (PQputline(conn, cstate->line_buf.data) || PQputnbytes(conn, "\n", 1))
+	{
+		pgfdw_report_error(ERROR, NULL, conn, false, cstate->line_buf.data);
+	}
+}
+
+/* Finish COPY FROM */
+static void
+postgresEndForeignCopyFrom(EState *estate, ResultRelInfo *rinfo)
+{
+	PGconn	   *conn = (PGconn *) rinfo->ri_FdwState;
+	PGresult   *res;
+
+	/* TODO: PQgetResult? */
+	if (PQendcopy(conn))
+	{
+		pgfdw_report_error(ERROR, NULL, conn, false, "end postgres_fdw copy from");
+	}
+	while ((res = PQgetResult(conn)) != NULL)
+	{
+		/* TODO: get error? */
+		PQclear(res);
+	}
+
+	ReleaseConnection(conn);
+}

contrib/postgres_fdw/postgres_fdw.h

@@ -14,6 +14,7 @@
 #define POSTGRES_FDW_H
 
 #include "foreign/foreign.h"
+#include "commands/copy.h"
 #include "lib/stringinfo.h"
 #include "nodes/relation.h"
 #include "utils/relcache.h"
@@ -177,6 +178,7 @@ extern void deparseSelectStmtForRel(StringInfo buf, PlannerInfo *root,
 						List *remote_conds, List *pathkeys, bool is_subquery,
 						List **retrieved_attrs, List **params_list);
 extern const char *get_jointype_name(JoinType jointype);
+extern void deparseCopyFromSql(StringInfo buf, Relation rel, CopyState cstate);
 
 /* in shippable.c */
 extern bool is_builtin(Oid objectId);