From 089594ce3c5a5f7de1a141684dc63d1cc687d9e2 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <bossartn@amazon.com>
Date: Fri, 19 Nov 2021 01:05:43 +0000
Subject: [PATCH v10 3/4] Add test archive module.

---
 src/test/modules/Makefile                          |   1 +
 src/test/modules/basic_archive/.gitignore          |   4 +
 src/test/modules/basic_archive/Makefile            |  20 +++
 src/test/modules/basic_archive/basic_archive.c     | 185 +++++++++++++++++++++
 src/test/modules/basic_archive/basic_archive.conf  |   3 +
 .../basic_archive/expected/basic_archive.out       |  29 ++++
 .../modules/basic_archive/sql/basic_archive.sql    |  22 +++
 7 files changed, 264 insertions(+)
 create mode 100644 src/test/modules/basic_archive/.gitignore
 create mode 100644 src/test/modules/basic_archive/Makefile
 create mode 100644 src/test/modules/basic_archive/basic_archive.c
 create mode 100644 src/test/modules/basic_archive/basic_archive.conf
 create mode 100644 src/test/modules/basic_archive/expected/basic_archive.out
 create mode 100644 src/test/modules/basic_archive/sql/basic_archive.sql

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index dffc79b2d9..b49e508a2c 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -5,6 +5,7 @@ top_builddir = ../../..
 include $(top_builddir)/src/Makefile.global
 
 SUBDIRS = \
+		  basic_archive \
 		  brin \
 		  commit_ts \
 		  delay_execution \
