From c180a56055bb9c61c4bb81d586aeb99241596457 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathandbossart@gmail.com>
Date: Fri, 22 Apr 2022 22:34:38 -0700
Subject: [PATCH v11 1/1] Add --schema and --exclude-schema options to
 vacuumdb.

These two new options can be used to either process all tables in
specific schemas or to skip processing all tables in specific
schemas.  This change also refactors the handling of invalid
combinations of command-line options to a new helper function.

Author: Gilles Darold
Reviewed-by: Justin Pryzby, Nathan Bossart
Discussion: https://postgr.es/m/929fbf3c-24b8-d454-811f-1d5898ab3e91%40migops.com
---
 doc/src/sgml/ref/vacuumdb.sgml        |  66 ++++++++++
 src/bin/scripts/t/100_vacuumdb.pl     |  35 ++++++
 src/bin/scripts/t/101_vacuumdb_all.pl |   3 +
 src/bin/scripts/vacuumdb.c            | 170 +++++++++++++++++++-------
 4 files changed, 233 insertions(+), 41 deletions(-)

diff --git a/doc/src/sgml/ref/vacuumdb.sgml b/doc/src/sgml/ref/vacuumdb.sgml
index 956c0f01cb..841aced3bd 100644
--- a/doc/src/sgml/ref/vacuumdb.sgml
+++ b/doc/src/sgml/ref/vacuumdb.sgml
@@ -39,6 +39,40 @@ PostgreSQL documentation
    <arg choice="opt"><replaceable>dbname</replaceable></arg>
   </cmdsynopsis>
 
+  <cmdsynopsis>
+   <command>vacuumdb</command>
+   <arg rep="repeat"><replaceable>connection-option</replaceable></arg>
+   <arg rep="repeat"><replaceable>option</replaceable></arg>
+
+   <arg choice="plain" rep="repeat">
+    <arg choice="opt">
+     <group choice="plain">
+       <arg choice="plain">
+        <arg choice="opt">
+         <group choice="plain">
+          <arg choice="plain"><option>-n</option></arg>
+          <arg choice="plain"><option>--schema</option></arg>
+         </group>
+         <replaceable>schema</replaceable>
+        </arg>
+       </arg>
+
+       <arg choice="plain">
+        <arg choice="opt">
+         <group choice="plain">
+          <arg choice="plain"><option>-N</option></arg>
+          <arg choice="plain"><option>--exclude-schema</option></arg>
+         </group>
+         <replaceable>schema</replaceable>
+        </arg>
+       </arg>
+     </group>
+    </arg>
+   </arg>
+
+   <arg choice="opt"><replaceable>dbname</replaceable></arg>
+  </cmdsynopsis>
+
   <cmdsynopsis>
    <command>vacuumdb</command>
    <arg rep="repeat"><replaceable>connection-option</replaceable></arg>
@@ -244,6 +278,30 @@ PostgreSQL documentation
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+      <term><option>-n <replaceable class="parameter">schema</replaceable></option></term>
+      <term><option>--schema=<replaceable class="parameter">schema</replaceable></option></term>
+      <listitem>
+       <para>
+        Clean or analyze all tables in
+        <replaceable class="parameter">schema</replaceable> only.  Multiple
+        schemas can be vacuumed by writing multiple <option>-n</option> switches.
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry>
+      <term><option>-N <replaceable class="parameter">schema</replaceable></option></term>
+      <term><option>--exclude-schema=<replaceable class="parameter">schema</replaceable></option></term>
+      <listitem>
+       <para>
+        Do not clean or analyze any tables in
+        <replaceable class="parameter">schema</replaceable>.  Multiple schemas
+        can be excluded by writing multiple <option>-N</option> switches.
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry>
       <term><option>--no-index-cleanup</option></term>
       <listitem>
@@ -619,6 +677,14 @@ PostgreSQL documentation
 <prompt>$ </prompt><userinput>vacuumdb --analyze --verbose --table='foo(bar)' xyzzy</userinput>
 </screen></para>
 
+   <para>
+    To clean all tables in the <literal>foo</literal> and <literal>bar</literal> schemas
+    in a database named <literal>xyzzy</literal>:
+<screen>
+<prompt>$ </prompt><userinput>vacuumdb --schema='foo' --schema='bar' xyzzy</userinput>
+</screen></para>
+
+
  </refsect1>
 
  <refsect1>
