Add --tablespace option to reindexdb

From ad14153ca7801d89bf20706660069b22a09788bc Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 26 Feb 2021 15:35:53 +0900
Subject: [PATCH] Add --tablespace option to reindexdb

---
 src/bin/scripts/reindexdb.c        | 113 ++++++++++++++++++++---------
 src/bin/scripts/t/090_reindexdb.pl |  69 ++++++++++++++++--
 doc/src/sgml/ref/reindexdb.sgml    |  10 +++
 3 files changed, 152 insertions(+), 40 deletions(-)

diff --git a/src/bin/scripts/reindexdb.c b/src/bin/scripts/reindexdb.c
index 9f072ac49a..d2b746f185 100644
--- a/src/bin/scripts/reindexdb.c
+++ b/src/bin/scripts/reindexdb.c
@@ -40,14 +40,15 @@ static void reindex_one_database(const ConnParams *cparams, ReindexType type,
 								 SimpleStringList *user_list,
 								 const char *progname,
 								 bool echo, bool verbose, bool concurrently,
-								 int concurrentCons);
+								 int concurrentCons, const char *tablespace);
 static void reindex_all_databases(ConnParams *cparams,
 								  const char *progname, bool echo,
 								  bool quiet, bool verbose, bool concurrently,
-								  int concurrentCons);
+								  int concurrentCons, const char *tablespace);
 static void run_reindex_command(PGconn *conn, ReindexType type,
 								const char *name, bool echo, bool verbose,
-								bool concurrently, bool async);
+								bool concurrently, bool async,
+								const char *tablespace);
 
 static void help(const char *progname);
 
@@ -72,6 +73,7 @@ main(int argc, char *argv[])
 		{"verbose", no_argument, NULL, 'v'},
 		{"concurrently", no_argument, NULL, 1},
 		{"maintenance-db", required_argument, NULL, 2},
+		{"tablespace", required_argument, NULL, 3},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -84,6 +86,7 @@ main(int argc, char *argv[])
 	const char *host = NULL;
 	const char *port = NULL;
 	const char *username = NULL;
+	const char *tablespace = NULL;
 	enum trivalue prompt_password = TRI_DEFAULT;
 	ConnParams	cparams;
 	bool		syscatalog = false;
@@ -164,6 +167,9 @@ main(int argc, char *argv[])
 			case 2:
 				maintenance_db = pg_strdup(optarg);
 				break;
+			case 3:
+				tablespace = pg_strdup(optarg);
+				break;
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
@@ -228,7 +234,7 @@ main(int argc, char *argv[])
 		cparams.dbname = maintenance_db;
 
 		reindex_all_databases(&cparams, progname, echo, quiet, verbose,
-							  concurrently, concurrentCons);
+							  concurrently, concurrentCons, tablespace);
 	}
 	else if (syscatalog)
 	{
@@ -268,7 +274,7 @@ main(int argc, char *argv[])
 
 		reindex_one_database(&cparams, REINDEX_SYSTEM, NULL,
 							 progname, echo, verbose,
-							 concurrently, 1);
+							 concurrently, 1, tablespace);
 	}
 	else
 	{
@@ -298,17 +304,17 @@ main(int argc, char *argv[])
 		if (schemas.head != NULL)
 			reindex_one_database(&cparams, REINDEX_SCHEMA, &schemas,
 								 progname, echo, verbose,
-								 concurrently, concurrentCons);
+								 concurrently, concurrentCons, tablespace);
 
 		if (indexes.head != NULL)
 			reindex_one_database(&cparams, REINDEX_INDEX, &indexes,
 								 progname, echo, verbose,
-								 concurrently, 1);
+								 concurrently, 1, tablespace);
 
 		if (tables.head != NULL)
 			reindex_one_database(&cparams, REINDEX_TABLE, &tables,
 								 progname, echo, verbose,
-								 concurrently, concurrentCons);
+								 concurrently, concurrentCons, tablespace);
 
 		/*
 		 * reindex database only if neither index nor table nor schema is
@@ -317,7 +323,7 @@ main(int argc, char *argv[])
 		if (indexes.head == NULL && tables.head == NULL && schemas.head == NULL)
 			reindex_one_database(&cparams, REINDEX_DATABASE, NULL,
 								 progname, echo, verbose,
-								 concurrently, concurrentCons);
+								 concurrently, concurrentCons, tablespace);
 	}
 
 	exit(0);
@@ -327,7 +333,8 @@ static void
 reindex_one_database(const ConnParams *cparams, ReindexType type,
 					 SimpleStringList *user_list,
 					 const char *progname, bool echo,
-					 bool verbose, bool concurrently, int concurrentCons)
+					 bool verbose, bool concurrently, int concurrentCons,
+					 const char *tablespace)
 {
 	PGconn	   *conn;
 	SimpleStringListCell *cell;
@@ -348,6 +355,14 @@ reindex_one_database(const ConnParams *cparams, ReindexType type,
 		exit(1);
 	}
 
+	if (tablespace && PQserverVersion(conn) < 140000)
+	{
+		PQfinish(conn);
+		pg_log_error("cannot use the \"%s\" option on server versions older than PostgreSQL %s",
+					 "tablespace", "14");
+		exit(1);
+	}
+
 	if (!parallel)
 	{
 		switch (process_type)
@@ -386,7 +401,8 @@ reindex_one_database(const ConnParams *cparams, ReindexType type,
 					pg_log_warning("cannot reindex system catalogs concurrently, skipping all");
 				else
 					run_reindex_command(conn, REINDEX_SYSTEM, PQdb(conn), echo,
-										verbose, concurrently, false);
+										verbose, concurrently, false,
+										tablespace);
 
 				/* Build a list of relations from the database */
 				process_list = get_parallel_object_list(conn, process_type,
@@ -468,7 +484,7 @@ reindex_one_database(const ConnParams *cparams, ReindexType type,
 
 		ParallelSlotSetHandler(free_slot, TableCommandResultHandler, NULL);
 		run_reindex_command(free_slot->connection, process_type, objname,
-							echo, verbose, concurrently, true);
+							echo, verbose, concurrently, true, tablespace);
 
 		cell = cell->next;
 	} while (cell != NULL);
