pg_backup_archiver.c

PostgreSQL7.4.6 for Linux

	ArchiveHandle *AH = (ArchiveHandle *) AHX;

	if (!AH->StartBlobPtr)
		die_horribly(AH, modulename, "large-object output not supported in chosen format\n");

	(*AH->StartBlobPtr) (AH, AH->currToc, oid);

	return 1;
}

/* Called by a dumper to signal end of a BLOB */
int
EndBlob(Archive *AHX, Oid oid)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;

	if (AH->EndBlobPtr)
		(*AH->EndBlobPtr) (AH, AH->currToc, oid);

	return 1;
}

/**********
 * BLOB Restoration
 **********/

/*
 * Called by a format handler before any blobs are restored
 */
void
StartRestoreBlobs(ArchiveHandle *AH)
{
	AH->blobCount = 0;
}

/*
 * Called by a format handler after all blobs are restored
 */
void
EndRestoreBlobs(ArchiveHandle *AH)
{
	if (AH->txActive)
	{
		ahlog(AH, 2, "committing large-object transactions\n");
		CommitTransaction(AH);
	}

	if (AH->blobTxActive)
		CommitTransactionXref(AH);

	ahlog(AH, 1, "restored %d large objects\n", AH->blobCount);
}


/*
 * Called by a format handler to initiate restoration of a blob
 */
void
StartRestoreBlob(ArchiveHandle *AH, Oid oid)
{
	Oid			loOid;

	AH->blobCount++;

	if (!AH->createdBlobXref)
	{
		if (!AH->connection)
			die_horribly(AH, modulename, "cannot restore large objects without a database connection\n");

		CreateBlobXrefTable(AH);
		AH->createdBlobXref = 1;
	}

	/* Initialize the LO Buffer */
	AH->lo_buf_used = 0;

	/*
	 * Start long-running TXs if necessary
	 */
	if (!AH->txActive)
	{
		ahlog(AH, 2, "starting large-object transactions\n");
		StartTransaction(AH);
	}
	if (!AH->blobTxActive)
		StartTransactionXref(AH);

	loOid = lo_creat(AH->connection, INV_READ | INV_WRITE);
	if (loOid == 0)
		die_horribly(AH, modulename, "could not create large object\n");

	ahlog(AH, 2, "restoring large object with OID %u as %u\n", oid, loOid);

	InsertBlobXref(AH, oid, loOid);

	AH->loFd = lo_open(AH->connection, loOid, INV_WRITE);
	if (AH->loFd == -1)
		die_horribly(AH, modulename, "could not open large object\n");

	AH->writingBlob = 1;
}

void
EndRestoreBlob(ArchiveHandle *AH, Oid oid)
{
	if (AH->lo_buf_used > 0)
	{
		/* Write remaining bytes from the LO buffer */
		size_t		res;

		res = lo_write(AH->connection, AH->loFd, (void *) AH->lo_buf, AH->lo_buf_used);

		ahlog(AH, 5, "wrote remaining %lu bytes of large-object data (result = %lu)\n",
			  (unsigned long) AH->lo_buf_used, (unsigned long) res);
		if (res != AH->lo_buf_used)
			die_horribly(AH, modulename, "could not write to large object (result: %lu, expected: %lu)\n",
				   (unsigned long) res, (unsigned long) AH->lo_buf_used);
		AH->lo_buf_used = 0;
	}

	lo_close(AH->connection, AH->loFd);
	AH->writingBlob = 0;

	/*
	 * Commit every BLOB_BATCH_SIZE blobs...
	 */
	if (((AH->blobCount / BLOB_BATCH_SIZE) * BLOB_BATCH_SIZE) == AH->blobCount)
	{
		ahlog(AH, 2, "committing large-object transactions\n");
		CommitTransaction(AH);
		CommitTransactionXref(AH);
	}
}

/***********
 * Sorting and Reordering
 ***********/

/*
 * Move TOC entries of the specified type to the START of the TOC.
 *
 * This is public, but if you use it anywhere but SortTocByObjectType,
 * you are risking breaking things.
 */
void
MoveToStart(Archive *AHX, const char *oType)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;
	TocEntry   *te = AH->toc->next;
	TocEntry   *newTe;

	while (te != AH->toc)
	{
		te->_moved = 0;
		te = te->next;
	}

	te = AH->toc->prev;
	while (te != AH->toc && !te->_moved)
	{
		newTe = te->prev;
		if (strcmp(te->desc, oType) == 0)
			_moveAfter(AH, AH->toc, te);
		te = newTe;
	}
}