diff --git a/src/test/modules/basic_archive/.gitignore b/src/test/modules/basic_archive/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/src/test/modules/basic_archive/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/basic_archive/Makefile b/src/test/modules/basic_archive/Makefile
new file mode 100644
index 0000000000..ffbf846b68
--- /dev/null
+++ b/src/test/modules/basic_archive/Makefile
@@ -0,0 +1,20 @@
+# src/test/modules/basic_archive/Makefile
+
+MODULES = basic_archive
+PGFILEDESC = "basic_archive - basic archive module"
+
+REGRESS = basic_archive
+REGRESS_OPTS = --temp-config $(top_srcdir)/src/test/modules/basic_archive/basic_archive.conf
+
+NO_INSTALLCHECK = 1
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/basic_archive
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/basic_archive/basic_archive.c b/src/test/modules/basic_archive/basic_archive.c
new file mode 100644
index 0000000000..84e923bb7e
--- /dev/null
+++ b/src/test/modules/basic_archive/basic_archive.c
@@ -0,0 +1,185 @@
+/*-------------------------------------------------------------------------
+ *
+ * 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
+ *	  src/test/modules/basic_archive/basic_archive.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include "miscadmin.h"
+#include "postmaster/pgarch.h"
+#include "storage/copydir.h"
+#include "storage/fd.h"
+#include "utils/guc.h"
+
+PG_MODULE_MAGIC;
+
+void _PG_init(void);
+void _PG_archive_module_init(ArchiveModuleCallbacks *cb);
+
+static char *archive_directory = NULL;
+
+static bool basic_archive_configured(void);
+static bool basic_archive_file(const char *file, const char *path);
+static bool check_archive_directory(char **newval, void **extra, GucSource source);
+
+/*
+ * _PG_init
+ *
+ * Defines the module's GUC.
+ */
+void
+_PG_init(void)
+{
+	DefineCustomStringVariable("basic_archive.archive_directory",
+							   gettext_noop("Archive file destination directory."),
+							   NULL,
+							   &archive_directory,
+							   "",
+							   PGC_SIGHUP,
+							   0,
+							   check_archive_directory, NULL, NULL);
+
+	EmitWarningsOnPlaceholders("basic_archive");
+}
+
+/*
+ * _PG_archive_module_init
+ *
+ * Returns the module's archiving callbacks.
+ */
+void
+_PG_archive_module_init(ArchiveModuleCallbacks *cb)
+{
+	AssertVariableIsOfType(&_PG_archive_module_init, ArchiveModuleInit);
+
+	cb->check_configured_cb = basic_archive_configured;
+	cb->archive_file_cb = basic_archive_file;
+}
+
+/*
+ * check_archive_directory
+ *
+ * Checks that the provided archive directory exists.
+ */
+static bool
+check_archive_directory(char **newval, void **extra, GucSource source)
+{
+	struct stat st;
+
+	/*
+	 * The default value is an empty string, so we have to accept that value.
+	 * Our check_configured callback also checks for this and prevents archiving
+	 * from proceeding if it is still empty.
+	 */
+	if (*newval == NULL || *newval[0] == '\0')
+		return true;
+
+	/*
+	 * Make sure the file paths won't be too long.  The docs indicate that the
+	 * file names to be archived can be up to 64 characters long.
+	 */
+	if (strlen(*newval) + 64 + 2 >= MAXPGPATH)
+	{
+		GUC_check_errdetail("archive directory too long");
+		return false;
+	}
+
+	/*
+	 * Do a basic sanity check that the specified archive directory exists.  It
+	 * could be removed at some point in the future, so we still need to be
+	 * prepared for it not to exist in the actual archiving logic.
+	 */
+	if (stat(*newval, &st) != 0 || !S_ISDIR(st.st_mode))
+	{
+		GUC_check_errdetail("specified archive directory does not exist");
+		return false;
+	}
+
+	return true;
+}
+
+/*
+ * basic_archive_configured
+ *
+ * Checks that archive_directory is not blank.
+ */
+static bool
+basic_archive_configured(void)
+{
+	return archive_directory != NULL && archive_directory[0] != '\0';
+}
+
+/*
+ * basic_archive_file
+ *
+ * Archives one file.
+ */
+static bool
+basic_archive_file(const char *file, const char *path)
+{
+	char destination[MAXPGPATH];
+	char temp[MAXPGPATH];
+	struct stat st;
+
+	ereport(DEBUG3,
+			(errmsg("archiving \"%s\" via basic_archive", file)));
+
+	snprintf(destination, MAXPGPATH, "%s/%s", archive_directory, file);
+	snprintf(temp, MAXPGPATH, "%s/%s", archive_directory, "archtemp");
+
+	/*
+	 * First, check if the file has already been archived.  If the archive file
+	 * already exists, something might be wrong, so we just fail.
+	 */
+	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(unconstify(char *, path), temp);
+
+	/*
+	 * Sync the temporary file to disk and move it to its final destination.
+	 */
+	(void) durable_rename_excl(temp, destination, ERROR);
+
+	ereport(DEBUG1,
+			(errmsg("archived \"%s\" via basic_archive", file)));
+
+	return true;
+}
diff --git a/src/test/modules/basic_archive/basic_archive.conf b/src/test/modules/basic_archive/basic_archive.conf
new file mode 100644
index 0000000000..b26b2d4144
--- /dev/null
+++ b/src/test/modules/basic_archive/basic_archive.conf
@@ -0,0 +1,3 @@
+archive_mode = 'on'
+archive_library = 'basic_archive'
+basic_archive.archive_directory = '.'
diff --git a/src/test/modules/basic_archive/expected/basic_archive.out b/src/test/modules/basic_archive/expected/basic_archive.out
new file mode 100644
index 0000000000..0015053e0f
--- /dev/null
+++ b/src/test/modules/basic_archive/expected/basic_archive.out
@@ -0,0 +1,29 @@
+CREATE TABLE test (a INT);
+SELECT 1 FROM pg_switch_wal();
+ ?column? 
+----------
+        1
+(1 row)
+
+DO $$
+DECLARE
+	archived bool;
+	loops int := 0;
+BEGIN
+	LOOP
+		archived := count(*) > 0 FROM pg_ls_dir('.', false, false) a
+			WHERE a ~ '^[0-9A-F]{24}$';
+		IF archived OR loops > 120 * 10 THEN EXIT; END IF;
+		PERFORM pg_sleep(0.1);
+		loops := loops + 1;
+	END LOOP;
+END
+$$;
+SELECT count(*) > 0 FROM pg_ls_dir('.', false, false) a
+	WHERE a ~ '^[0-9A-F]{24}$';
+ ?column? 
+----------
+ t
+(1 row)
+
+DROP TABLE test;
diff --git a/src/test/modules/basic_archive/sql/basic_archive.sql b/src/test/modules/basic_archive/sql/basic_archive.sql
new file mode 100644
index 0000000000..14e236d57a
--- /dev/null
+++ b/src/test/modules/basic_archive/sql/basic_archive.sql
@@ -0,0 +1,22 @@
+CREATE TABLE test (a INT);
+SELECT 1 FROM pg_switch_wal();
+
+DO $$
+DECLARE
+	archived bool;
+	loops int := 0;
+BEGIN
+	LOOP
+		archived := count(*) > 0 FROM pg_ls_dir('.', false, false) a
+			WHERE a ~ '^[0-9A-F]{24}$';
+		IF archived OR loops > 120 * 10 THEN EXIT; END IF;
+		PERFORM pg_sleep(0.1);
+		loops := loops + 1;
+	END LOOP;
+END
+$$;
+
+SELECT count(*) > 0 FROM pg_ls_dir('.', false, false) a
+	WHERE a ~ '^[0-9A-F]{24}$';
+
+DROP TABLE test;
-- 
2.16.6

