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.

Author: arssher

contrib/postgres_fdw/connection.c

@@ -62,6 +62,7 @@ typedef 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 */
 } ConnCacheEntry;
 
 /*
@@ -114,7 +115,8 @@ static bool pgfdw_get_cleanup_result(PGconn *conn, TimestampTz endtime,
  * (not even on error), we need this flag to cue manual cleanup.
  */
 PGconn *
-GetConnection(UserMapping *user, bool will_prep_stmt)
+GetConnectionCopyFrom(UserMapping *user, bool will_prep_stmt,
+					  bool **copy_from_started)
 {
 	bool		found;
 	ConnCacheEntry *entry;
@@ -200,6 +202,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));
@@ -222,9 +225,17 @@ 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->conn;
 }
 
+PGconn *
+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.
  */

contrib/postgres_fdw/deparse.c

@@ -3179,10 +3179,14 @@ get_relation_column_alias_ids(Var *node, RelOptInfo *foreignrel,
  * Deparse COPY FROM
  */
 void
-deparseCopyFromSql(StringInfo buf, Relation rel, CopyState cstate)
+deparseCopyFromSql(StringInfo buf, Relation rel, CopyState cstate,
+				   const char *dest_relname)
 {
 	appendStringInfoString(buf, "COPY ");
-	deparseRelation(buf, rel);
+	if (dest_relname == NULL)
+		deparseRelation(buf, rel);
+	else
+		appendStringInfoString(buf, dest_relname);
 	appendStringInfoString(buf, " FROM STDIN WITH (");
 
 	/* TODO: deparse column names */

contrib/postgres_fdw/postgres_fdw.c

@@ -354,7 +354,8 @@ static void postgresGetForeignUpperPaths(PlannerInfo *root,
 							 RelOptInfo *output_rel);
 static void postgresBeginForeignCopyFrom(EState *estate,
 										 ResultRelInfo *rinfo,
-										 CopyState cstate);
+										 CopyState cstate,
+										 const char *dest_relname);
 static void postgresForeignNextCopyFrom(EState *estate,
 										ResultRelInfo *rinfo,
 										CopyState cstate);
@@ -5225,7 +5226,7 @@ _PG_init(void)
 /* Begin COPY FROM to foreign table */
 static void
 postgresBeginForeignCopyFrom(EState *estate, ResultRelInfo *rinfo,
-							 CopyState cstate)
+							 CopyState cstate, const char *dest_relname)
 {
 	Relation		rel = rinfo->ri_RelationDesc;
 	RangeTblEntry	*rte;
@@ -5235,6 +5236,7 @@ postgresBeginForeignCopyFrom(EState *estate, ResultRelInfo *rinfo,
 	StringInfoData 	sql;
 	PGconn	   *conn;
 	PGresult   *res;
+	bool *copy_from_started;
 
 	/*
 	 * Identify which user to do the remote access as.  This should match what
@@ -5246,30 +5248,38 @@ postgresBeginForeignCopyFrom(EState *estate, ResultRelInfo *rinfo,
 	/* Get info about foreign table. */
 	table = GetForeignTable(RelationGetRelid(rel));
 	user = GetUserMapping(userid, table->serverid);
+	rinfo->ri_FdwState = user;
 
-	/* Open connection */
-	conn = GetConnection(user, false);
-	rinfo->ri_FdwState = conn;
+	/* Get (open, if not yet) connection */
+	conn = GetConnectionCopyFrom(user, false, &copy_from_started);
+	/* We already did COPY FROM to this server */
+	if (*copy_from_started)
+		return;
 
 	/* deparse COPY stmt */
 	initStringInfo(&sql);
-	deparseCopyFromSql(&sql, rel, cstate);
+	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)
 {
-	PGconn	   *conn = (PGconn *) rinfo->ri_FdwState;
+	bool		*copy_from_started;
+	UserMapping *user = (UserMapping *) rinfo->ri_FdwState;
+	PGconn		*conn = 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))
@@ -5282,19 +5292,24 @@ postgresForeignNextCopyFrom(EState *estate, ResultRelInfo *rinfo,
 static void
 postgresEndForeignCopyFrom(EState *estate, ResultRelInfo *rinfo)
 {
-	PGconn	   *conn = (PGconn *) rinfo->ri_FdwState;
-	PGresult   *res;
+	bool		*copy_from_started;
+	UserMapping *user = (UserMapping *) rinfo->ri_FdwState;
+	PGconn		*conn = GetConnectionCopyFrom(user, false, &copy_from_started);
+	PGresult	*res;
 
-	/* TODO: PQgetResult? */
-	if (PQendcopy(conn))
+	if (*copy_from_started)
 	{
-		pgfdw_report_error(ERROR, NULL, conn, false, "end postgres_fdw copy from");
-	}
-	while ((res = PQgetResult(conn)) != NULL)
-	{
-		/* TODO: get error? */
-		PQclear(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);
+		}
+		*copy_from_started = false;
+		ReleaseConnection(conn);
 	}
-
-	ReleaseConnection(conn);
 }

