From 80aada5822c7cda6f097f08bba98e069fa20ea77 Mon Sep 17 00:00:00 2001
From: Mark Dilger <mark.dilger@enterprisedb.com>
Date: Sat, 17 Apr 2021 12:40:11 -0700
Subject: [PATCH v2] Adding --install-missing option to pg_amcheck

To ease the checking of databases where amcheck is not yet
installed, adding an option that will create the extension if it
does not yet exist.
---
 doc/src/sgml/ref/pg_amcheck.sgml | 17 ++++++++++++
 src/bin/pg_amcheck/pg_amcheck.c  | 44 ++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/doc/src/sgml/ref/pg_amcheck.sgml b/doc/src/sgml/ref/pg_amcheck.sgml
index d01e26faa8..a38e3ce7b3 100644
--- a/doc/src/sgml/ref/pg_amcheck.sgml
+++ b/doc/src/sgml/ref/pg_amcheck.sgml
@@ -217,6 +217,23 @@ PostgreSQL documentation
      </listitem>
     </varlistentry>
 
+    <varlistentry>
+     <term><option>--install-missing</option></term>
+     <term><option>--install-missing=<replaceable class="parameter">schema</replaceable></option></term>
+     <listitem>
+      <para>
+	   Install any missing extensions that are required to check the
+	   database(s).  If not yet installed, each extension's objects will be
+	   installed into the given
+	   <replaceable class="parameter">schema</replaceable>, or if not specified
+	   into schema <literal>pg_catalog</literal>.
+      </para>
+      <para>
+	   At present, the only required extension is <xref linkend="amcheck"/>.
+      </para>
+     </listitem>
+    </varlistentry>
+
     <varlistentry>
      <term><option>-j <replaceable class="parameter">num</replaceable></option></term>
      <term><option>--jobs=<replaceable class="parameter">num</replaceable></option></term>
diff --git a/src/bin/pg_amcheck/pg_amcheck.c b/src/bin/pg_amcheck/pg_amcheck.c
index 587a79a1a6..f9e5ff0d4f 100644
--- a/src/bin/pg_amcheck/pg_amcheck.c
+++ b/src/bin/pg_amcheck/pg_amcheck.c
@@ -61,6 +61,13 @@ typedef struct AmcheckOptions
 	bool		show_progress;
 	int			jobs;
 
+	/*
+	 * Whether to install missing extensions, and optionally the name of the
+	 * schema in which to install the extension's objects.
+	 */
+	bool		install_missing;
+	char	   *install_schema;
+
 	/* Objects to check or not to check, as lists of PatternInfo structs. */
 	PatternInfoArray include;
 	PatternInfoArray exclude;
@@ -109,6 +116,8 @@ static AmcheckOptions opts = {
 	.strict_names = true,
 	.show_progress = false,
 	.jobs = 1,
+	.install_missing = false,
+	.install_schema = NULL,
 	.include = {NULL, 0},
 	.exclude = {NULL, 0},
 	.excludetbl = false,
@@ -259,6 +268,7 @@ main(int argc, char *argv[])
 		{"no-strict-names", no_argument, NULL, 10},
 		{"heapallindexed", no_argument, NULL, 11},
 		{"parent-check", no_argument, NULL, 12},
+		{"install-missing", optional_argument, NULL, 13},
 
 		{NULL, 0, NULL, 0}
 	};
@@ -440,6 +450,11 @@ main(int argc, char *argv[])
 			case 12:
 				opts.parent_check = true;
 				break;
+			case 13:
+				opts.install_missing = true;
+				if (optarg)
+					opts.install_schema = pg_strdup(optarg);
+				break;
 			default:
 				fprintf(stderr,
 						_("Try \"%s --help\" for more information.\n"),
@@ -549,6 +564,34 @@ main(int argc, char *argv[])
 			conn = connectDatabase(&cparams, progname, opts.echo, false, true);
 		}
 
+		/*
+		 * Optionally install amcheck if not already installed in this
+		 * database.
+		 */
+		if (opts.install_missing && opts.install_schema != NULL)
+		{
+			char *schema;
+			char *install_sql;
+
+			/*
+			 * Must re-escape the schema name for each database, as the
+			 * escaping rules may change.
+			 */
+			schema = PQescapeIdentifier(conn, opts.install_schema,
+										strlen(opts.install_schema));
+			install_sql = psprintf("CREATE EXTENSION IF NOT EXISTS amcheck WITH SCHEMA %s",
+								   schema);
+
+			executeCommand(conn, install_sql, opts.echo);
+			pfree(install_sql);
+			pfree(schema);
+		}
+		else if (opts.install_missing)
+		{
+			executeCommand(conn, "CREATE EXTENSION IF NOT EXISTS amcheck WITH SCHEMA pg_catalog",
+			opts.echo);
+		}
+
 		/*
 		 * Verify that amcheck is installed for this next database.  User
 		 * error could result in a database not having amcheck that should
@@ -1159,6 +1202,7 @@ help(const char *progname)
 	printf(_("  -V, --version                  output version information, then exit\n"));
 	printf(_("  -P, --progress                 show progress information\n"));
 	printf(_("  -?, --help                     show this help, then exit\n"));
+	printf(_("      --install-missing          install missing extensions\n"));
 
 	printf(_("\nReport bugs to <%s>.\n"), PACKAGE_BUGREPORT);
 	printf(_("%s home page: <%s>\n"), PACKAGE_NAME, PACKAGE_URL);
-- 
2.21.1 (Apple Git-122.3)

