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 63e62e42693136bba7e1f028567b3acfc8f5a1f3..91c7006b192701a147e4a72586bf8d8e2533f359 100644 (file)
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -345,10 +345,15 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
                                                 * mode 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;
+
                                                /*
                                                 * If we just restored blobs, fix references in
                                                 * previously-loaded tables; otherwise, if we

diff --git a/src/bin/pg_dump/pg_backup_archiver.h b/src/bin/pg_dump/pg_backup_archiver.h
index fe9bbe04ff1a2c0849bf54cf4dbc33d0397def9d..29590aa6027b4cce525f315ceda8f58f805bff4d 100644 (file)
--- a/src/bin/pg_dump/pg_backup_archiver.h
+++ b/src/bin/pg_dump/pg_backup_archiver.h
@@ -164,6 +164,13 @@ typedef enum
        STAGE_FINALIZING
 } ArchiverStage;
 
+typedef enum
+{
+       REQ_SCHEMA = 1,
+       REQ_DATA = 2,
+       REQ_ALL = REQ_SCHEMA + REQ_DATA
+} teReqs;
+
 typedef struct _archiveHandle
 {
        Archive         public;                 /* Public part of archive */
@@ -237,7 +244,8 @@ typedef struct _archiveHandle
        int                     blobTxActive;   /* Flag set if TX active on blobConnection */
        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 */
 
@@ -321,7 +329,7 @@ extern void WriteToc(ArchiveHandle *AH);
 extern void ReadToc(ArchiveHandle *AH);
 extern void WriteDataChunks(ArchiveHandle *AH);
 
-extern int     TocIDRequired(ArchiveHandle *AH, DumpId id, RestoreOptions *ropt);
+extern teReqs TocIDRequired(ArchiveHandle *AH, DumpId id, RestoreOptions *ropt);
 extern bool checkSeek(FILE *fp);
 
 /*

diff --git a/src/bin/pg_dump/pg_backup_db.c b/src/bin/pg_dump/pg_backup_db.c
index b412cf0e0522abaa07995b40aae0eaf9954a59b5..2d49cbed1fe34f901efe9fd04bf33b368154686c 100644 (file)
--- a/src/bin/pg_dump/pg_backup_db.c
+++ b/src/bin/pg_dump/pg_backup_db.c
@@ -319,7 +319,7 @@ _executeSqlCommand(ArchiveHandle *AH, PGconn *conn, PQExpBuffer qry, char *desc)
                        if (conn != AH->connection)
                                die_horribly(AH, modulename, "COPY command executed in non-primary connection\n");
 
-                       AH->pgCopyIn = 1;
+                       AH->pgCopyIn = true;
                }
                else
                {
@@ -406,7 +406,7 @@ _sendCopyLine(ArchiveHandle *AH, char *qry, char *eos)
         *---------
         */
 
-       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);
@@ -417,10 +417,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;
@@ -658,7 +658,10 @@ 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 we are in CopyIn mode *or* if the upper-layers believe we're doing
+                * a COPY then call sendCopyLine.  If we're not actually in CopyIn mode
+                * the sendCopyLine will just dump the data coming in */
+               if (AH->pgCopyIn || AH->writingCopyData)
                        qry = _sendCopyLine(AH, qry, eos);
                else
                        qry = _sendSQLLine(AH, qry, eos);