@@ -492,8 +508,12 @@ finish:
 
 static void
 run_reindex_command(PGconn *conn, ReindexType type, const char *name,
-					bool echo, bool verbose, bool concurrently, bool async)
+					bool echo, bool verbose, bool concurrently, bool async,
+					const char *tablespace)
 {
+	const char *paren = "(";
+	const char *comma = ", ";
+	const char *sep = paren;
 	PQExpBufferData sql;
 	bool		status;
 
@@ -505,7 +525,20 @@ run_reindex_command(PGconn *conn, ReindexType type, const char *name,
 	appendPQExpBufferStr(&sql, "REINDEX ");
 
 	if (verbose)
-		appendPQExpBufferStr(&sql, "(VERBOSE) ");
+	{
+		appendPQExpBuffer(&sql, "%sVERBOSE", sep);
+		sep = comma;
+	}
+
+	if (tablespace)
+	{
+		Assert(PQserverVersion(conn) >= 140000);
+		appendPQExpBuffer(&sql, "%sTABLESPACE %s", sep, fmtId(tablespace));
+		sep = comma;
+	}
+
+	if (sep != paren)
+		appendPQExpBufferStr(&sql, ") ");
 
 	/* object type */
 	switch (type)
@@ -527,8 +560,16 @@ run_reindex_command(PGconn *conn, ReindexType type, const char *name,
 			break;
 	}
 
+	/*
+	 * Parenthesized grammar is only supported for CONCURRENTLY since
+	 * PostgreSQL 14.  Since 12, CONCURRENTLY can be specified after the
+	 * object type.
+	 */
 	if (concurrently)
+	{
+		Assert(PQserverVersion(conn) >= 120000);
 		appendPQExpBufferStr(&sql, "CONCURRENTLY ");
+	}
 
 	/* object name */
 	switch (type)
@@ -716,7 +757,8 @@ get_parallel_object_list(PGconn *conn, ReindexType type,
 static void
 reindex_all_databases(ConnParams *cparams,
 					  const char *progname, bool echo, bool quiet, bool verbose,
-					  bool concurrently, int concurrentCons)
+					  bool concurrently, int concurrentCons,
+					  const char *tablespace)
 {
 	PGconn	   *conn;
 	PGresult   *result;
@@ -740,7 +782,7 @@ reindex_all_databases(ConnParams *cparams,
 
 		reindex_one_database(cparams, REINDEX_DATABASE, NULL,
 							 progname, echo, verbose, concurrently,
-							 concurrentCons);
+							 concurrentCons, tablespace);
 	}
 
 	PQclear(result);
@@ -753,26 +795,27 @@ help(const char *progname)
 	printf(_("Usage:\n"));
 	printf(_("  %s [OPTION]... [DBNAME]\n"), progname);
 	printf(_("\nOptions:\n"));
-	printf(_("  -a, --all                 reindex all databases\n"));
-	printf(_("      --concurrently        reindex concurrently\n"));
-	printf(_("  -d, --dbname=DBNAME       database to reindex\n"));
-	printf(_("  -e, --echo                show the commands being sent to the server\n"));
-	printf(_("  -i, --index=INDEX         recreate specific index(es) only\n"));
-	printf(_("  -j, --jobs=NUM            use this many concurrent connections to reindex\n"));
-	printf(_("  -q, --quiet               don't write any messages\n"));
-	printf(_("  -s, --system              reindex system catalogs\n"));
-	printf(_("  -S, --schema=SCHEMA       reindex specific schema(s) only\n"));
-	printf(_("  -t, --table=TABLE         reindex specific table(s) only\n"));
-	printf(_("  -v, --verbose             write a lot of output\n"));
-	printf(_("  -V, --version             output version information, then exit\n"));
-	printf(_("  -?, --help                show this help, then exit\n"));
+	printf(_("  -a, --all                    reindex all databases\n"));
+	printf(_("      --concurrently           reindex concurrently\n"));
+	printf(_("  -d, --dbname=DBNAME          database to reindex\n"));
+	printf(_("  -e, --echo                   show the commands being sent to the server\n"));
+	printf(_("  -i, --index=INDEX            recreate specific index(es) only\n"));
+	printf(_("  -j, --jobs=NUM               use this many concurrent connections to reindex\n"));
+	printf(_("  -q, --quiet                  don't write any messages\n"));
+	printf(_("  -s, --system                 reindex system catalogs\n"));
+	printf(_("  -S, --schema=SCHEMA          reindex specific schema(s) only\n"));
+	printf(_("  -t, --table=TABLE            reindex specific table(s) only\n"));
+	printf(_("      --tablespace=TABLESPACE  reindex on specified tablespace\n"));
+	printf(_("  -v, --verbose                write a lot of output\n"));
+	printf(_("  -V, --version                output version information, then exit\n"));
+	printf(_("  -?, --help                   show this help, then exit\n"));
 	printf(_("\nConnection options:\n"));
-	printf(_("  -h, --host=HOSTNAME       database server host or socket directory\n"));
-	printf(_("  -p, --port=PORT           database server port\n"));
-	printf(_("  -U, --username=USERNAME   user name to connect as\n"));
-	printf(_("  -w, --no-password         never prompt for password\n"));
-	printf(_("  -W, --password            force password prompt\n"));
-	printf(_("  --maintenance-db=DBNAME   alternate maintenance database\n"));
+	printf(_("  -h, --host=HOSTNAME          database server host or socket directory\n"));
+	printf(_("  -p, --port=PORT              database server port\n"));
+	printf(_("  -U, --username=USERNAME      user name to connect as\n"));
+	printf(_("  -w, --no-password            never prompt for password\n"));
+	printf(_("  -W, --password               force password prompt\n"));
+	printf(_("  --maintenance-db=DBNAME      alternate maintenance database\n"));
 	printf(_("\nRead the description of the SQL command REINDEX for details.\n"));
 	printf(_("\nReport bugs to <%s>.\n"), PACKAGE_BUGREPORT);
 	printf(_("%s home page: <%s>\n"), PACKAGE_NAME, PACKAGE_URL);
diff --git a/src/bin/scripts/t/090_reindexdb.pl b/src/bin/scripts/t/090_reindexdb.pl
index 87417c86ff..6946268209 100644
--- a/src/bin/scripts/t/090_reindexdb.pl
+++ b/src/bin/scripts/t/090_reindexdb.pl
@@ -3,7 +3,7 @@ use warnings;
 
 use PostgresNode;
 use TestLib;
-use Test::More tests => 44;
+use Test::More tests => 54;
 
 program_help_ok('reindexdb');
 program_version_ok('reindexdb');
@@ -15,17 +15,38 @@ $node->start;
 
 $ENV{PGOPTIONS} = '--client-min-messages=WARNING';
 
+# Create a tablespace for testing.
+my $ts = $node->basedir . '/regress_reindex_tbspace';
+mkdir $ts or die "cannot create directory $ts";
+# this takes care of WIN-specific path issues
+my $ets     = TestLib::perl2host($ts);
+my $tbspace = 'reindex_tbspace';
+$node->safe_psql('postgres', "CREATE TABLESPACE $tbspace LOCATION '$ets';");
+
 $node->issues_sql_like(
 	[ 'reindexdb', 'postgres' ],
 	qr/statement: REINDEX DATABASE postgres;/,
 	'SQL REINDEX run');
 
+# Use text as data type to get a toast table.
 $node->safe_psql('postgres',
-	'CREATE TABLE test1 (a int); CREATE INDEX test1x ON test1 (a);');
+	'CREATE TABLE test1 (a text); CREATE INDEX test1x ON test1 (a);');
+# Collect toast table and index names of this relation, for later use.
+my $toast_table = $node->safe_psql('postgres',
+	"SELECT reltoastrelid::regclass FROM pg_class WHERE oid = 'test1'::regclass;"
+);
+my $toast_index = $node->safe_psql('postgres',
+	"SELECT indexrelid::regclass FROM pg_index WHERE indrelid = '$toast_table'::regclass;"
+);
+
 $node->issues_sql_like(
 	[ 'reindexdb', '-t', 'test1', 'postgres' ],
 	qr/statement: REINDEX TABLE public\.test1;/,
 	'reindex specific table');
+$node->issues_sql_like(
+	[ 'reindexdb', '-t', 'test1', '--tablespace', $tbspace, 'postgres' ],
+	qr/statement: REINDEX \(TABLESPACE $tbspace\) TABLE public\.test1;/,
+	'reindex specific table on tablespace');
 $node->issues_sql_like(
 	[ 'reindexdb', '-i', 'test1x', 'postgres' ],
 	qr/statement: REINDEX INDEX public\.test1x;/,
@@ -42,6 +63,13 @@ $node->issues_sql_like(
 	[ 'reindexdb', '-v', '-t', 'test1', 'postgres' ],
 	qr/statement: REINDEX \(VERBOSE\) TABLE public\.test1;/,
 	'reindex with verbose output');
+$node->issues_sql_like(
+	[
+		'reindexdb', '-v', '-t', 'test1', '--tablespace', $tbspace,
+		'postgres'
+	],
+	qr/statement: REINDEX \(VERBOSE, TABLESPACE $tbspace\) TABLE public\.test1;/,
+	'reindex with verbose output and tablespace');
 
 # the same with --concurrently
 $node->issues_sql_like(
@@ -64,9 +92,40 @@ $node->issues_sql_like(
 $node->command_fails([ 'reindexdb', '--concurrently', '-s', 'postgres' ],
 	'reindex system tables concurrently');
 $node->issues_sql_like(
-	[ 'reindexdb', '-v', '-t', 'test1', 'postgres' ],
-	qr/statement: REINDEX \(VERBOSE\) TABLE public\.test1;/,
-	'reindex with verbose output');
+	[ 'reindexdb', '--concurrently', '-v', '-t', 'test1', 'postgres' ],
+	qr/statement: REINDEX \(VERBOSE\) TABLE CONCURRENTLY public\.test1;/,
+	'reindex concurrently with verbose output');
+$node->issues_sql_like(
+	[
+		'reindexdb', '--concurrently', '-v',     '-t',
+		'test1',     '--tablespace',   $tbspace, 'postgres'
+	],
+	qr/statement: REINDEX \(VERBOSE, TABLESPACE $tbspace\) TABLE CONCURRENTLY public\.test1;/,
+	'reindex concurrently with verbose output and tablespace');
+
+# REINDEX TABLESPACE on toast indexes and tables fails.  This is not
+# part of the main regression test suite as these have unpredictable
+# names, and CONCURRENTLY cannot be used in transaction blocks, preventing
+# the use of TRY/CATCH blocks in a custom function to filter error
+# messages.
+$node->command_fails(
+	[ 'reindexdb', '-t', $toast_table, '--tablespace', $tbspace, 'postgres' ],
+	'reindex toast table with tablespace');
+$node->command_fails(
+	[
+		'reindexdb',    '--concurrently', '-t', $toast_table,
+		'--tablespace', $tbspace,         'postgres'
+	],
+	'reindex toast table concurrently with tablespace');
+$node->command_fails(
+	[ 'reindexdb', '-i', $toast_index, '--tablespace', $tbspace, 'postgres' ],
+	'reindex toast index with tablespace');
+$node->command_fails(
+	[
+		'reindexdb',    '--concurrently', '-i', $toast_index,
+		'--tablespace', $tbspace,         'postgres'
+	],
+	'reindex toast index concurrently with tablespace');
 
 # connection strings
 $node->command_ok([qw(reindexdb --echo --table=pg_am dbname=template1)],
diff --git a/doc/src/sgml/ref/reindexdb.sgml b/doc/src/sgml/ref/reindexdb.sgml
index a3b0f7ce31..6b1784584d 100644
--- a/doc/src/sgml/ref/reindexdb.sgml
+++ b/doc/src/sgml/ref/reindexdb.sgml
@@ -237,6 +237,16 @@ PostgreSQL documentation
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+      <term><option>--tablespace=<replaceable class="parameter">tablespace</replaceable></option></term>
+      <listitem>
+       <para>
+        Specifies the tablespace to reindex on. (This name is processed as
+        a double-quoted identifier.)
+       </para>
+      </listitem>
+     </varlistentry>
+
     <varlistentry>
      <term><option>-v</option></term>
      <term><option>--verbose</option></term>
-- 
2.30.1