1 files changed, 16 insertions, 21 deletions

diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index b0d9784b3fc..695b4afa861 100644
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -63,6 +63,7 @@ static TocEntry *getTocEntryByDumpId(ArchiveHandle *AH, DumpId id);
 static void _moveAfter(ArchiveHandle *AH, TocEntry *pos, TocEntry *te);
 static int	_discoverArchiveFormat(ArchiveHandle *AH);
 
+static int	RestoringToDB(ArchiveHandle *AH);
 static void dump_lo_buf(ArchiveHandle *AH);
 static void _write_msg(const char *modulename, const char *fmt, va_list ap);
 static void _die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt, va_list ap);
@@ -379,14 +380,7 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 							  te->tag);
 
 						/*
-						 * If we have a copy statement, use it. As of V1.3,
-						 * these are separate to allow easy import from
-						 * withing a database connection. Pre 1.3 archives can
-						 * not use DB connections and are sent to output only.
-						 *
-						 * For V1.3+, the table data MUST have a copy
-						 * statement so that we can go into appropriate mode
-						 * with libpq.
+						 * If we have a copy statement, use it.
 						 */
 						if (te->copyStmt && strlen(te->copyStmt) > 0)
 						{
@@ -396,7 +390,15 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 
 						(*AH->PrintTocDataPtr) (AH, te, ropt);
 
-						AH->writingCopyData = false;
+						/*
+						 * Terminate COPY if needed.
+						 */
+						if (AH->writingCopyData)
+						{
+							if (RestoringToDB(AH))
+								EndDBCopyMode(AH, te);
+							AH->writingCopyData = false;
+						}
 
 						_enableTriggersIfNecessary(AH, te, ropt);
 					}
@@ -1038,17 +1040,13 @@ ahprintf(ArchiveHandle *AH, const char *fmt,...)
 {
 	char	   *p = NULL;
 	va_list		ap;
-	int			bSize = strlen(fmt) + 256;		/* Should be enough */
+	int			bSize = strlen(fmt) + 256;		/* Usually enough */
 	int			cnt = -1;
 
 	/*
 	 * This is paranoid: deal with the possibility that vsnprintf is willing
-	 * to ignore trailing null
-	 */
-
-	/*
-	 * or returns > 0 even if string does not fit. It may be the case that it
-	 * returns cnt = bufsize
+	 * to ignore trailing null or returns > 0 even if string does not fit.
+	 * It may be the case that it returns cnt = bufsize.
 	 */
 	while (cnt < 0 || cnt >= (bSize - 1))
 	{
@@ -1128,7 +1126,7 @@ dump_lo_buf(ArchiveHandle *AH)
 
 
 /*
- *	Write buffer to the output file (usually stdout). This is user for
+ *	Write buffer to the output file (usually stdout). This is used for
  *	outputting 'restore' scripts etc. It is even possible for an archive
  *	format to create a custom output routine to 'fake' a restore if it
  *	wants to generate a script (see TAR output).
@@ -1180,7 +1178,7 @@ ahwrite(const void *ptr, size_t size, size_t nmemb, ArchiveHandle *AH)
 		 * connected then send it to the DB.
 		 */
 		if (RestoringToDB(AH))
-			return ExecuteSqlCommandBuf(AH, (void *) ptr, size * nmemb);		/* Always 1, currently */
+			return ExecuteSqlCommandBuf(AH, (const char *) ptr, size * nmemb);
 		else
 		{
 			res = fwrite((void *) ptr, size, nmemb, AH->OF);
@@ -1730,9 +1728,6 @@ _allocAH(const char *FileSpec, const ArchiveFormat fmt,
 	AH->mode = mode;
 	AH->compression = compression;
 
-	AH->pgCopyBuf = createPQExpBuffer();
-	AH->sqlBuf = createPQExpBuffer();
-
 	/* Open stdout with no compression for AH output handle */
 	AH->gzOut = 0;
 	AH->OF = stdout;