OSDN Git Service

Improve pg_dump and psql to use libpq's newer COPY support routines,
authorTom Lane <tgl@sss.pgh.pa.us>
Fri, 3 Mar 2006 23:38:30 +0000 (23:38 +0000)
committerTom Lane <tgl@sss.pgh.pa.us>
Fri, 3 Mar 2006 23:38:30 +0000 (23:38 +0000)
instead of the old deprecated ones.
Volkan Yazici, with some editorializing by moi.

src/bin/pg_dump/pg_backup_db.c
src/bin/pg_dump/pg_dump.c
src/bin/psql/common.c
src/bin/psql/copy.c

index 531ddb9..77c4086 100644 (file)
@@ -5,7 +5,7 @@
  *     Implements the basic DB functions used by the archiver.
  *
  * IDENTIFICATION
- *       $PostgreSQL: pgsql/src/bin/pg_dump/pg_backup_db.c,v 1.69 2006/02/12 06:11:50 momjian Exp $
+ *       $PostgreSQL: pgsql/src/bin/pg_dump/pg_backup_db.c,v 1.70 2006/03/03 23:38:29 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -391,22 +391,29 @@ _sendCopyLine(ArchiveHandle *AH, char *qry, char *eos)
         * enter COPY mode; this allows us to behave reasonably when trying
         * to continue after an error in a COPY command.
         */
