1 files changed, 46 insertions, 329 deletions

diff --git a/src/bin/pg_dump/pg_backup_db.c b/src/bin/pg_dump/pg_backup_db.c
index 37d1b742e41..1ce0d3164d2 100644
--- a/src/bin/pg_dump/pg_backup_db.c
+++ b/src/bin/pg_dump/pg_backup_db.c
@@ -14,26 +14,19 @@
 #include "dumputils.h"
 
 #include <unistd.h>
-
 #include <ctype.h>
-
 #ifdef HAVE_TERMIOS_H
 #include <termios.h>
 #endif
 
 
+#define DB_MAX_ERR_STMT 128
+
 static const char *modulename = gettext_noop("archiver (db)");
 
 static void _check_database_version(ArchiveHandle *AH);
 static PGconn *_connectDB(ArchiveHandle *AH, const char *newdbname, const char *newUser);
 static void notice_processor(void *arg, const char *message);
-static char *_sendSQLLine(ArchiveHandle *AH, char *qry, char *eos);
-static char *_sendCopyLine(ArchiveHandle *AH, char *qry, char *eos);
-
-static bool _isIdentChar(unsigned char c);
-static bool _isDQChar(unsigned char c, bool atStart);
-
-#define DB_MAX_ERR_STMT 128
 
 static int
 _parse_version(ArchiveHandle *AH, const char *versionString)
@@ -330,8 +323,10 @@ notice_processor(void *arg, const char *message)
 }
 
 
