From edd11eacc2ae4fdb868e831f5fabdf3668640417 Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Tue, 7 Jul 2020 19:14:21 -0400
Subject: [PATCH v2 2/2] pg_dump: lock all rels, not just RELKIND_RELATION

---
 src/bin/pg_dump/pg_backup.h    |  2 ++
 src/bin/pg_dump/pg_backup_db.c | 64 ++++++++++++++++++++++++++++++++++
 src/bin/pg_dump/pg_backup_db.h |  2 ++
 src/bin/pg_dump/pg_dump.c      | 17 +++++----
 4 files changed, 78 insertions(+), 7 deletions(-)

diff --git a/src/bin/pg_dump/pg_backup.h b/src/bin/pg_dump/pg_backup.h
index a6a8e6f2fd..6c79319164 100644
--- a/src/bin/pg_dump/pg_backup.h
+++ b/src/bin/pg_dump/pg_backup.h
@@ -198,6 +198,8 @@ typedef struct Archive
 	int			minRemoteVersion;	/* allowable range */
 	int			maxRemoteVersion;
 
+	bool		hasGenericLockTable;	/* can LOCK TABLE do non-table rels */
+
 	int			numWorkers;		/* number of parallel processes */
 	char	   *sync_snapshot_id;	/* sync snapshot id for parallel operation */
 
diff --git a/src/bin/pg_dump/pg_backup_db.c b/src/bin/pg_dump/pg_backup_db.c
index 5ba43441f5..a9eba97b84 100644
--- a/src/bin/pg_dump/pg_backup_db.c
+++ b/src/bin/pg_dump/pg_backup_db.c
@@ -531,6 +531,70 @@ EndDBCopyMode(Archive *AHX, const char *tocEntryTag)
 	}
 }
 
+/*
+ * Does LOCK TABLE work on non-table relations on this server?
+ *
+ * Note: assumes it is called out of any transaction
+ */
+bool
+IsLockTableGeneric(Archive *AHX)
+{
+	ArchiveHandle *AH = (ArchiveHandle *) AHX;
+	PGresult *res;
+	char	 *sqlstate;
+	bool	retval;
+
+	if (AHX->remoteVersion >= 140000)
+		return true;
+	else if (AHX->remoteVersion < 90500)
+		return false;
+
+	StartTransaction(AHX);
+
+	/*
+	 * Try a LOCK TABLE on a well-known non-table catalog; WRONG_OBJECT_TYPE
+	 * tells us that this server doesn't support locking non-table rels, while
+	 * LOCK_NOT_AVAILABLE and INSUFFICIENT_PRIVILEGE tell us that it does.
+	 * Report anything else as a fatal problem.
+	 */
+#define ERRCODE_INSUFFICIENT_PRIVILEGE	"42501"
+#define ERRCODE_WRONG_OBJECT_TYPE	"42809"
+#define ERRCODE_LOCK_NOT_AVAILABLE	"55P03"
+	res = PQexec(AH->connection,
+				 "LOCK TABLE pg_catalog.pg_aggregate_fnoid_index IN ACCESS SHARE MODE NOWAIT");
+	switch (PQresultStatus(res))
+	{
+		case PGRES_COMMAND_OK:
+			retval = true;
+			break;
+		case PGRES_FATAL_ERROR:
+			sqlstate = PQresultErrorField(res, PG_DIAG_SQLSTATE);
+			if (strcmp(sqlstate, ERRCODE_WRONG_OBJECT_TYPE) == 0)
+			{
+				retval = false;
+				break;
+			}
+			else if (strcmp(sqlstate, ERRCODE_LOCK_NOT_AVAILABLE) == 0 ||
+					 strcmp(sqlstate, ERRCODE_INSUFFICIENT_PRIVILEGE) == 0)
+			{
+				retval = true;
+				break;
+			}
+			/* else, falls through */
+		default:
+			warn_or_exit_horribly(AH, "LOCK TABLE failed for \"%s\": %s",
+								  "pg_catalog.pg_stats",
+								  PQerrorMessage(AH->connection));
+			retval = false;		/* not reached */
+			break;
+	}
+	PQclear(res);
+
+	CommitTransaction(AHX);
+
+	return retval;
+}
+
 void
 StartTransaction(Archive *AHX)
 {
diff --git a/src/bin/pg_dump/pg_backup_db.h b/src/bin/pg_dump/pg_backup_db.h
index 8888dd34b9..355beec1f7 100644
--- a/src/bin/pg_dump/pg_backup_db.h
+++ b/src/bin/pg_dump/pg_backup_db.h
@@ -20,6 +20,8 @@ extern PGresult *ExecuteSqlQueryForSingleRow(Archive *fout, const char *query);
 
 extern void EndDBCopyMode(Archive *AHX, const char *tocEntryTag);
 
+extern bool IsLockTableGeneric(Archive *AHX);
+
 extern void StartTransaction(Archive *AHX);
 extern void CommitTransaction(Archive *AHX);
 
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index ff45e3fb8c..03f9d4d9e8 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -1176,6 +1176,9 @@ setup_connection(Archive *AH, const char *dumpencoding,
 			ExecuteSqlStatement(AH, "SET row_security = off");
 	}
 
+	/* Detect whether LOCK TABLE can handle non-table relations */
+	AH->hasGenericLockTable = IsLockTableGeneric(AH);
+
 	/*
 	 * Start transaction-snapshot mode transaction to dump consistent data.
 	 */
@@ -6843,16 +6846,16 @@ getTables(Archive *fout, int *numTables)
 		 * assume our lock on the child is enough to prevent schema
 		 * alterations to parent tables.
 		 *
-		 * NOTE: it'd be kinda nice to lock other relations too, not only
-		 * plain or partitioned tables, but the backend doesn't presently
-		 * allow that.
-		 *
-		 * We only need to lock the table for certain components; see
+		 * We only need to lock the relation for certain components; see
 		 * pg_dump.h
+		 *
+		 * On server versions that support it, we lock all relations not just
+		 * plain tables.
 		 */
 		if (tblinfo[i].dobj.dump &&
-			(tblinfo[i].relkind == RELKIND_RELATION ||
-			 tblinfo->relkind == RELKIND_PARTITIONED_TABLE) &&
+			(fout->hasGenericLockTable ||
+			 tblinfo[i].relkind == RELKIND_PARTITIONED_TABLE ||
+			 tblinfo[i].relkind == RELKIND_RELATION) &&
 			(tblinfo[i].dobj.dump & DUMP_COMPONENTS_REQUIRING_LOCK))
 		{
 			resetPQExpBuffer(query);
-- 
2.20.1