contrib/postgres_fdw/postgres_fdw.h

@@ -117,6 +117,8 @@ extern void reset_transmission_modes(int nestlevel);
 
 /* in connection.c */
 extern PGconn *GetConnection(UserMapping *user, bool will_prep_stmt);
+extern PGconn *GetConnectionCopyFrom(UserMapping *user, bool will_prep_stmt,
+									 bool **copy_from_started);
 extern void ReleaseConnection(PGconn *conn);
 extern unsigned int GetCursorNumber(PGconn *conn);
 extern unsigned int GetPrepStmtNumber(PGconn *conn);
@@ -178,7 +180,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);
+extern void deparseCopyFromSql(StringInfo buf, Relation rel, CopyState cstate,
+							   const char *dest_relname);
 
 /* in shippable.c */
 extern bool is_builtin(Oid objectId);

src/backend/commands/copy.c

@@ -181,7 +181,7 @@ static bool CopyGetInt32(CopyState cstate, int32 *val);
 static void CopySendInt16(CopyState cstate, int16 val);
 static bool CopyGetInt16(CopyState cstate, int16 *val);
 static void InitForeignCopyFrom(EState *estate, ResultRelInfo *resultRelInfo,
-								CopyState cstate);
+								CopyState cstate, char *dest_relname);
 
 
 /*
@@ -2354,12 +2354,12 @@ CopyFrom(CopyState cstate)
 		/* If some partitions are foreign tables, init copy on remote end */
 		for (i = 0; i < num_partitions; i++)
 		{
-			InitForeignCopyFrom(estate, partitions + i, cstate);
+			InitForeignCopyFrom(estate, partitions + i, cstate, NULL);
 		}
 	}
 
 	/* If we are copying to foreign table, init it */
-	InitForeignCopyFrom(estate, resultRelInfo, cstate);
+	InitForeignCopyFrom(estate, resultRelInfo, cstate, NULL);
 
 	/*
 	 * It's more efficient to prepare a bunch of tuples for insertion, and
@@ -4727,7 +4727,7 @@ CreateCopyDestReceiver(void)
 }
 
 static void InitForeignCopyFrom(EState *estate, ResultRelInfo *resultRelInfo,
-								CopyState cstate)
+								CopyState cstate, char *dest_relname)
 {
 	if (resultRelInfo->ri_FdwRoutine)
 	{
@@ -4739,6 +4739,6 @@ static void InitForeignCopyFrom(EState *estate, ResultRelInfo *resultRelInfo,
 					 errmsg("FDW adapter for relation \"%s\" doesn't support COPY FROM",
 							RelationGetRelationName(resultRelInfo->ri_RelationDesc))));
 		resultRelInfo->ri_FdwRoutine->
-			BeginForeignCopyFrom(estate, resultRelInfo, cstate);
+			BeginForeignCopyFrom(estate, resultRelInfo, cstate, dest_relname);
 	}
 }

src/include/foreign/fdwapi.h

@@ -162,7 +162,8 @@ typedef bool (*IsForeignScanParallelSafe_function) (PlannerInfo *root,
 
 typedef void (*BeginForeignCopyFrom_function) (EState *estate,
 											   ResultRelInfo *rinfo,
-											   CopyState cstate);
+											   CopyState cstate,
+											   const char *dest_relname);
 /*
  * Currently we support only text and csv format and pass each row in
  * cstate->line_buf. We should also pass binary data and/or deformed tuple.