From 4ee95216fed740ba45237a308a4c1c9c383fba68 Mon Sep 17 00:00:00 2001
From: John Hsu <47648066+JohnHVancouver@users.noreply.github.com>
Date: Tue, 25 Aug 2020 20:01:49 +0000
Subject: [PATCH] pg_dump - Iterate through all publication tables in
 pg_publication_rel at once

Avoid querying on a per table basis so that we don't need to check if every
table has a publication
---
 src/bin/pg_dump/common.c  |  2 +-
 src/bin/pg_dump/pg_dump.c | 93 +++++++++++++++------------------------
 src/bin/pg_dump/pg_dump.h |  3 +-
 3 files changed, 37 insertions(+), 61 deletions(-)

diff --git a/src/bin/pg_dump/common.c b/src/bin/pg_dump/common.c
index 08239dde4f..63b12c7218 100644
--- a/src/bin/pg_dump/common.c
+++ b/src/bin/pg_dump/common.c
@@ -250,7 +250,7 @@ getSchemaData(Archive *fout, int *numTablesPtr)
 	getPublications(fout);
 
 	pg_log_info("reading publication membership");
-	getPublicationTables(fout, tblinfo, numTables);
+	getPublicationTables(fout);
 
 	pg_log_info("reading subscriptions");
 	getSubscriptions(fout);
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 2cb3f9b083..9272e5e341 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -4044,7 +4044,7 @@ dumpPublication(Archive *fout, PublicationInfo *pubinfo)
  *	  get information about publication membership for dumpable tables.
  */
 void