-/* Public interface */
-/* Convenience function to send a query. Monitors result to handle COPY statements */
+/*
+ * Convenience function to send a query.
+ * Monitors result to detect COPY statements
+ */
 static void
 ExecuteSqlCommand(ArchiveHandle *AH, const char *qry, const char *desc)
 {
@@ -348,6 +343,7 @@ ExecuteSqlCommand(ArchiveHandle *AH, const char *qry, const char *desc)
 	{
 		case PGRES_COMMAND_OK:
 		case PGRES_TUPLES_OK:
+		case PGRES_EMPTY_QUERY:
 			/* A-OK */
 			break;
 		case PGRES_COPY_IN:
@@ -372,78 +368,56 @@ ExecuteSqlCommand(ArchiveHandle *AH, const char *qry, const char *desc)
 	PQclear(res);
 }
 
+
 /*
- * Used by ExecuteSqlCommandBuf to send one buffered line when running a COPY command.
+ * Implement ahwrite() for direct-to-DB restore
  */
-static char *
-_sendCopyLine(ArchiveHandle *AH, char *qry, char *eos)
+int
+ExecuteSqlCommandBuf(ArchiveHandle *AH, const char *buf, size_t bufLen)
 {
-	size_t		loc;			/* Location of next newline */
-	int			pos = 0;		/* Current position */
-	int			sPos = 0;		/* Last pos of a slash char */
-	int			isEnd = 0;
-
-	/* loop to find unquoted newline ending the line of COPY data */
-	for (;;)
+	if (AH->writingCopyData)
 	{
-		loc = strcspn(&qry[pos], "\n") + pos;
-
-		/* If no match, then wait */
-		if (loc >= (eos - qry)) /* None found */
-		{
-			appendBinaryPQExpBuffer(AH->pgCopyBuf, qry, (eos - qry));
-			return eos;
-		}
-
 		/*
-		 * fprintf(stderr, "Found cr at %d, prev char was %c, next was %c\n",
-		 * loc, qry[loc-1], qry[loc+1]);
+		 * 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.
 		 */
-
-		/* Count the number of preceding slashes */
-		sPos = loc;
-		while (sPos > 0 && qry[sPos - 1] == '\\')
-			sPos--;
-
-		sPos = loc - sPos;
-
+		if (AH->pgCopyIn &&
+			PQputCopyData(AH->connection, buf, bufLen) <= 0)
+			die_horribly(AH, modulename, "error returned by PQputCopyData: %s",
+						 PQerrorMessage(AH->connection));
+	}
+	else
+	{
 		/*
-		 * If an odd number of preceding slashes, then \n was escaped so set
-		 * the next search pos, and loop (if any left).
+		 * In most cases the data passed to us will be a null-terminated
+		 * string, but if it's not, we have to add a trailing null.
 		 */
-		if ((sPos & 1) == 1)
+		if (buf[bufLen] == '\0')
+			ExecuteSqlCommand(AH, buf, "could not execute query");
+		else
 		{
-			/* fprintf(stderr, "cr was escaped\n"); */
-			pos = loc + 1;
-			if (pos >= (eos - qry))
-			{
-				appendBinaryPQExpBuffer(AH->pgCopyBuf, qry, (eos - qry));
-				return eos;
-			}
+			char   *str = (char *) malloc(bufLen + 1);
+
+			if (!str)
+				die_horribly(AH, modulename, "out of memory\n");
+			memcpy(str, buf, bufLen);
+			str[bufLen] = '\0';
+			ExecuteSqlCommand(AH, str, "could not execute query");
+			free(str);
 		}
-		else
-			break;
 	}
 
-	/* We found an unquoted newline */
-	qry[loc] = '\0';
-	appendPQExpBuffer(AH->pgCopyBuf, "%s\n", qry);
-	isEnd = (strcmp(AH->pgCopyBuf->data, "\\.\n") == 0);
-
-	/*
-	 * 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 (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);
+	return 1;
+}
 
-	if (isEnd && AH->pgCopyIn)
+/*
+ * Terminate a COPY operation during direct-to-DB restore
+ */
+void
+EndDBCopyMode(ArchiveHandle *AH, TocEntry *te)
+{
+	if (AH->pgCopyIn)
 	{
 		PGresult   *res;
 
@@ -454,240 +428,12 @@ _sendCopyLine(ArchiveHandle *AH, char *qry, char *eos)
 		/* 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));
+			warn_or_die_horribly(AH, modulename, "COPY failed for table \"%s\": %s",
+								 te->tag, PQerrorMessage(AH->connection));
 		PQclear(res);
 
 		AH->pgCopyIn = false;
 	}
-
-	return qry + loc + 1;
-}
-
-/*
- * Used by ExecuteSqlCommandBuf to send one buffered line of SQL
- * (not data for the copy command).
- */
-static char *
-_sendSQLLine(ArchiveHandle *AH, char *qry, char *eos)
-{
-	/*
-	 * The following is a mini state machine to assess the end of an SQL
-	 * statement. It really only needs to parse good SQL, or at least that's
-	 * the theory... End-of-statement is assumed to be an unquoted,
-	 * un-commented semi-colon that's not within any parentheses.
-	 *
-	 * Note: the input can be split into bufferloads at arbitrary boundaries.
-	 * Therefore all state must be kept in AH->sqlparse, not in local
-	 * variables of this routine.  We assume that AH->sqlparse was filled with
-	 * zeroes when created.
-	 */
-	for (; qry < eos; qry++)
-	{
-		switch (AH->sqlparse.state)
-		{
-			case SQL_SCAN:		/* Default state == 0, set in _allocAH */
-				if (*qry == ';' && AH->sqlparse.braceDepth == 0)
-				{
-					/*
-					 * We've found the end of a statement. Send it and reset
-					 * the buffer.
-					 */
-					appendPQExpBufferChar(AH->sqlBuf, ';');		/* inessential */
-					ExecuteSqlCommand(AH, AH->sqlBuf->data,
-									  "could not execute query");
-					resetPQExpBuffer(AH->sqlBuf);
-					AH->sqlparse.lastChar = '\0';
-
-					/*
-					 * Remove any following newlines - so that embedded COPY
-					 * commands don't get a starting newline.
-					 */
-					qry++;
-					while (qry < eos && *qry == '\n')
-						qry++;
-
-					/* We've finished one line, so exit */
-					return qry;
-				}
-				else if (*qry == '\'')
-				{
-					if (AH->sqlparse.lastChar == 'E')
-						AH->sqlparse.state = SQL_IN_E_QUOTE;
-					else
-						AH->sqlparse.state = SQL_IN_SINGLE_QUOTE;
-					AH->sqlparse.backSlash = false;
-				}
-				else if (*qry == '"')
-				{
-					AH->sqlparse.state = SQL_IN_DOUBLE_QUOTE;
-				}
-
-				/*
-				 * Look for dollar-quotes. We make the assumption that
-				 * $-quotes will not have an ident character just before them
-				 * in pg_dump output.  XXX is this good enough?
-				 */
-				else if (*qry == '$' && !_isIdentChar(AH->sqlparse.lastChar))
-				{
-					AH->sqlparse.state = SQL_IN_DOLLAR_TAG;
-					/* initialize separate buffer with possible tag */
-					if (AH->sqlparse.tagBuf == NULL)
-						AH->sqlparse.tagBuf = createPQExpBuffer();
-					else
-						resetPQExpBuffer(AH->sqlparse.tagBuf);
-					appendPQExpBufferChar(AH->sqlparse.tagBuf, *qry);
-				}
-				else if (*qry == '-' && AH->sqlparse.lastChar == '-')
-					AH->sqlparse.state = SQL_IN_SQL_COMMENT;
-				else if (*qry == '*' && AH->sqlparse.lastChar == '/')
-					AH->sqlparse.state = SQL_IN_EXT_COMMENT;
-				else if (*qry == '(')
-					AH->sqlparse.braceDepth++;
-				else if (*qry == ')')
-					AH->sqlparse.braceDepth--;
-				break;
-
-			case SQL_IN_SQL_COMMENT:
-				if (*qry == '\n')
-					AH->sqlparse.state = SQL_SCAN;
-				break;
-
-			case SQL_IN_EXT_COMMENT:
-
-				/*
-				 * This isn't fully correct, because we don't account for
-				 * nested slash-stars, but pg_dump never emits such.
-				 */
-				if (AH->sqlparse.lastChar == '*' && *qry == '/')
-					AH->sqlparse.state = SQL_SCAN;
-				break;
-
-			case SQL_IN_SINGLE_QUOTE:
-				/* We needn't handle '' specially */
-				if (*qry == '\'' && !AH->sqlparse.backSlash)
-					AH->sqlparse.state = SQL_SCAN;
-				else if (*qry == '\\')
-					AH->sqlparse.backSlash = !AH->sqlparse.backSlash;
-				else
-					AH->sqlparse.backSlash = false;
-				break;
-
-			case SQL_IN_E_QUOTE:
-
-				/*
-				 * Eventually we will need to handle '' specially, because
-				 * after E'...''... we should still be in E_QUOTE state.
-				 *
-				 * XXX problem: how do we tell whether the dump was made by a
-				 * version that thinks backslashes aren't special in non-E
-				 * literals??
-				 */
-				if (*qry == '\'' && !AH->sqlparse.backSlash)
-					AH->sqlparse.state = SQL_SCAN;
-				else if (*qry == '\\')
-					AH->sqlparse.backSlash = !AH->sqlparse.backSlash;
-				else
-					AH->sqlparse.backSlash = false;
-				break;
-
-			case SQL_IN_DOUBLE_QUOTE:
-				/* We needn't handle "" specially */
-				if (*qry == '"')
-					AH->sqlparse.state = SQL_SCAN;
-				break;
-
-			case SQL_IN_DOLLAR_TAG:
-				if (*qry == '$')
-				{
-					/* Do not add the closing $ to tagBuf */
-					AH->sqlparse.state = SQL_IN_DOLLAR_QUOTE;
-					AH->sqlparse.minTagEndPos = AH->sqlBuf->len + AH->sqlparse.tagBuf->len + 1;
-				}
-				else if (_isDQChar(*qry, (AH->sqlparse.tagBuf->len == 1)))
-				{
-					/* Valid, so add to tag */
-					appendPQExpBufferChar(AH->sqlparse.tagBuf, *qry);
-				}
-				else
-				{
-					/*
-					 * Ooops, we're not really in a dollar-tag.  Valid tag
-					 * chars do not include the various chars we look for in
-					 * this state machine, so it's safe to just jump from this
-					 * state back to SCAN.	We have to back up the qry pointer
-					 * so that the current character gets rescanned in SCAN
-					 * state; and then "continue" so that the bottom-of-loop
-					 * actions aren't done yet.
-					 */
-					AH->sqlparse.state = SQL_SCAN;
-					qry--;
-					continue;
-				}
-				break;
-
-			case SQL_IN_DOLLAR_QUOTE:
-
-				/*
-				 * If we are at a $, see whether what precedes it matches
-				 * tagBuf.	(Remember that the trailing $ of the tag was not
-				 * added to tagBuf.)  However, don't compare until we have
-				 * enough data to be a possible match --- this is needed to
-				 * avoid false match on '$a$a$...'
-				 */
-				if (*qry == '$' &&
-					AH->sqlBuf->len >= AH->sqlparse.minTagEndPos &&
-					strcmp(AH->sqlparse.tagBuf->data,
-						   AH->sqlBuf->data + AH->sqlBuf->len - AH->sqlparse.tagBuf->len) == 0)
-					AH->sqlparse.state = SQL_SCAN;
-				break;
-		}
-
-		appendPQExpBufferChar(AH->sqlBuf, *qry);
-		AH->sqlparse.lastChar = *qry;
-	}
-
-	/*
-	 * If we get here, we've processed entire bufferload with no complete SQL
-	 * stmt
-	 */
-	return eos;
-}
-
-
-/* Convenience function to send one or more queries. Monitors result to handle COPY statements */
-int
-ExecuteSqlCommandBuf(ArchiveHandle *AH, void *qryv, size_t bufLen)
-{
-	char	   *qry = (char *) qryv;
-	char	   *eos = qry + bufLen;
-
-	/*
-	 * fprintf(stderr, "\n\n*****\n Buffer:\n\n%s\n*******************\n\n",
-	 * qry);
-	 */
-
-	/* Could switch between command and COPY IN mode at each line */
-	while (qry < eos)
-	{
-		/*
-		 * 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);
-	}
-
-	return 1;
 }
 
 void
@@ -728,32 +474,3 @@ DropBlobIfExists(ArchiveHandle *AH, Oid oid)
 				 oid, oid);
 	}
 }
-
-static bool
-_isIdentChar(unsigned char c)
-{
-	if ((c >= 'a' && c <= 'z')
-		|| (c >= 'A' && c <= 'Z')
-		|| (c >= '0' && c <= '9')
-		|| (c == '_')
-		|| (c == '$')
-		|| (c >= (unsigned char) '\200')		/* no need to check <= \377 */
-		)
-		return true;
-	else
-		return false;
-}
-
-static bool
-_isDQChar(unsigned char c, bool atStart)
-{
-	if ((c >= 'a' && c <= 'z')
-		|| (c >= 'A' && c <= 'Z')
-		|| (c == '_')
-		|| (!atStart && c >= '0' && c <= '9')
-		|| (c >= (unsigned char) '\200')		/* no need to check <= \377 */
-		)
-		return true;
-	else
-		return false;
-}