/*
 * Move TOC entries of the specified type to the end of the TOC.
 *
 * This is public, but if you use it anywhere but SortTocByObjectType,
 * you are risking breaking things.
 */
void
MoveToEnd(Archive *AHX, const char *oType)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;
	TocEntry   *te = AH->toc->next;
	TocEntry   *newTe;

	while (te != AH->toc)
	{
		te->_moved = 0;
		te = te->next;
	}

	te = AH->toc->next;
	while (te != AH->toc && !te->_moved)
	{
		newTe = te->next;
		if (strcmp(te->desc, oType) == 0)
			_moveBefore(AH, AH->toc, te);
		te = newTe;
	}
}

/*
 * Sort TOC by object type (items of same type keep same relative order)
 *
 * This is factored out to ensure that pg_dump and pg_restore stay in sync
 * about the standard ordering.
 */
void
SortTocByObjectType(Archive *AH)
{
	/*
	 * Procedural languages have to be declared just after database and
	 * schema creation, before they are used.
	 */
	MoveToStart(AH, "ACL LANGUAGE");
	MoveToStart(AH, "PROCEDURAL LANGUAGE");
	MoveToStart(AH, "FUNC PROCEDURAL LANGUAGE");
	MoveToStart(AH, "SCHEMA");
	MoveToStart(AH, "<Init>");
	/* Database entries *must* be at front (see also pg_restore.c) */
	MoveToStart(AH, "DATABASE");

	MoveToEnd(AH, "TABLE DATA");
	MoveToEnd(AH, "BLOBS");
	MoveToEnd(AH, "INDEX");
	MoveToEnd(AH, "CONSTRAINT");
	MoveToEnd(AH, "FK CONSTRAINT");
	MoveToEnd(AH, "TRIGGER");
	MoveToEnd(AH, "RULE");
	MoveToEnd(AH, "SEQUENCE SET");

	/*
	 * Moving all comments to end is annoying, but must do it for comments
	 * on stuff we just moved, and we don't seem to have quite enough
	 * dependency structure to get it really right...
	 */
	MoveToEnd(AH, "COMMENT");
}

/*
 * Sort TOC by OID
 */
/* Public */
void
SortTocByOID(Archive *AHX)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;

	_SortToc(AH, _tocSortCompareByOIDNum);
}

/*
 * Sort TOC by ID
 */
/* Public */
void
SortTocByID(Archive *AHX)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;

	_SortToc(AH, _tocSortCompareByIDNum);
}

void
SortTocFromFile(Archive *AHX, RestoreOptions *ropt)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;
	FILE	   *fh;
	char		buf[1024];
	char	   *cmnt;
	char	   *endptr;
	int			id;
	TocEntry   *te;
	TocEntry   *tePrev;
	int			i;

	/* Allocate space for the 'wanted' array, and init it */
	ropt->idWanted = (int *) malloc(sizeof(int) * AH->tocCount);
	for (i = 0; i < AH->tocCount; i++)
		ropt->idWanted[i] = 0;

	ropt->limitToList = 1;

	/* Mark all entries as 'not moved' */
	te = AH->toc->next;
	while (te != AH->toc)
	{
		te->_moved = 0;
		te = te->next;
	}

	/* Set prev entry as head of list */
	tePrev = AH->toc;

	/* Setup the file */
	fh = fopen(ropt->tocFile, PG_BINARY_R);
	if (!fh)
		die_horribly(AH, modulename, "could not open TOC file\n");

	while (fgets(buf, 1024, fh) != NULL)
	{
		/* Find a comment */
		cmnt = strchr(buf, ';');
		if (cmnt == buf)
			continue;

		/* End string at comment */
		if (cmnt != NULL)
			cmnt[0] = '\0';

		/* Skip if all spaces */
		if (strspn(buf, " \t") == strlen(buf))
			continue;

		/* Get an ID */
		id = strtol(buf, &endptr, 10);
		if (endptr == buf)
		{
			write_msg(modulename, "WARNING: line ignored: %s\n", buf);
			continue;
		}

		/* Find TOC entry */
		te = _getTocEntry(AH, id);
		if (!te)
			die_horribly(AH, modulename, "could not find entry for ID %d\n", id);

		ropt->idWanted[id - 1] = 1;

		_moveAfter(AH, tePrev, te);
		tePrev = te;
	}

	if (fclose(fh) != 0)
		die_horribly(AH, modulename, "could not close TOC file: %s\n", strerror(errno));
}