diff --git a/src/bin/scripts/t/100_vacuumdb.pl b/src/bin/scripts/t/100_vacuumdb.pl
index 96a818a3c1..e4aac53249 100644
--- a/src/bin/scripts/t/100_vacuumdb.pl
+++ b/src/bin/scripts/t/100_vacuumdb.pl
@@ -103,6 +103,8 @@ $node->safe_psql(
   CREATE TABLE funcidx (x int);
   INSERT INTO funcidx VALUES (0),(1),(2),(3);
   CREATE INDEX i0 ON funcidx ((f1(x)));
+  CREATE SCHEMA "Foo";
+  CREATE TABLE "Foo".bar(id int);
 |);
 $node->command_ok([qw|vacuumdb -Z --table="need""q(uot"(")x") postgres|],
 	'column list');
@@ -146,5 +148,38 @@ $node->issues_sql_like(
 	[ 'vacuumdb', '--min-xid-age', '2147483001', 'postgres' ],
 	qr/GREATEST.*relfrozenxid.*2147483001/,
 	'vacuumdb --table --min-xid-age');
+$node->issues_sql_like(
+	[ 'vacuumdb', '--schema', '"Foo"', 'postgres' ],
+	qr/VACUUM "Foo".bar/,
+	'vacuumdb --schema schema only');
+$node->issues_sql_like(
+	[ 'vacuumdb', '--exclude-schema', '"Foo"', 'postgres' ],
+	qr/(?:(?!VACUUM "Foo".bar).)*/,
+	'vacuumdb --exclude-schema schema');
+$node->command_fails(
+	[ 'vacuumdb', '-n', 'pg_catalog', '-t', 'pg_class', 'postgres' ],
+	'cannot use options -n and -t at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-n', 'pg_catalog', '-N', '"Foo"', 'postgres' ],
+	'cannot use options -n and -N at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-a', '-N', '"Foo"' ],
+	'cannot use options -a and -N at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-a', '-n', '"Foo"' ],
+	'cannot use options -a and -n at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-a', '-t', '"Foo".bar' ],
+	'cannot use options -a and -t at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-a', '-n', '"Foo"', 'postgres' ],
+	'cannot use options -a and -n at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-a', '-d', 'postgres' ],
+	'cannot use options -a and -d at the same time');
+$node->command_fails(
+	[ 'vacuumdb', '-a', 'postgres' ],
+	'cannot use option -a and a dbname as argument at the same time');
+
 
 done_testing();
diff --git a/src/bin/scripts/t/101_vacuumdb_all.pl b/src/bin/scripts/t/101_vacuumdb_all.pl
index 1dcf411767..ccb05166df 100644
--- a/src/bin/scripts/t/101_vacuumdb_all.pl
+++ b/src/bin/scripts/t/101_vacuumdb_all.pl
@@ -15,5 +15,8 @@ $node->issues_sql_like(
 	[ 'vacuumdb', '-a' ],
 	qr/statement: VACUUM.*statement: VACUUM/s,
 	'vacuum all databases');
+$node->command_fails(
+	[ 'vacuumdb', '-a', '-n', 'pg_catalog' ],
+	'cannot vacuum specific schema(s) in all databases');
 
 done_testing();
diff --git a/src/bin/scripts/vacuumdb.c b/src/bin/scripts/vacuumdb.c
index 92f1ffe147..9ff81e5a17 100644
--- a/src/bin/scripts/vacuumdb.c
+++ b/src/bin/scripts/vacuumdb.c
@@ -46,11 +46,23 @@ typedef struct vacuumingOptions
 	bool		process_toast;
 } vacuumingOptions;
 
+/* object filter options */
+typedef enum
+{
+	OBJFILTER_NONE = 0,					/* no filter used */
+	OBJFILTER_ALL_DBS = (1 << 0),		/* -a | --all */
+	OBJFILTER_DATABASE = (1 << 1),		/* -d | --dbname */
+	OBJFILTER_TABLE = (1 << 2),			/* -t | --table */
+	OBJFILTER_SCHEMA = (1 << 3),		/* -n | --schema */
+	OBJFILTER_SCHEMA_EXCLUDE = (1 << 4)	/* -N | --exclude-schema */
+} VacObjFilter;
+
+VacObjFilter objfilter = OBJFILTER_NONE;
 
 static void vacuum_one_database(ConnParams *cparams,
 								vacuumingOptions *vacopts,
 								int stage,
-								SimpleStringList *tables,
+								SimpleStringList *objects,
 								int concurrentCons,
 								const char *progname, bool echo, bool quiet);
 