-       if (AH->pgCopyIn && PQputline(AH->connection, AH->pgCopyBuf->data) != 0)
-               die_horribly(AH, modulename, "error returned by PQputline: %s",
+       if (AH->pgCopyIn &&
+               PQputCopyData(AH->connection, AH->pgCopyBuf->data,
+                                         AH->pgCopyBuf->len) <= 0)
+               die_horribly(AH, modulename, "error returned by PQputCopyData: %s",
                                         PQerrorMessage(AH->connection));
 
        resetPQExpBuffer(AH->pgCopyBuf);
 
-       /*
-        * fprintf(stderr, "Buffer is '%s'\n", AH->pgCopyBuf->data);
-        */
-
-       if (isEnd)
+       if (isEnd && AH->pgCopyIn)
        {
-               if (AH->pgCopyIn && PQendcopy(AH->connection) != 0)
-                       die_horribly(AH, modulename, "error returned by PQendcopy: %s",
+               PGresult   *res;
+
+               if (PQputCopyEnd(AH->connection, NULL) <= 0)
+                       die_horribly(AH, modulename, "error returned by PQputCopyEnd: %s",
                                                 PQerrorMessage(AH->connection));
 
+               /* Check command status and return to normal libpq state */
+               res = PQgetResult(AH->connection);
+               if (PQresultStatus(res) != PGRES_COMMAND_OK)
+                       warn_or_die_horribly(AH, modulename, "COPY failed: %s",
+                                                                PQerrorMessage(AH->connection));
+               PQclear(res);
+
                AH->pgCopyIn = false;
        }
 
index 4eb2caf..535104d 100644 (file)
@@ -12,7 +12,7 @@
  *     by PostgreSQL
  *
  * IDENTIFICATION
- *       $PostgreSQL: pgsql/src/bin/pg_dump/pg_dump.c,v 1.431 2006/03/02 01:18:25 tgl Exp $
+ *       $PostgreSQL: pgsql/src/bin/pg_dump/pg_dump.c,v 1.432 2006/03/03 23:38:29 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -831,8 +831,6 @@ selectDumpableObject(DumpableObject *dobj)
  *       to be dumped.
  */
 
-#define COPYBUFSIZ             8192
-
 static int
 dumpTableData_copy(Archive *fout, void *dcontext)
 {
@@ -844,8 +842,7 @@ dumpTableData_copy(Archive *fout, void *dcontext)
        PQExpBuffer q = createPQExpBuffer();
        PGresult   *res;
        int                     ret;
-       bool            copydone;
-       char            copybuf[COPYBUFSIZ];
+       char       *copybuf;
        const char *column_list;
 
        if (g_verbose)
@@ -886,33 +883,19 @@ dumpTableData_copy(Archive *fout, void *dcontext)
        }
        res = PQexec(g_conn, q->data);
        check_sql_result(res, g_conn, q->data, PGRES_COPY_OUT);
+       PQclear(res);
 
-       copydone = false;
-
-       while (!copydone)
+       for (;;)
        {
-               ret = PQgetline(g_conn, copybuf, COPYBUFSIZ);
+               ret = PQgetCopyData(g_conn, &copybuf, 0);
+
+               if (ret < 0)
+                       break;                          /* done or error */
 
-               if (copybuf[0] == '\\' &&
-                       copybuf[1] == '.' &&
-                       copybuf[2] == '\0')
+               if (copybuf)
                {
-                       copydone = true;        /* don't print this... */
-               }
-               else
-               {
-                       archputs(copybuf, fout);
-                       switch (ret)
-                       {
-                               case EOF:
-                                       copydone = true;
-                                       /* FALLTHROUGH */
-                               case 0:
-                                       archputs("\n", fout);
-                                       break;
-                               case 1:
-                                       break;
-                       }
+                       WriteData(fout, copybuf, ret);
+                       PQfreemem(copybuf);
                }
 
                /*
@@ -920,7 +903,7 @@ dumpTableData_copy(Archive *fout, void *dcontext)
                 *
                 * There was considerable discussion in late July, 2000 regarding
                 * slowing down pg_dump when backing up large tables. Users with both
-                * slow & fast (muti-processor) machines experienced performance
+                * slow & fast (multi-processor) machines experienced performance
                 * degradation when doing a backup.
                 *
                 * Initial attempts based on sleeping for a number of ms for each ms
@@ -957,16 +940,20 @@ dumpTableData_copy(Archive *fout, void *dcontext)
        }
        archprintf(fout, "\\.\n\n\n");
 
-       ret = PQendcopy(g_conn);
-       if (ret != 0)
+       if (ret == -2)
        {
-               write_msg(NULL, "SQL command to dump the contents of table \"%s\" failed: PQendcopy() failed.\n", classname);
+               /* copy data transfer failed */
+               write_msg(NULL, "Dumping the contents of table \"%s\" failed: PQgetCopyData() failed.\n", classname);
                write_msg(NULL, "Error message from server: %s", PQerrorMessage(g_conn));
                write_msg(NULL, "The command was: %s\n", q->data);
                exit_nicely();
        }
 
+       /* Check command status and return to normal libpq state */
+       res = PQgetResult(g_conn);
+       check_sql_result(res, g_conn, q->data, PGRES_COMMAND_OK);
        PQclear(res);
+
        destroyPQExpBuffer(q);
        return 1;
 }
index 8f6d035..7e436d5 100644 (file)
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2005, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.112 2006/02/12 03:30:21 tgl Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.113 2006/03/03 23:38:30 tgl Exp $
  */
 #include "postgres_fe.h"
 #include "common.h"
@@ -685,7 +685,10 @@ AcceptResult(const PGresult *result, const char *query)
                                break;
 
                        case PGRES_COPY_OUT:
-                               /* keep cancel connection for copy out state */
+                               /*
+                                * Keep cancel connection active during copy out state.
+                                * The matching ResetCancelConn() is in handleCopyOut.
+                                */
                                SetCancelConn();
                                break;
 
@@ -702,6 +705,7 @@ AcceptResult(const PGresult *result, const char *query)
                        psql_error("%s", error);
 
                ReportSyntaxErrorPosition(result, query);
+
                CheckConnection();
        }
 
@@ -720,6 +724,9 @@ AcceptResult(const PGresult *result, const char *query)
  * is true; nothing special is done when start_xact is false.  Typically,
  * start_xact = false is used for SELECTs and explicit BEGIN/COMMIT commands.
  *
+ * Caller is responsible for handling the ensuing processing if a COPY
+ * command is sent.
+ *
  * Note: we don't bother to check PQclientEncoding; it is assumed that no
  * caller uses this path to issue "SET CLIENT_ENCODING".
  */