/**********************
 * 'Convenience functions that look like standard IO functions
 * for writing data when in dump mode.
 **********************/

/* Public */
int
archputs(const char *s, Archive *AH)
{
	return WriteData(AH, s, strlen(s));
}

/* Public */
int
archputc(const char c, Archive *AH)
{
	return WriteData(AH, &c, 1);
}

/* Public */
int
archprintf(Archive *AH, const char *fmt,...)
{
	char	   *p = NULL;
	va_list		ap;
	int			bSize = strlen(fmt) + 256;
	int			cnt = -1;

	/*
	 * This is paranoid: deal with the possibility that vsnprintf is
	 * willing to ignore trailing null
	 */

	/*
	 * or returns > 0 even if string does not fit. It may be the case that
	 * it returns cnt = bufsize
	 */
	while (cnt < 0 || cnt >= (bSize - 1))
	{
		if (p != NULL)
			free(p);
		bSize *= 2;
		p = (char *) malloc(bSize);
		if (p == NULL)
			exit_horribly(AH, modulename, "out of memory\n");
		va_start(ap, fmt);
		cnt = vsnprintf(p, bSize, fmt, ap);
		va_end(ap);
	}
	WriteData(AH, p, cnt);
	free(p);
	return cnt;
}


/*******************************
 * Stuff below here should be 'private' to the archiver routines
 *******************************/

OutputContext
SetOutput(ArchiveHandle *AH, char *filename, int compression)
{
	OutputContext sav;
	int			fn;

	/* Replace the AH output file handle */
	sav.OF = AH->OF;
	sav.gzOut = AH->gzOut;

	if (filename)
		fn = -1;
	else if (AH->FH)
		fn = fileno(AH->FH);
	else if (AH->fSpec)
	{
		fn = -1;
		filename = AH->fSpec;
	}
	else
		fn = fileno(stdout);

	/* If compression explicitly requested, use gzopen */
#ifdef HAVE_LIBZ
	if (compression != 0)
	{
		char		fmode[10];

		/* Don't use PG_BINARY_x since this is zlib */
		sprintf(fmode, "wb%d", compression);
		if (fn >= 0)
			AH->OF = gzdopen(dup(fn), fmode);
		else
			AH->OF = gzopen(filename, fmode);
		AH->gzOut = 1;
	}
	else
#endif
	{							/* Use fopen */
		if (fn >= 0)
			AH->OF = fdopen(dup(fn), PG_BINARY_W);
		else
			AH->OF = fopen(filename, PG_BINARY_W);
		AH->gzOut = 0;
	}

	if (!AH->OF)
		die_horribly(AH, modulename, "could not open output file: %s\n", strerror(errno));

	return sav;
}

void
ResetOutput(ArchiveHandle *AH, OutputContext sav)
{
	int			res;

	if (AH->gzOut)
		res = GZCLOSE(AH->OF);
	else
		res = fclose(AH->OF);

	if (res != 0)
		die_horribly(AH, modulename, "could not close output file: %s\n",
					 strerror(errno));

	AH->gzOut = sav.gzOut;
	AH->OF = sav.OF;
}



/*
 *	Print formatted text to the output file (usually stdout).
 */
int
ahprintf(ArchiveHandle *AH, const char *fmt,...)
{
	char	   *p = NULL;
	va_list		ap;
	int			bSize = strlen(fmt) + 256;		/* Should be enough */
	int			cnt = -1;

	/*
	 * This is paranoid: deal with the possibility that vsnprintf is
	 * willing to ignore trailing null
	 */

	/*
	 * or returns > 0 even if string does not fit. It may be the case that
	 * it returns cnt = bufsize
	 */
	while (cnt < 0 || cnt >= (bSize - 1))
	{
		if (p != NULL)
			free(p);
		bSize *= 2;
		p = (char *) malloc(bSize);
		if (p == NULL)
			die_horribly(AH, modulename, "out of memory\n");
		va_start(ap, fmt);
		cnt = vsnprintf(p, bSize, fmt, ap);
		va_end(ap);
	}
	ahwrite(p, 1, cnt, AH);
	free(p);
	return cnt;
}

