v4-0002-Always-use-a-catalog-query-to-discover-tables-to-.patch

From 98803895b2422a4e7ef78b922fb195e116b02110 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <bossartn@amazon.com>
Date: Tue, 22 Jan 2019 23:08:01 +0000
Subject: [PATCH v4 2/3] Always use a catalog query to discover tables to
 process in vacuumdb.

---
 src/bin/scripts/common.c   |   2 +-
 src/bin/scripts/common.h   |   3 +
 src/bin/scripts/vacuumdb.c | 149 +++++++++++++++++++++++++++------------------
 3 files changed, 93 insertions(+), 61 deletions(-)

diff --git a/src/bin/scripts/common.c b/src/bin/scripts/common.c
index 4215bc3d6e..493408d15d 100644
--- a/src/bin/scripts/common.c
+++ b/src/bin/scripts/common.c
@@ -265,7 +265,7 @@ executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
  * finish using them, pg_free(*table).  *columns is a pointer into "spec",
  * possibly to its NUL terminator.
  */
-static void
+void
 split_table_columns_spec(const char *spec, int encoding,
 						 char **table, const char **columns)
 {
diff --git a/src/bin/scripts/common.h b/src/bin/scripts/common.h
index 7c3888cefd..f3a7616bae 100644
--- a/src/bin/scripts/common.h
+++ b/src/bin/scripts/common.h
@@ -48,6 +48,9 @@ extern void executeCommand(PGconn *conn, const char *query,
 extern bool executeMaintenanceCommand(PGconn *conn, const char *query,
 						  bool echo);
 
+extern void split_table_columns_spec(const char *spec, int encoding,
+						 char **table, const char **columns);
+
 extern void appendQualifiedRelation(PQExpBuffer buf, const char *name,
 						PGconn *conn, const char *progname, bool echo);
 
diff --git a/src/bin/scripts/vacuumdb.c b/src/bin/scripts/vacuumdb.c
index ec7d0a326a..b72e65c36d 100644
--- a/src/bin/scripts/vacuumdb.c
+++ b/src/bin/scripts/vacuumdb.c
@@ -19,6 +19,7 @@
 #include "catalog/pg_class_d.h"
 
 #include "common.h"
+#include "fe_utils/connect.h"
 #include "fe_utils/simple_list.h"
 #include "fe_utils/string_utils.h"
 
@@ -62,9 +63,7 @@ static void vacuum_all_databases(vacuumingOptions *vacopts,
 					 const char *progname, bool echo, bool quiet);
 
 static void prepare_vacuum_command(PQExpBuffer sql, PGconn *conn,
-					   vacuumingOptions *vacopts, const char *table,
-					   bool table_pre_qualified,
-					   const char *progname, bool echo);
+					   vacuumingOptions *vacopts, const char *table);
 
 static void run_vacuum_command(PGconn *conn, const char *sql, bool echo,
 				   const char *table, const char *progname, bool async);
@@ -359,13 +358,18 @@ vacuum_one_database(const char *dbname, vacuumingOptions *vacopts,
 					const char *progname, bool echo, bool quiet)
 {
 	PQExpBufferData sql;
+	PQExpBufferData buf;
+	PQExpBufferData catalog_query;
+	PGresult   *res;
 	PGconn	   *conn;
 	SimpleStringListCell *cell;
 	ParallelSlot *slots;
 	SimpleStringList dbtables = {NULL, NULL};
 	int			i;
+	int			ntups;
 	bool		failed = false;
 	bool		parallel = concurrentCons > 1;
+	bool		first_table = true;
 	const char *stage_commands[] = {
 		"SET default_statistics_target=1; SET vacuum_cost_delay=0;",
 		"SET default_statistics_target=10; RESET vacuum_cost_delay;",
@@ -410,54 +414,91 @@ vacuum_one_database(const char *dbname, vacuumingOptions *vacopts,
 		fflush(stdout);
 	}
 
-	initPQExpBuffer(&sql);
-
 	/*
-	 * If a table list is not provided and we're using multiple connections,
-	 * prepare the list of tables by querying the catalogs.
+	 * Prepare the list of tables to process by querying the catalogs.
+	 *
+	 * Since we execute the constructed query with the default search_path
+	 * (which could be unsafe), everything in this query must be fully
+	 * qualified.
 	 */
-	if (parallel && (!tables || !tables->head))
+	initPQExpBuffer(&catalog_query);
+	appendPQExpBuffer(&catalog_query,
+					  "SELECT c.relname, ns.nspname"
+					  " FROM pg_catalog.pg_class c\n"
+					  " JOIN pg_catalog.pg_namespace ns"
+					  " ON c.relnamespace OPERATOR(pg_catalog.=) ns.oid\n"
+					  " LEFT JOIN pg_catalog.pg_class t"
+					  " ON c.reltoastrelid OPERATOR(pg_catalog.=) t.oid\n"
+					  " WHERE c.relkind OPERATOR(pg_catalog.=) ANY (array["
+					  CppAsString2(RELKIND_RELATION) ", "
+					  CppAsString2(RELKIND_MATVIEW) "])\n");
+
+	cell = tables ? tables->head : NULL;
+	while (cell != NULL)
 	{
-		PQExpBufferData buf;
-		PGresult   *res;
-		int			ntups;
-
-		initPQExpBuffer(&buf);
-
-		res = executeQuery(conn,
-						   "SELECT c.relname, ns.nspname"
-						   " FROM pg_class c, pg_namespace ns\n"
-						   " WHERE relkind IN ("
-						   CppAsString2(RELKIND_RELATION) ", "
-						   CppAsString2(RELKIND_MATVIEW) ")"
-						   " AND c.relnamespace = ns.oid\n"
-						   " ORDER BY c.relpages DESC;",
-						   progname, echo);
-
-		ntups = PQntuples(res);
-		for (i = 0; i < ntups; i++)
-		{
-			appendPQExpBufferStr(&buf,
-								 fmtQualifiedId(PQgetvalue(res, i, 1),
-												PQgetvalue(res, i, 0)));
+		char	   *just_table;
+		const char *just_columns;
 
-			simple_string_list_append(&dbtables, buf.data);
-			resetPQExpBuffer(&buf);
+		split_table_columns_spec(cell->val, PQclientEncoding(conn),
+								 &just_table, &just_columns);
+
+		if (first_table)
+		{
+			appendPQExpBuffer(&catalog_query, " AND (\n  c.oid OPERATOR(pg_catalog.=) ");
+			first_table = false;
 		}
+		else
+			appendPQExpBuffer(&catalog_query, "  OR c.oid OPERATOR(pg_catalog.=) ");
 
-		termPQExpBuffer(&buf);
-		tables = &dbtables;
+		appendStringLiteralConn(&catalog_query, just_table, conn);
+		appendPQExpBuffer(&catalog_query, "::pg_catalog.regclass\n");
 
-		/*
-		 * If there are more connections than vacuumable relations, we don't
-		 * need to use them all.
-		 */
+		pg_free(just_table);
+
+		cell = cell->next;
+		if (cell == NULL)
+			appendPQExpBuffer(&catalog_query, " )\n");
+	}
+
+	appendPQExpBuffer(&catalog_query, " ORDER BY c.relpages DESC;");
+	executeCommand(conn, "RESET search_path;", progname, echo);
+	res = executeQuery(conn, catalog_query.data, progname, echo);
+	termPQExpBuffer(&catalog_query);
+	PQclear(executeQuery(conn, ALWAYS_SECURE_SEARCH_PATH_SQL,
+						 progname, echo));
+
+	ntups = PQntuples(res);
+	if (ntups == 0)
+	{
+		PQclear(res);
+		PQfinish(conn);
+		return;
+	}
+
+	initPQExpBuffer(&buf);
+	for (i = 0; i < ntups; i++)
+	{
+		appendPQExpBufferStr(&buf,
+							 fmtQualifiedId(PQgetvalue(res, i, 1),
+											PQgetvalue(res, i, 0)));
+
+		simple_string_list_append(&dbtables, buf.data);
+		resetPQExpBuffer(&buf);
+	}
+	termPQExpBuffer(&buf);
+
+	/*
+	 * If there are more connections than vacuumable relations, we don't need
+	 * to use them all.
+	 */
+	if (parallel)
+	{
 		if (concurrentCons > ntups)
 			concurrentCons = ntups;
 		if (concurrentCons <= 1)
 			parallel = false;
-		PQclear(res);
 	}
+	PQclear(res);
 
 	/*
 	 * Setup the database connections. We reuse the connection we already have
@@ -493,10 +534,12 @@ vacuum_one_database(const char *dbname, vacuumingOptions *vacopts,
 						   stage_commands[stage], progname, echo);
 	}
 
-	cell = tables ? tables->head : NULL;
+	initPQExpBuffer(&sql);
+
+	cell = dbtables.head;
 	do
 	{
-		const char *tabname = cell ? cell->val : NULL;
+		const char *tabname = cell->val;
 		ParallelSlot *free_slot;
 
 		if (CancelRequested)
@@ -529,12 +572,7 @@ vacuum_one_database(const char *dbname, vacuumingOptions *vacopts,
 		else
 			free_slot = slots;
 
-		/*
-		 * Prepare the vacuum command.  Note that in some cases this requires
-		 * query execution, so be sure to use the free connection.
-		 */
-		prepare_vacuum_command(&sql, free_slot->connection, vacopts, tabname,
-							   tables == &dbtables, progname, echo);
+		prepare_vacuum_command(&sql, free_slot->connection, vacopts, tabname);
 
 		/*
 		 * Execute the vacuum.  If not in parallel mode, this terminates the
@@ -544,8 +582,7 @@ vacuum_one_database(const char *dbname, vacuumingOptions *vacopts,
 		run_vacuum_command(free_slot->connection, sql.data,
 						   echo, tabname, progname, parallel);
 
-		if (cell)
-			cell = cell->next;
+		cell = cell->next;
 	} while (cell != NULL);
 
 	if (parallel)
@@ -658,9 +695,7 @@ vacuum_all_databases(vacuumingOptions *vacopts,
  */
 static void
 prepare_vacuum_command(PQExpBuffer sql, PGconn *conn,
-					   vacuumingOptions *vacopts, const char *table,
-					   bool table_pre_qualified,
-					   const char *progname, bool echo)
+					   vacuumingOptions *vacopts, const char *table)
 {
 	const char *paren = " (";
 	const char *comma = ", ";
@@ -753,14 +788,8 @@ prepare_vacuum_command(PQExpBuffer sql, PGconn *conn,
 		}
 	}
 
-	if (table)
-	{
-		appendPQExpBufferChar(sql, ' ');
-		if (table_pre_qualified)
-			appendPQExpBufferStr(sql, table);
-		else
-			appendQualifiedRelation(sql, table, conn, progname, echo);
-	}
+	appendPQExpBufferChar(sql, ' ');
+	appendPQExpBufferStr(sql, table);
 	appendPQExpBufferChar(sql, ';');
 }
 
-- 
2.16.5