Improve pg_dump and psql to use libpq's newer COPY support routines,

instead of the old deprecated ones.
Volkan Yazici, with some editorializing by moi.

src/bin/pg_dump/pg_backup_db.c

@@ -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;
 	}
 

src/bin/pg_dump/pg_dump.c

@@ -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;
 }

src/bin/psql/common.c

@@ -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".
  */

src/bin/psql/copy.c

@@ -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;
 }