COPY FROM to postgres_fdw implementation for pg_shardman

Cherry-picked from af234df.

This is the commit message #2:

Tracking on which foreign servers we already started COPY.

If foreign server holds several partitions, COPY FROM to local root partition
will try to perform several copies at the same time through one connection,
obviously without much success. Now we track that and start/end COPY only once.

We also allow to pass destination relation name which may be different from
foreing table -- so we can copy into foreign root partition in shardman. This
is pretty narrow solution. However, keeping several connections to the same
foreign server requires significant changes, especially in 2pc handling, so
staying here for now.

This is the commit message #3:

Allow COPY FROM to par8d table even if some FDW parts can't do that.

This behaviour was broken in patches allowing COPY FROM to FDW tables.

This is the commit message #4:

COPY FROM deparse more complete, PG_SHARDMAN macro.

Now column names, FORCE NULL and FORCE NOT NULL are deparsed too.
PG_SHARDMAN macro ensures this PG contains patches for Postgres.

This is the commit message #5:

Disable COPY FROM to foreign parts, because no generic impl exists.

This is the commit message #6:

Fix COPY FROM deparse, forgotten comma for FORCE_NULL etc.

Author: arssher

contrib/postgres_fdw/connection.c

@@ -18,7 +18,9 @@
 #include "access/htup_details.h"
 #include "catalog/pg_user_mapping.h"
 #include "access/xact.h"
+#include "access/transam.h"
 #include "access/xlog.h" /* GetSystemIdentifier() */
+#include "libpq-int.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "pgstat.h"
@@ -63,7 +65,8 @@ struct ConnCacheEntry
 	bool		invalidated;	/* true if reconnect is pending */
 	uint32		server_hashvalue;	/* hash value of foreign server OID */
 	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
-};
+	bool		copy_from_started;	/* COPY FROM in progress on this conn */
+} ;
 
 /*
  * Connection cache (initialized on first use)
@@ -129,7 +132,8 @@ static bool pgfdw_get_cleanup_result(ConnCacheEntry *entry, TimestampTz endtime,
  * (not even on error), we need this flag to cue manual cleanup.
  */
 ConnCacheEntry *
-GetConnection(UserMapping *user, bool will_prep_stmt)
+GetConnectionCopyFrom(UserMapping *user, bool will_prep_stmt,
+					  bool **copy_from_started)
 {
 	bool		found;
 	ConnCacheEntry *entry;
@@ -224,6 +228,7 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 		entry->have_error = false;
 		entry->changing_xact_state = false;
 		entry->invalidated = false;
+		entry->copy_from_started = false;
 		entry->server_hashvalue =
 			GetSysCacheHashValue1(FOREIGNSERVEROID,
 								  ObjectIdGetDatum(server->serverid));
@@ -246,6 +251,9 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
 
+	if (copy_from_started)
+		*copy_from_started = &(entry->copy_from_started);
+
 	return entry;
 }
 
@@ -255,6 +263,12 @@ ConnectionEntryGetConn(ConnCacheEntry *entry)
 	return entry->conn;
 }
 
+ConnCacheEntry *
+GetConnection(UserMapping *user, bool will_prep_stmt)
+{
+	return GetConnectionCopyFrom(user, will_prep_stmt, NULL);
+}
+
 /*
  * Connect to remote server using specified server and user mapping properties.
  */
@@ -1292,21 +1306,40 @@ pgfdw_cancel_query(ConnCacheEntry *entry)
 	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"
@@ -3323,3 +3325,107 @@ get_relation_column_alias_ids(Var *node, RelOptInfo *foreignrel,
 	/* Shouldn't get here */
 	elog(ERROR, "unexpected expression in subquery output");
 }
+
+/*
+ * Deparse COPY FROM into given buf.
+ */
+void
+deparseCopyFromSql(StringInfo buf, Relation rel, CopyState cstate,
+				   const char *dest_relname)
+{
+	ListCell *cur;
+
+	appendStringInfoString(buf, "COPY ");
+	if (dest_relname == NULL)
+		deparseRelation(buf, rel);
+	else
+		appendStringInfoString(buf, dest_relname);
+
+	if (cstate->binary)
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot copy to postgres_fdw table \"%s\" in binary format ",
+						RelationGetRelationName(rel))));
+	}
+
+	/* deparse column names */
+	if (cstate->attnumlist != NIL)
+	{
+		bool first = true;
+
+		appendStringInfoString(buf, " (");
+		foreach(cur, cstate->attnumlist)
+		{
+			int attnum = lfirst_int(cur);
+			char *attname;
+
+			if (!first)
+				appendStringInfoString(buf, ", ");
+			first = false;
+
+			attname = get_attname(rel->rd_id, attnum, false);
+			appendStringInfoString(buf, quote_identifier(attname));
+		}
+		appendStringInfoString(buf, " )");
+	}
+
+	appendStringInfoString(buf, " FROM STDIN WITH (");
+	if (cstate->csv_mode)
+	{
+		appendStringInfoString(buf, " FORMAT csv ");
+		appendStringInfo(buf, ", QUOTE '%c'", *(cstate->quote));
+		appendStringInfo(buf, ", ESCAPE '%c'", *(cstate->escape));
+		if (cstate->force_notnull != NIL)
+		{
+			bool first = true;
+
+			appendStringInfoString(buf, ", FORCE_NOT_NULL (");
+			foreach(cur, cstate->force_notnull)
+			{
+				char *attname = strVal(lfirst(cur));
+
+				if (!first)
+					appendStringInfoString(buf, ", ");
+				first = false;
+
+				appendStringInfoString(buf, quote_identifier(attname));
+			}
+			appendStringInfoString(buf, " )");
+		}
+		if (cstate->force_null != NIL)
+		{
+			bool first = true;
+
+			appendStringInfoString(buf, ", FORCE_NULL (");
+			foreach(cur, cstate->force_null)
+			{
+				char *attname = strVal(lfirst(cur));
+
+				if (!first)
+					appendStringInfoString(buf, ", ");
+				first = false;
+
+				appendStringInfoString(buf, quote_identifier(attname));
+			}
+			appendStringInfoString(buf, " )");
+		}
+	}
+	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