-getPublicationTables(Archive *fout, TableInfo tblinfo[], int numTables)
+getPublicationTables(Archive *fout)
 {
 	PQExpBuffer query;
 	PGresult   *res;
@@ -4053,8 +4053,8 @@ getPublicationTables(Archive *fout, TableInfo tblinfo[], int numTables)
 	int			i_tableoid;
 	int			i_oid;
 	int			i_pubname;
+	int			i_relid;
 	int			i,
-				j,
 				ntups;
 
 	if (dopt->no_publications || fout->remoteVersion < 100000)
@@ -4062,73 +4062,50 @@ getPublicationTables(Archive *fout, TableInfo tblinfo[], int numTables)
 
 	query = createPQExpBuffer();
 
-	for (i = 0; i < numTables; i++)
-	{
-		TableInfo  *tbinfo = &tblinfo[i];
-
-		/*
-		 * Only regular and partitioned tables can be added to publications.
-		 */
-		if (tbinfo->relkind != RELKIND_RELATION &&
-			tbinfo->relkind != RELKIND_PARTITIONED_TABLE)
-			continue;
+	appendPQExpBuffer(query,
+						"SELECT pr.tableoid, pr.oid, pr.prrelid, p.pubname "
+						"FROM pg_publication_rel pr, pg_publication p "
+						"WHERE p.oid = pr.prpubid");
 
-		/*
-		 * Ignore publication membership of tables whose definitions are not
-		 * to be dumped.
-		 */
-		if (!(tbinfo->dobj.dump & DUMP_COMPONENT_DEFINITION))
-			continue;
+	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
 
-		pg_log_info("reading publication membership for table \"%s.%s\"",
-					tbinfo->dobj.namespace->dobj.name,
-					tbinfo->dobj.name);
+	ntups = PQntuples(res);
 
-		resetPQExpBuffer(query);
+	pubrinfo = pg_malloc(ntups * sizeof(PublicationRelInfo));
+	for (i = 0; i < ntups; i++)
+	{
+		TableInfo *tbInfo;
+		Oid tbRelid;
 
-		/* Get the publication membership for the table. */
-		appendPQExpBuffer(query,
-						  "SELECT pr.tableoid, pr.oid, p.pubname "
-						  "FROM pg_publication_rel pr, pg_publication p "
-						  "WHERE pr.prrelid = '%u'"
-						  "  AND p.oid = pr.prpubid",
-						  tbinfo->dobj.catId.oid);
-		res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+		i_tableoid = PQfnumber(res, "tableoid");
+		i_oid = PQfnumber(res, "oid");
+		i_pubname = PQfnumber(res, "pubname");
+		i_relid = PQfnumber(res, "prrelid");
+		tbRelid = atooid(PQgetvalue(res, i, i_relid));
 
-		ntups = PQntuples(res);
+		tbInfo = findTableByOid(tbRelid);
 
-		if (ntups == 0)
-		{
-			/*
-			 * Table is not member of any publications. Clean up and return.
-			 */
-			PQclear(res);
+		if (tbInfo == NULL)
 			continue;
-		}
 
-		i_tableoid = PQfnumber(res, "tableoid");
-		i_oid = PQfnumber(res, "oid");
-		i_pubname = PQfnumber(res, "pubname");
+		if (!(tbInfo->dobj.dump & DUMP_COMPONENT_DEFINITION))
+			continue;
 
-		pubrinfo = pg_malloc(ntups * sizeof(PublicationRelInfo));
+		pubrinfo[i].dobj.objType = DO_PUBLICATION_REL;
+		pubrinfo[i].dobj.catId.tableoid =
+			atooid(PQgetvalue(res, i, i_tableoid));
+		pubrinfo[i].dobj.catId.oid = atooid(PQgetvalue(res, i, i_oid));
+		AssignDumpId(&pubrinfo[i].dobj);
+		pubrinfo[i].dobj.namespace = tbInfo->dobj.namespace;
+		pubrinfo[i].dobj.name = tbInfo->dobj.name;
+		pubrinfo[i].pubname = pg_strdup(PQgetvalue(res, i, i_pubname));
+		pubrinfo[i].pubtable = tbInfo;
 
-		for (j = 0; j < ntups; j++)
-		{
-			pubrinfo[j].dobj.objType = DO_PUBLICATION_REL;
-			pubrinfo[j].dobj.catId.tableoid =
-				atooid(PQgetvalue(res, j, i_tableoid));
-			pubrinfo[j].dobj.catId.oid = atooid(PQgetvalue(res, j, i_oid));
-			AssignDumpId(&pubrinfo[j].dobj);
-			pubrinfo[j].dobj.namespace = tbinfo->dobj.namespace;
-			pubrinfo[j].dobj.name = tbinfo->dobj.name;
-			pubrinfo[j].pubname = pg_strdup(PQgetvalue(res, j, i_pubname));
-			pubrinfo[j].pubtable = tbinfo;
-
-			/* Decide whether we want to dump it */
-			selectDumpablePublicationTable(&(pubrinfo[j].dobj), fout);
-		}
-		PQclear(res);
+		/* Decide whether we want to dump it */
+		selectDumpablePublicationTable(&(pubrinfo[i].dobj), fout);
 	}
+
+	PQclear(res);
 	destroyPQExpBuffer(query);
 }
 
diff --git a/src/bin/pg_dump/pg_dump.h b/src/bin/pg_dump/pg_dump.h
index da97b731b1..1b43cabea0 100644
--- a/src/bin/pg_dump/pg_dump.h
+++ b/src/bin/pg_dump/pg_dump.h
@@ -716,8 +716,7 @@ extern void processExtensionTables(Archive *fout, ExtensionInfo extinfo[],
 extern EventTriggerInfo *getEventTriggers(Archive *fout, int *numEventTriggers);
 extern void getPolicies(Archive *fout, TableInfo tblinfo[], int numTables);
 extern void getPublications(Archive *fout);
-extern void getPublicationTables(Archive *fout, TableInfo tblinfo[],
-								 int numTables);
+extern void getPublicationTables(Archive *fout);
 extern void getSubscriptions(Archive *fout);
 
 #endif							/* PG_DUMP_H */
-- 
2.17.2

