From 853c54d86dcaeca9ad6496fa84323b0982741227 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <bossartn@amazon.com>
Date: Wed, 20 Oct 2021 21:51:08 +0000
Subject: [PATCH v4 2/2] Add an example of a basic archiving module.

---
 contrib/Makefile                      |   1 +
 contrib/basic_archive/Makefile        |  15 +++
 contrib/basic_archive/basic_archive.c | 218 ++++++++++++++++++++++++++++++++++
 3 files changed, 234 insertions(+)
 create mode 100644 contrib/basic_archive/Makefile
 create mode 100644 contrib/basic_archive/basic_archive.c

diff --git a/contrib/Makefile b/contrib/Makefile
index f27e458482..aff834ebaa 100644
--- a/contrib/Makefile
+++ b/contrib/Makefile
@@ -9,6 +9,7 @@ SUBDIRS = \
 		amcheck		\
 		auth_delay	\
 		auto_explain	\
+		basic_archive	\
 		bloom		\
 		btree_gin	\
 		btree_gist	\
diff --git a/contrib/basic_archive/Makefile b/contrib/basic_archive/Makefile
new file mode 100644
index 0000000000..ea6b460889
--- /dev/null
+++ b/contrib/basic_archive/Makefile
@@ -0,0 +1,15 @@
+# contrib/basic_archive/Makefile
+
+MODULES = basic_archive
+PGFILEDESC = "basic_archive - basic archive module"
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/basic_archive
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/contrib/basic_archive/basic_archive.c b/contrib/basic_archive/basic_archive.c
new file mode 100644
index 0000000000..ed78c96e16
--- /dev/null
+++ b/contrib/basic_archive/basic_archive.c
@@ -0,0 +1,218 @@
+/*-------------------------------------------------------------------------
+ *
+ * basic_archive.c
+ *
+ * This file demonstrates a basic archive library implementation that is
+ * roughly equivalent to the following shell command:
+ *
+ * 		test ! -f /path/to/dest && cp /path/to/src /path/to/dest
+ *
+ * One notable difference between this module and the shell command above
+ * is that this module first copies the file to a temporary destination,
+ * syncs it to disk, and then durably moves it to the final destination.
+ *
+ * Copyright (c) 2021, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *	  contrib/basic_archive/basic_archive.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include "fmgr.h"
+#include "miscadmin.h"
+#include "postmaster/pgarch.h"
+#include "storage/fd.h"
+#include "utils/guc.h"
+
+PG_MODULE_MAGIC;
+
+void	_PG_init(void);
+bool	basic_archive(const char *file, const char *path);
+
+static char *archive_directory = NULL;
+
+static bool check_archive_directory(char **newval, void **extra, GucSource source);
+static void copy_file(const char *src, const char *dst);
+
+void
+_PG_init(void)
+{
+	if (!process_shared_preload_libraries_in_progress)
+		ereport(ERROR,
+				(errmsg("\"basic_archive\" can only be loaded via "
+						"\"shared_preload_libraries\"")));
+
+	if (PG_archive != shell_archive)
+		ereport(ERROR,
+				(errmsg("custom archive function already loaded by another module")));
+
+	DefineCustomStringVariable("basic_archive.archive_directory",
+							   gettext_noop("Archive file destination directory."),
+							   NULL,
+							   &archive_directory,
+							   "",
+							   PGC_POSTMASTER,
+							   GUC_NOT_IN_SAMPLE,
+							   check_archive_directory, NULL, NULL);
+
+	EmitWarningsOnPlaceholders("basic_archive");
+
+	PG_archive = basic_archive;
+}
+
+static bool
+check_archive_directory(char **newval, void **extra, GucSource source)
+{
+	struct stat st;
+
+	if (*newval == NULL || *newval[0] == '\0')
+		return true;
+
+	if (stat(*newval, &st) != 0 || !S_ISDIR(st.st_mode))
+	{
+		GUC_check_errdetail("specified archive directory does not exist");
+		return false;
+	}
+
+	return true;
+}
+
+bool
+basic_archive(const char *file, const char *path)
+{
+	char destination[MAXPGPATH];
+	char temp[MAXPGPATH];
+	struct stat st;
+
+	Assert(file);
+	Assert(path);
+
+	ereport(DEBUG3,
+			(errmsg("archiving \"%s\" via basic_archive", file)));
+
+	if (archive_directory == NULL || archive_directory[0] == '\0')
+	{
+		ereport(WARNING,
+				(errmsg("\"basic_archive.archive_directory\" not specified")));
+		return false;
+	}
+
+#define TEMP_FILE_NAME ("archtemp")
+
+	if (strlen(archive_directory) + Max(strlen(file), strlen(TEMP_FILE_NAME)) + 2 >= MAXPGPATH)
+	{
+		ereport(WARNING,
+				(errmsg("archive destination path too long")));
+		return false;
+	}
+
+	snprintf(destination, MAXPGPATH, "%s/%s", archive_directory, file);
+	snprintf(temp, MAXPGPATH, "%s/%s", archive_directory, TEMP_FILE_NAME);
+
+	/*
+	 * First, check if the file has already been archived.  If it has,
+	 * just fail because something might be wrong.
+	 */
+	if (stat(destination, &st) == 0)
+	{
+		ereport(WARNING,
+				(errmsg("archive file \"%s\" already exists",
+						destination)));
+		return false;
+	}
+	else if (errno != ENOENT)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not stat file \"%s\": %m",
+						destination)));
+
+	/*
+	 * Remove pre-existing temporary file, if one exists.
+	 */
+	if (unlink(temp) != 0 && errno != ENOENT)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not unlink file \"%s\": %m",
+						temp)));
+
+	/*
+	 * Copy the file to its temporary destination.
+	 */
+	copy_file(path, temp);
+
+	/*
+	 * Sync the temporary file to disk and move it to its final
+	 * destination.
+	 */
+	(void) durable_rename(temp, destination, ERROR);
+
+	ereport(DEBUG1,
+			(errmsg("archived \"%s\" via basic_archive", file)));
+
+	return true;
+}
+
+static void
+copy_file(const char *src, const char *dst)
+{
+	int srcfd;
+	int dstfd;
+	int nbytes;
+	char *buf;
+
+#define COPY_BUF_SIZE (64 * 1024)
+
+	buf = palloc(COPY_BUF_SIZE);
+
+	srcfd = OpenTransientFile(src, O_RDONLY | PG_BINARY);
+	if (srcfd < 0)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not open file \"%s\": %m", src)));
+
+	dstfd = OpenTransientFile(dst, O_RDWR | O_CREAT | O_EXCL | PG_BINARY);
+	if (dstfd < 0)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not open file \"%s\": %m", dst)));
+
+	for (;;)
+	{
+		nbytes = read(srcfd, buf, COPY_BUF_SIZE);
+		if (nbytes < 0)
+			ereport(ERROR,
+					(errcode_for_file_access(),
+					 errmsg("could not read file \"%s\": %m", src)));
+
+		if (nbytes == 0)
+			break;
+
+		errno = 0;
+		if ((int) write(dstfd, buf, nbytes) != nbytes)
+		{
+			/* if write didn't set errno, assume problem is no disk space */
+			if (errno == 0)
+				errno = ENOSPC;
+			ereport(ERROR,
+					(errcode_for_file_access(),
+					 errmsg("could not write to file \"%s\": %m", dst)));
+		}
+	}
+
+	if (CloseTransientFile(dstfd) != 0)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not close file \"%s\": %m", dst)));
+
+	if (CloseTransientFile(srcfd) != 0)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not close file \"%s\": %m", src)));
+
+	pfree(buf);
+}
-- 
2.16.6

