1 files changed, 24 insertions, 7 deletions

diff --git a/src/bin/pg_dump/pg_backup_db.c b/src/bin/pg_dump/pg_backup_db.c
index 785eaee304a..00426700431 100644
--- a/src/bin/pg_dump/pg_backup_db.c
+++ b/src/bin/pg_dump/pg_backup_db.c
@@ -90,6 +90,9 @@ ReconnectToServer(ArchiveHandle *AH, const char *dbname, const char *username)
 
 	newConn = _connectDB(AH, newdbname, newusername);
 
+	/* Update ArchiveHandle's connCancel before closing old connection */
+	set_archive_cancel_info(AH, newConn);
+
 	PQfinish(AH->connection);
 	AH->connection = newConn;
 
@@ -311,6 +314,9 @@ ConnectDatabase(Archive *AHX,
 	_check_database_version(AH);
 
 	PQsetNoticeProcessor(AH->connection, notice_processor, NULL);
+
+	/* arrange for SIGINT to issue a query cancel on this connection */
+	set_archive_cancel_info(AH, AH->connection);
 }
 
 /*
@@ -321,19 +327,25 @@ void
 DisconnectDatabase(Archive *AHX)
 {
 	ArchiveHandle *AH = (ArchiveHandle *) AHX;
-	PGcancel   *cancel;
 	char		errbuf[1];
 
 	if (!AH->connection)
 		return;
 
-	if (PQtransactionStatus(AH->connection) == PQTRANS_ACTIVE)
+	if (AH->connCancel)
 	{
-		if ((cancel = PQgetCancel(AH->connection)))
-		{
-			PQcancel(cancel, errbuf, sizeof(errbuf));
-			PQfreeCancel(cancel);
-		}
+		/*
+		 * If we have an active query, send a cancel before closing.  This is
+		 * of no use for a normal exit, but might be helpful during
+		 * exit_horribly().
+		 */
+		if (PQtransactionStatus(AH->connection) == PQTRANS_ACTIVE)
+			PQcancel(AH->connCancel, errbuf, sizeof(errbuf));
+
+		/*
+		 * Prevent signal handler from sending a cancel after this.
+		 */
+		set_archive_cancel_info(AH, NULL);
 	}
 
 	PQfinish(AH->connection);
@@ -588,6 +600,11 @@ EndDBCopyMode(ArchiveHandle *AH, TocEntry *te)
 								  te->tag, PQerrorMessage(AH->connection));
 		PQclear(res);
 
+		/* Do this to ensure we've pumped libpq back to idle state */
+		if (PQgetResult(AH->connection) != NULL)
+			write_msg(NULL, "WARNING: unexpected extra results during COPY of table \"%s\"\n",
+					  te->tag);
+
 		AH->pgCopyIn = false;
 	}
 }