index bd1763c..6dd1b52 100644 (file)
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2005, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/copy.c,v 1.58 2005/10/15 02:49:40 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/copy.c,v 1.59 2006/03/03 23:38:30 tgl Exp $
  */
 #include "postgres_fe.h"
 #include "copy.h"
@@ -568,7 +568,9 @@ do_copy(const char *args)
                        break;
                default:
                        success = false;
-                       psql_error("\\copy: unexpected response (%d)\n", PQresultStatus(result));
+                       psql_error("\\copy: unexpected response (%d)\n",
+                                          PQresultStatus(result));
+                       break;
        }
 
        PQclear(result);
@@ -586,85 +588,92 @@ do_copy(const char *args)
 }
 
 
-#define COPYBUFSIZ 8192                        /* size doesn't matter */
-
+/*
+ * Functions for handling COPY IN/OUT data transfer.
+ *
+ * If you want to use COPY TO STDOUT/FROM STDIN in your application,
+ * this is the code to steal ;)
+ */
 
 /*
  * handleCopyOut
- * receives data as a result of a COPY ... TO stdout command
+ * receives data as a result of a COPY ... TO STDOUT command
  *
- * If you want to use COPY TO in your application, this is the code to steal :)
+ * conn should be a database connection that you just issued COPY TO on
+ * and got back a PGRES_COPY_OUT result.
+ * copystream is the file stream for the data to go to.
  *
- * conn should be a database connection that you just called COPY TO on
- * (and which gave you PGRES_COPY_OUT back);
- * copystream is the file stream you want the output to go to
+ * result is true if successful, false if not.
  */
 bool
 handleCopyOut(PGconn *conn, FILE *copystream)
 {
-       bool            copydone = false;               /* haven't started yet */
-       char            copybuf[COPYBUFSIZ];
-       int                     ret;
+       bool    OK = true;
+       char    *buf;
+       int              ret;
+       PGresult *res;
 
-       while (!copydone)
+       for (;;)
        {
-               ret = PQgetline(conn, copybuf, COPYBUFSIZ);
+               ret = PQgetCopyData(conn, &buf, 0);
 
-               if (copybuf[0] == '\\' &&
-                       copybuf[1] == '.' &&
-                       copybuf[2] == '\0')
-               {
-                       copydone = true;        /* we're at the end */
-               }
-               else
+               if (ret < 0)
+                       break;                          /* done or error */
+
+               if (buf)
                {
-                       fputs(copybuf, copystream);
-                       switch (ret)
-                       {
-                               case EOF:
-                                       copydone = true;
-                                       /* FALLTHROUGH */
-                               case 0:
-                                       fputc('\n', copystream);
-                                       break;
-                               case 1:
-                                       break;
-                       }
+                       fputs(buf, copystream);
+                       PQfreemem(buf);
                }
        }
+
        fflush(copystream);
-       ret = !PQendcopy(conn);
-       ResetCancelConn();
-       return ret;
-}
 
+       if (ret == -2)
+       {
+               psql_error("COPY data transfer failed: %s", PQerrorMessage(conn));
+               OK = false;
+       }
+
+       /* Check command status and return to normal libpq state */
+       res = PQgetResult(conn);
+       if (PQresultStatus(res) != PGRES_COMMAND_OK)
+       {
+               psql_error("%s", PQerrorMessage(conn));
+               OK = false;
+       }
+       PQclear(res);
 
+       /* Disable cancel connection (see AcceptResult in common.c) */
+       ResetCancelConn();
+       
+       return OK;
+}
 
 /*
  * handleCopyIn
- * receives data as a result of a COPY ... FROM stdin command
+ * sends data to complete a COPY ... FROM STDIN command
  *
- * Again, if you want to use COPY FROM in your application, copy this.
+ * conn should be a database connection that you just issued COPY FROM on
+ * and got back a PGRES_COPY_IN result.
+ * copystream is the file stream to read the data from.
  *
- * conn should be a database connection that you just called COPY FROM on
- * (and which gave you PGRES_COPY_IN back);
- * copystream is the file stream you want the input to come from
+ * result is true if successful, false if not.
  */
 