@@ -362,6 +362,15 @@ static void postgresGetForeignUpperPaths(PlannerInfo *root,
 							 RelOptInfo *input_rel,
 							 RelOptInfo *output_rel,
 							 void *extra);
+static void postgresBeginForeignCopyFrom(EState *estate,
+										 ResultRelInfo *rinfo,
+										 CopyState cstate,
+										 ResultRelInfo *parent_rinfo);
+static void postgresForeignNextCopyFrom(EState *estate,
+										ResultRelInfo *rinfo,
+										CopyState cstate);
+static void postgresEndForeignCopyFrom(EState *estate,
+									   ResultRelInfo *rinfo);
 
 /*
  * Helper functions
@@ -509,6 +518,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);
 }
 
@@ -5821,6 +5835,114 @@ find_em_expr_for_rel(EquivalenceClass *ec, RelOptInfo *rel)
 	return NULL;
 }
 
+/*
+ * Begin COPY FROM to foreign table. Currently we do it in a bit perverted
+ * way: we redirect COPY FROM to parent table on foreign server, assuming it
+ * exists in public schema (as in shardman), and let it direct tuples to
+ * proper partitions. Otherwise we would have to modify logic of managing
+ * connections and keep many connections open to one server from one backend.
+ * This probably should not be used outside pg_shardman.
+ */
+static void
+postgresBeginForeignCopyFrom(EState *estate, ResultRelInfo *rinfo,
+							 CopyState cstate, ResultRelInfo *parent_rinfo)
+{
+	Relation		rel = rinfo->ri_RelationDesc;
+	RangeTblEntry	*rte;
+	Oid				userid;
+	ForeignTable	*table;
+	UserMapping		*user;
+	StringInfoData 	sql;
+	PGconn	   		*conn;
+	PGresult   		*res;
+	bool			*copy_from_started;
+	char 			*dest_relname;
+
+	/*
+	 * 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);
+	rinfo->ri_FdwState = user;
+
+	/* Get (open, if not yet) connection */
+	conn = ConnectionEntryGetConn(
+		GetConnectionCopyFrom(user, false, &copy_from_started));
+	/* We already did COPY FROM to this server */
+	if (*copy_from_started)
+		return;
+
+	/* deparse COPY stmt */
+	dest_relname = psprintf(
+		"public.%s", quote_identifier(RelationGetRelationName(
+										  parent_rinfo == NULL ?
+										  rinfo->ri_RelationDesc :
+										  parent_rinfo->ri_RelationDesc)));
+	initStringInfo(&sql);
+	deparseCopyFromSql(&sql, rel, cstate, dest_relname);
+
+	res = PQexec(conn, sql.data);
+	if (PQresultStatus(res) != PGRES_COPY_IN)
+	{
+		pgfdw_report_error(ERROR, res, conn, true, sql.data);
+	}
+	PQclear(res);
+
+	*copy_from_started = true;
+}
+
+/* COPY FROM next row to foreign table */
+static void
+postgresForeignNextCopyFrom(EState *estate, ResultRelInfo *rinfo,
+							CopyState cstate)
+{
+	bool		*copy_from_started;
+	UserMapping *user = (UserMapping *) rinfo->ri_FdwState;
+	PGconn		*conn = ConnectionEntryGetConn(
+		GetConnectionCopyFrom(user, false, &copy_from_started));
+
+	Assert(copy_from_started);
+	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)
+{
+	bool		*copy_from_started;
+	UserMapping *user = (UserMapping *) rinfo->ri_FdwState;
+	ConnCacheEntry *conn_entry = GetConnectionCopyFrom(user, false,
+													   &copy_from_started);
+	PGconn		*conn = ConnectionEntryGetConn(conn_entry);
+	PGresult	*res;
+
+	if (*copy_from_started)
+	{
+		/* 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);
+		}
+		*copy_from_started = false;
+		ReleaseConnection(conn_entry);
+	}
+}
+
 void
 _PG_init(void)
 {

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"
@@ -122,6 +123,8 @@ extern void reset_transmission_modes(int nestlevel);
 
 /* in connection.c */
 extern ConnCacheEntry *GetConnection(UserMapping *user, bool will_prep_stmt);
+extern ConnCacheEntry *GetConnectionCopyFrom(UserMapping *user, bool will_prep_stmt,
+											 bool **copy_from_started);
 extern PGconn *ConnectionEntryGetConn(ConnCacheEntry *entry);
 extern void ReleaseConnection(ConnCacheEntry *entry);
 extern unsigned int GetCursorNumber(ConnCacheEntry *entry);
@@ -186,6 +189,8 @@ 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,
+							   const char *dest_relname);
 
 /* in shippable.c */
 extern bool is_builtin(Oid objectId);