void
ahlog(ArchiveHandle *AH, int level, const char *fmt,...)
{
	va_list		ap;

	if (AH->debugLevel < level && (!AH->public.verbose || level > 1))
		return;

	va_start(ap, fmt);
	_write_msg(NULL, fmt, ap);
	va_end(ap);
}

/*
 * Single place for logic which says 'We are restoring to a direct DB connection'.
 */
int
RestoringToDB(ArchiveHandle *AH)
{
	return (AH->ropt && AH->ropt->useDB && AH->connection);
}

/*
 *	Write buffer to the output file (usually stdout). This is user for
 *	outputting 'restore' scripts etc. It is even possible for an archive
 *	format to create a custom output routine to 'fake' a restore if it
 *	wants to generate a script (see TAR output).
 */
int
ahwrite(const void *ptr, size_t size, size_t nmemb, ArchiveHandle *AH)
{
	size_t		res;

	if (AH->writingBlob)
	{
		if (AH->lo_buf_used + size * nmemb > AH->lo_buf_size)
		{
			/* Split LO buffer */
			size_t		remaining = AH->lo_buf_size - AH->lo_buf_used;
			size_t		slack = nmemb * size - remaining;

			memcpy((char *) AH->lo_buf + AH->lo_buf_used, ptr, remaining);
			res = lo_write(AH->connection, AH->loFd, AH->lo_buf, AH->lo_buf_size);
			ahlog(AH, 5, "wrote %lu bytes of large object data (result = %lu)\n",
				  (unsigned long) AH->lo_buf_size, (unsigned long) res);
			if (res != AH->lo_buf_size)
				die_horribly(AH, modulename,
							 "could not write to large object (result: %lu, expected: %lu)\n",
				   (unsigned long) res, (unsigned long) AH->lo_buf_size);
			memcpy(AH->lo_buf, (char *) ptr + remaining, slack);
			AH->lo_buf_used = slack;
		}
		else
		{
			/* LO Buffer is still large enough, buffer it */
			memcpy((char *) AH->lo_buf + AH->lo_buf_used, ptr, size * nmemb);
			AH->lo_buf_used += size * nmemb;
		}

		return size * nmemb;
	}
	else if (AH->gzOut)
	{
		res = GZWRITE((void *) ptr, size, nmemb, AH->OF);
		if (res != (nmemb * size))
			die_horribly(AH, modulename, "could not write to compressed archive\n");
		return res;
	}
	else if (AH->CustomOutPtr)
	{
		res = AH->CustomOutPtr (AH, ptr, size * nmemb);

		if (res != (nmemb * size))
			die_horribly(AH, modulename, "could not write to custom output routine\n");
		return res;
	}
	else
	{
		/*
		 * If we're doing a restore, and it's direct to DB, and we're
		 * connected then send it to the DB.
		 */
		if (RestoringToDB(AH))
			return ExecuteSqlCommandBuf(AH, (void *) ptr, size * nmemb);		/* Always 1, currently */
		else
		{
			res = fwrite((void *) ptr, size, nmemb, AH->OF);
			if (res != nmemb)
				die_horribly(AH, modulename, "could not write to output file (%lu != %lu)\n",
							 (unsigned long) res, (unsigned long) nmemb);
			return res;
		}
	}
}

/* Common exit code */
static void
_write_msg(const char *modulename, const char *fmt, va_list ap)
{
	if (modulename)
		fprintf(stderr, "%s: [%s] ", progname, gettext(modulename));
	else
		fprintf(stderr, "%s: ", progname);
	vfprintf(stderr, gettext(fmt), ap);
}

void
write_msg(const char *modulename, const char *fmt,...)
{
	va_list		ap;

	va_start(ap, fmt);
	_write_msg(modulename, fmt, ap);
	va_end(ap);
}


static void
_die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt, va_list ap)
{
	_write_msg(modulename, fmt, ap);

	if (AH)
	{
		if (AH->public.verbose)
			write_msg(NULL, "*** aborted because of error\n");
		if (AH->connection)
			PQfinish(AH->connection);
		if (AH->blobConnection)
			PQfinish(AH->blobConnection);
	}

	exit(1);
}

/* External use */
void
exit_horribly(Archive *AH, const char *modulename, const char *fmt,...)
{
	va_list		ap;

	va_start(ap, fmt);
	_die_horribly((ArchiveHandle *) AH, modulename, fmt, ap);
}

/* Archiver use (just different arg declaration) */
void
die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt,...)
{
	va_list		ap;