+/* read chunk size for COPY IN - size is not critical */
+#define COPYBUFSIZ 8192
+
 bool
 handleCopyIn(PGconn *conn, FILE *copystream)
 {
+       bool            OK = true;
        const char *prompt;
        bool            copydone = false;
        bool            firstload;
        bool            linedone;
-       bool            saw_cr = false;
-       char            copybuf[COPYBUFSIZ];
-       char       *s;
-       int                     bufleft;
-       int                     c = 0;
-       int                     ret;
-       unsigned int linecount = 0;
+       char            buf[COPYBUFSIZ];
+       PGresult *res;
 
        /* Prompt if interactive input */
        if (isatty(fileno(copystream)))
@@ -684,64 +693,65 @@ handleCopyIn(PGconn *conn, FILE *copystream)
                        fputs(prompt, stdout);
                        fflush(stdout);
                }
+               
                firstload = true;
                linedone = false;
 
                while (!linedone)
                {                                               /* for each bufferload in line ... */
-                       /* Fetch string until \n, EOF, or buffer full */
-                       s = copybuf;
-                       for (bufleft = COPYBUFSIZ - 1; bufleft > 0; bufleft--)
-                       {
-                               c = getc(copystream);
-                               if (c == EOF)
-                               {
-                                       linedone = true;
-                                       break;
-                               }
-                               *s++ = c;
-                               if (c == '\n')
-                               {
-                                       linedone = true;
-                                       break;
-                               }
-                               if (c == '\r')
-                                       saw_cr = true;
-                       }
-                       *s = '\0';
-                       /* EOF with empty line-so-far? */
-                       if (c == EOF && s == copybuf && firstload)
+                       int             linelen;
+
+                       if (!fgets(buf, COPYBUFSIZ, copystream))
                        {
-                               /*
-                                * We are guessing a little bit as to the right line-ending
-                                * here...
-                                */
-                               if (saw_cr)
-                                       PQputline(conn, "\\.\r\n");
-                               else
-                                       PQputline(conn, "\\.\n");
+                               if (ferror(copystream))
+                                       OK = false;
                                copydone = true;
-                               if (pset.cur_cmd_interactive)
-                                       puts("\\.");
                                break;
                        }
-                       /* No, so pass the data to the backend */
-                       PQputline(conn, copybuf);
-                       /* Check for line consisting only of \. */
+
+                       linelen = strlen(buf);
+
+                       /* current line is done? */
+                       if (linelen > 0 && buf[linelen-1] == '\n')
+                               linedone = true;
+
+                       /* check for EOF marker, but not on a partial line */
                        if (firstload)
                        {
-                               if (strcmp(copybuf, "\\.\n") == 0 ||
-                                       strcmp(copybuf, "\\.\r\n") == 0)
+                               if (strcmp(buf, "\\.\n") == 0 ||
+                                       strcmp(buf, "\\.\r\n") == 0)
                                {
                                        copydone = true;
                                        break;
                                }
+                               
                                firstload = false;
                        }
+                       
+                       if (PQputCopyData(conn, buf, linelen) <= 0)
+                       {
+                               OK = false;
+                               copydone = true;
+                               break;
+                       }
                }
-               linecount++;
+               
+               pset.lineno++;
+       }
+
+       /* Terminate data transfer */
+       if (PQputCopyEnd(conn,
+                                        OK ? NULL : _("aborted due to read failure")) <= 0)
+               OK = false;
+
+       /* Check command status and return to normal libpq state */
+       res = PQgetResult(conn);
+       if (PQresultStatus(res) != PGRES_COMMAND_OK)
+       {
+               psql_error("%s", PQerrorMessage(conn));
+               OK = false;
        }
-       ret = !PQendcopy(conn);
-       pset.lineno += linecount;
-       return ret;
+       PQclear(res);
+
+       return OK;
 }