Fix pg_restore to properly discard COPY data when trying to continue

after an error in a COPY statement.  Formerly it thought the COPY data
was SQL commands, and got quite confused.

Stephen Frost

diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index 99b055dcc0390c6240e76252bf05935814c46f16..11c1ce8271f1d99ad4984da4641b14ff81026308 100644 (file)
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -330,10 +330,15 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
                                                 * with libpq.
                                                 */
                                                if (te->copyStmt && strlen(te->copyStmt) > 0)
+                                               {
                                                        ahprintf(AH, "%s", te->copyStmt);
+                                                       AH->writingCopyData = true;
+                                               }
 
                                                (*AH->PrintTocDataPtr) (AH, te, ropt);
 
+                                               AH->writingCopyData = false;
+
                                                _enableTriggersIfNecessary(AH, te, ropt);
                                        }
                                }

diff --git a/src/bin/pg_dump/pg_backup_archiver.h b/src/bin/pg_dump/pg_backup_archiver.h
index 10e306648b18f2bc61b4fc3d2ed1d5169f3cc9bf..9bdf48d6d3f05889e01f32bb71ed708076fab515 100644 (file)
--- a/src/bin/pg_dump/pg_backup_archiver.h
+++ b/src/bin/pg_dump/pg_backup_archiver.h
@@ -239,7 +239,8 @@ typedef struct _archiveHandle
        PGconn     *connection;
        int                     connectToDB;    /* Flag to indicate if direct DB connection is
                                                                 * required */
-       int                     pgCopyIn;               /* Currently in libpq 'COPY IN' mode. */
+       bool            writingCopyData;        /* True when we are sending COPY data */
+       bool            pgCopyIn;               /* Currently in libpq 'COPY IN' mode. */
        PQExpBuffer pgCopyBuf;          /* Left-over data from incomplete lines in
                                                                 * COPY IN */
 

diff --git a/src/bin/pg_dump/pg_backup_db.c b/src/bin/pg_dump/pg_backup_db.c
index 0eb0d744e8b5cbacef9808f99e7fb7beb7f54cdd..29ed7f9d8ffd332c26a97afb0b1b61f1aaa5af38 100644 (file)
--- a/src/bin/pg_dump/pg_backup_db.c
+++ b/src/bin/pg_dump/pg_backup_db.c
@@ -302,7 +302,7 @@ ExecuteSqlCommand(ArchiveHandle *AH, PQExpBuffer qry, char *desc)
        {
                if (PQresultStatus(res) == PGRES_COPY_IN)
                {
-                       AH->pgCopyIn = 1;
+                       AH->pgCopyIn = true;
                }
                else
                {
@@ -383,13 +383,12 @@ _sendCopyLine(ArchiveHandle *AH, char *qry, char *eos)
        appendPQExpBuffer(AH->pgCopyBuf, "%s\n", qry);
        isEnd = (strcmp(AH->pgCopyBuf->data, "\\.\n") == 0);
 
-       /*---------
-        * fprintf(stderr, "Sending '%s' via
-        *              COPY (at end = %d)\n\n", AH->pgCopyBuf->data, isEnd);
-        *---------
+       /*
+        * Note that we drop the data on the floor if libpq has failed to
+        * enter COPY mode; this allows us to behave reasonably when trying
+        * to continue after an error in a COPY command.
         */
-
-       if (PQputline(AH->connection, AH->pgCopyBuf->data) != 0)
+       if (AH->pgCopyIn && PQputline(AH->connection, AH->pgCopyBuf->data) != 0)
                die_horribly(AH, modulename, "error returned by PQputline\n");
 
        resetPQExpBuffer(AH->pgCopyBuf);
@@ -400,10 +399,10 @@ _sendCopyLine(ArchiveHandle *AH, char *qry, char *eos)
 
        if (isEnd)
        {
-               if (PQendcopy(AH->connection) != 0)
+               if (AH->pgCopyIn && PQendcopy(AH->connection) != 0)
                        die_horribly(AH, modulename, "error returned by PQendcopy\n");
 
-               AH->pgCopyIn = 0;
+               AH->pgCopyIn = false;
        }
 
        return qry + loc + 1;
@@ -615,7 +614,18 @@ ExecuteSqlCommandBuf(ArchiveHandle *AH, void *qryv, size_t bufLen)
        /* Could switch between command and COPY IN mode at each line */
        while (qry < eos)
        {
-               if (AH->pgCopyIn)
+               /*
+                * If libpq is in CopyIn mode *or* if the archive structure shows we
+                * are sending COPY data, treat the data as COPY data.  The pgCopyIn
+                * check is only needed for backwards compatibility with ancient
+                * archive files that might just issue a COPY command without marking
+                * it properly.  Note that in an archive entry that has a copyStmt,
+                * all data up to the end of the entry will go to _sendCopyLine, and
+                * therefore will be dropped if libpq has failed to enter COPY mode.
+                * Also, if a "\." data terminator is found, anything remaining in the
+                * archive entry will be dropped.
+                */
+               if (AH->pgCopyIn || AH->writingCopyData)
                        qry = _sendCopyLine(AH, qry, eos);
                else
                        qry = _sendSQLLine(AH, qry, eos);