@@ -68,6 +80,8 @@ static void run_vacuum_command(PGconn *conn, const char *sql, bool echo,
 
 static void help(const char *progname);
 
+void check_objfilter(VacObjFilter curr_option);
+
 /* For analyze-in-stages mode */
 #define ANALYZE_NO_STAGE	-1
 #define ANALYZE_NUM_STAGES	3
@@ -94,6 +108,8 @@ main(int argc, char *argv[])
 		{"verbose", no_argument, NULL, 'v'},
 		{"jobs", required_argument, NULL, 'j'},
 		{"parallel", required_argument, NULL, 'P'},
+		{"schema", required_argument, NULL, 'n'},
+		{"exclude-schema", required_argument, NULL, 'N'},
 		{"maintenance-db", required_argument, NULL, 2},
 		{"analyze-in-stages", no_argument, NULL, 3},
 		{"disable-page-skipping", no_argument, NULL, 4},
@@ -121,8 +137,7 @@ main(int argc, char *argv[])
 	bool		quiet = false;
 	vacuumingOptions vacopts;
 	bool		analyze_in_stages = false;
-	bool		alldb = false;
-	SimpleStringList tables = {NULL, NULL};
+	SimpleStringList objects = {NULL, NULL};
 	int			concurrentCons = 1;
 	int			tbl_count = 0;
 
@@ -140,7 +155,7 @@ main(int argc, char *argv[])
 
 	handle_help_version_opts(argc, argv, "vacuumdb", help);
 
-	while ((c = getopt_long(argc, argv, "h:p:U:wWeqd:zZFat:fvj:P:", long_options, &optindex)) != -1)
+	while ((c = getopt_long(argc, argv, "h:p:U:wWeqd:zZFat:fvj:P:n:N:", long_options, &optindex)) != -1)
 	{
 		switch (c)
 		{
@@ -166,6 +181,7 @@ main(int argc, char *argv[])
 				quiet = true;
 				break;
 			case 'd':
+				check_objfilter(OBJFILTER_DATABASE);
 				dbname = pg_strdup(optarg);
 				break;
 			case 'z':
@@ -178,11 +194,12 @@ main(int argc, char *argv[])
 				vacopts.freeze = true;
 				break;
 			case 'a':
-				alldb = true;
+				check_objfilter(OBJFILTER_ALL_DBS);
 				break;
 			case 't':
 				{
-					simple_string_list_append(&tables, optarg);
+					check_objfilter(OBJFILTER_TABLE);
+					simple_string_list_append(&objects, optarg);
 					tbl_count++;
 					break;
 				}
@@ -202,6 +219,18 @@ main(int argc, char *argv[])
 									  &vacopts.parallel_workers))
 					exit(1);
 				break;
+			case 'n':
+				{
+					check_objfilter(OBJFILTER_SCHEMA);
+					simple_string_list_append(&objects, optarg);
+					break;
+				}
+			case 'N':
+				{
+					check_objfilter(OBJFILTER_SCHEMA_EXCLUDE);
+					simple_string_list_append(&objects, optarg);
+					break;
+				}
 			case 2:
 				maintenance_db = pg_strdup(optarg);
 				break;
@@ -249,6 +278,7 @@ main(int argc, char *argv[])
 	 */
 	if (optind < argc && dbname == NULL)
 	{
+		check_objfilter(OBJFILTER_DATABASE);
 		dbname = argv[optind];
 		optind++;
 	}
@@ -316,13 +346,8 @@ main(int argc, char *argv[])
 	if (tbl_count && (concurrentCons > tbl_count))
 		concurrentCons = tbl_count;
 
-	if (alldb)
+	if (objfilter & OBJFILTER_ALL_DBS)
 	{
-		if (dbname)
-			pg_fatal("cannot vacuum all databases and a specific one at the same time");
-		if (tables.head != NULL)
-			pg_fatal("cannot vacuum specific table(s) in all databases");
-
 		cparams.dbname = maintenance_db;
 
 		vacuum_all_databases(&cparams, &vacopts,
@@ -352,7 +377,7 @@ main(int argc, char *argv[])
 			{
 				vacuum_one_database(&cparams, &vacopts,
 									stage,
-									&tables,
+									&objects,
 									concurrentCons,
 									progname, echo, quiet);
 			}
@@ -360,7 +385,7 @@ main(int argc, char *argv[])
 		else
 			vacuum_one_database(&cparams, &vacopts,
 								ANALYZE_NO_STAGE,
-								&tables,
+								&objects,
 								concurrentCons,
 								progname, echo, quiet);
 	}
@@ -368,6 +393,45 @@ main(int argc, char *argv[])
 	exit(0);
 }
 
+/*
+ * Verify that the filters used at command line are compatible.
+ */
+void
+check_objfilter(VacObjFilter curr_option)
+{
+	VacObjFilter	invalid_combo;
+
+	objfilter |= curr_option;
+
+	invalid_combo = (OBJFILTER_ALL_DBS | OBJFILTER_DATABASE);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot vacuum all databases and a specific one at the same time");
+
+	invalid_combo = (OBJFILTER_ALL_DBS | OBJFILTER_TABLE);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot vacuum specific table(s) in all databases");
+
+	invalid_combo = (OBJFILTER_ALL_DBS | OBJFILTER_SCHEMA);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot vacuum specific schema(s) in all databases");
+
+	invalid_combo = (OBJFILTER_ALL_DBS | OBJFILTER_SCHEMA_EXCLUDE);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot exclude specific schema(s) in all databases");
+
+	invalid_combo = (OBJFILTER_TABLE | OBJFILTER_SCHEMA);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot vacuum all tables in schema(s) and specific table(s) at the same time");
+
+	invalid_combo = (OBJFILTER_TABLE | OBJFILTER_SCHEMA_EXCLUDE);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot vacuum specific table(s) and exclude schema(s) at the same time");
+
+	invalid_combo = (OBJFILTER_SCHEMA | OBJFILTER_SCHEMA_EXCLUDE);
+	if ((objfilter & invalid_combo) == invalid_combo)
+		pg_fatal("cannot vacuum all tables in schema(s) and exclude schema(s) at the same time");
+}
+
 /*
  * vacuum_one_database
  *
@@ -385,7 +449,7 @@ static void
 vacuum_one_database(ConnParams *cparams,
 					vacuumingOptions *vacopts,
 					int stage,
-					SimpleStringList *tables,
+					SimpleStringList *objects,
 					int concurrentCons,
 					const char *progname, bool echo, bool quiet)
 {
@@ -400,7 +464,7 @@ vacuum_one_database(ConnParams *cparams,
 	int			i;
 	int			ntups;
 	bool		failed = false;
-	bool		tables_listed = false;
+	bool		objects_listed = false;
 	bool		has_where = false;
 	const char *initcmd;
 	const char *stage_commands[] = {
@@ -499,31 +563,41 @@ vacuum_one_database(ConnParams *cparams,
 	 * catalog query will fail.
 	 */
 	initPQExpBuffer(&catalog_query);
-	for (cell = tables ? tables->head : NULL; cell; cell = cell->next)
+	for (cell = objects ? objects->head : NULL; cell; cell = cell->next)
 	{
-		char	   *just_table;
-		const char *just_columns;
+		char	   *just_table = NULL;
+		const char *just_columns = NULL;
 
-		/*
-		 * Split relation and column names given by the user, this is used to
-		 * feed the CTE with values on which are performed pre-run validity
-		 * checks as well.  For now these happen only on the relation name.
-		 */
-		splitTableColumnsSpec(cell->val, PQclientEncoding(conn),
-							  &just_table, &just_columns);
-
-		if (!tables_listed)
+		if (!objects_listed)
 		{
 			appendPQExpBufferStr(&catalog_query,
-								 "WITH listed_tables (table_oid, column_list) "
+								 "WITH listed_objects (object_oid, column_list) "
 								 "AS (\n  VALUES (");
-			tables_listed = true;
+			objects_listed = true;
 		}
 		else
 			appendPQExpBufferStr(&catalog_query, ",\n  (");
 
-		appendStringLiteralConn(&catalog_query, just_table, conn);
-		appendPQExpBufferStr(&catalog_query, "::pg_catalog.regclass, ");
+		if (objfilter & (OBJFILTER_SCHEMA | OBJFILTER_SCHEMA_EXCLUDE))
+		{
+			appendStringLiteralConn(&catalog_query, cell->val, conn);
+			appendPQExpBufferStr(&catalog_query, "::pg_catalog.regnamespace, ");
+		}
+
+		if (objfilter & OBJFILTER_TABLE)
+		{
+			/*
+			 * Split relation and column names given by the user, this is used
+			 * to feed the CTE with values on which are performed pre-run
+			 * validity checks as well.  For now these happen only on the
+			 * relation name.
+			 */
+			splitTableColumnsSpec(cell->val, PQclientEncoding(conn),
+								  &just_table, &just_columns);
+
+			appendStringLiteralConn(&catalog_query, just_table, conn);
+			appendPQExpBufferStr(&catalog_query, "::pg_catalog.regclass, ");
+		}
 
 		if (just_columns && just_columns[0] != '\0')
 			appendStringLiteralConn(&catalog_query, just_columns, conn);
@@ -536,13 +610,13 @@ vacuum_one_database(ConnParams *cparams,
 	}
 
 	/* Finish formatting the CTE */
-	if (tables_listed)
+	if (objects_listed)
 		appendPQExpBufferStr(&catalog_query, "\n)\n");
 
 	appendPQExpBufferStr(&catalog_query, "SELECT c.relname, ns.nspname");
 
-	if (tables_listed)
-		appendPQExpBufferStr(&catalog_query, ", listed_tables.column_list");
+	if (objects_listed)
+		appendPQExpBufferStr(&catalog_query, ", listed_objects.column_list");
 
 	appendPQExpBufferStr(&catalog_query,
 						 " FROM pg_catalog.pg_class c\n"
@@ -551,10 +625,22 @@ vacuum_one_database(ConnParams *cparams,
 						 " LEFT JOIN pg_catalog.pg_class t"
 						 " ON c.reltoastrelid OPERATOR(pg_catalog.=) t.oid\n");
 
-	/* Used to match the tables listed by the user */
-	if (tables_listed)
-		appendPQExpBufferStr(&catalog_query, " JOIN listed_tables"
-							 " ON listed_tables.table_oid OPERATOR(pg_catalog.=) c.oid\n");
+	/* Used to match the tables or schemas listed by the user */
+	if (objects_listed)
+	{
+		appendPQExpBufferStr(&catalog_query, " JOIN listed_objects"
+							 " ON listed_objects.object_oid ");
+
+		if (objfilter & OBJFILTER_SCHEMA_EXCLUDE)
+			appendPQExpBufferStr(&catalog_query, "OPERATOR(pg_catalog.!=) ");
+		else
+			appendPQExpBufferStr(&catalog_query, "OPERATOR(pg_catalog.=) ");
+
+		if (objfilter & OBJFILTER_TABLE)
+			appendPQExpBufferStr(&catalog_query, "c.oid\n");
+		else
+			appendPQExpBufferStr(&catalog_query, "ns.oid\n");
+	}
 
 	/*
 	 * If no tables were listed, filter for the relevant relation types.  If
@@ -562,7 +648,7 @@ vacuum_one_database(ConnParams *cparams,
 	 * Instead, let the server decide whether a given relation can be
 	 * processed in which case the user will know about it.
 	 */
-	if (!tables_listed)
+	if ((objfilter & OBJFILTER_TABLE) == 0)
 	{
 		appendPQExpBufferStr(&catalog_query, " WHERE c.relkind OPERATOR(pg_catalog.=) ANY (array["
 							 CppAsString2(RELKIND_RELATION) ", "
@@ -633,7 +719,7 @@ vacuum_one_database(ConnParams *cparams,
 							 fmtQualifiedId(PQgetvalue(res, i, 1),
 											PQgetvalue(res, i, 0)));
 
-		if (tables_listed && !PQgetisnull(res, i, 2))
+		if (objects_listed && !PQgetisnull(res, i, 2))
 			appendPQExpBufferStr(&buf, PQgetvalue(res, i, 2));
 
 		simple_string_list_append(&dbtables, buf.data);
@@ -977,6 +1063,8 @@ help(const char *progname)
 	printf(_("      --no-index-cleanup          don't remove index entries that point to dead tuples\n"));
 	printf(_("      --no-process-toast          skip the TOAST table associated with the table to vacuum\n"));
 	printf(_("      --no-truncate               don't truncate empty pages at the end of the table\n"));
+	printf(_("  -n, --schema=PATTERN            vacuum tables in the specified schema(s) only\n"));
+	printf(_("  -N, --exclude-schema=PATTERN    do not vacuum tables in the specified schema(s)\n"));
 	printf(_("  -P, --parallel=PARALLEL_WORKERS use this many background workers for vacuum, if available\n"));
 	printf(_("  -q, --quiet                     don't write any messages\n"));
 	printf(_("      --skip-locked               skip relations that cannot be immediately locked\n"));
-- 
2.25.1

