From 8ef006b34462a51cf6c52bc1059eb7b38e26147d Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Fri, 28 Jun 2024 21:31:57 -0500
Subject: [PATCH v3 6/7] use new pg_upgrade async API to parallelize reporting
 extension updates

---
 src/bin/pg_upgrade/version.c | 82 ++++++++++++++++++------------------
 1 file changed, 41 insertions(+), 41 deletions(-)

diff --git a/src/bin/pg_upgrade/version.c b/src/bin/pg_upgrade/version.c
index 2de6dffccd..12783bb2ba 100644
--- a/src/bin/pg_upgrade/version.c
+++ b/src/bin/pg_upgrade/version.c
@@ -139,6 +139,42 @@ old_9_6_invalidate_hash_indexes(ClusterInfo *cluster, bool check_mode)
 		check_ok();
 }
 
+static char *
+report_extension_updates_query(DbInfo *dbinfo, void *arg)
+{
+	return pg_strdup("SELECT name "
+					 "FROM pg_available_extensions "
+					 "WHERE installed_version != default_version");
+}
+
+static void
+report_extension_updates_result(DbInfo *dbinfo, PGresult *res, void *arg)
+{
+	bool		db_used = false;
+	int			ntups = PQntuples(res);
+	int			i_name = PQfnumber(res, "name");
+	char	   *output_path = "update_extensions.sql";
+	FILE	  **script = (FILE **) arg;
+
+	for (int rowno = 0; rowno < ntups; rowno++)
+	{
+		if (*script == NULL && (*script = fopen_priv(output_path, "w")) == NULL)
+			pg_fatal("could not open file \"%s\": %m", output_path);
+		if (!db_used)
+		{
+			PQExpBufferData connectbuf;
+
+			initPQExpBuffer(&connectbuf);
+			appendPsqlMetaConnect(&connectbuf, dbinfo->db_name);
+			fputs(connectbuf.data, *script);
+			termPQExpBuffer(&connectbuf);
+			db_used = true;
+		}
+		fprintf(*script, "ALTER EXTENSION %s UPDATE;\n",
+				quote_identifier(PQgetvalue(res, rowno, i_name)));
+	}
+}
+
 /*
  * report_extension_updates()
  *	Report extensions that should be updated.
@@ -146,53 +182,17 @@ old_9_6_invalidate_hash_indexes(ClusterInfo *cluster, bool check_mode)
 void
 report_extension_updates(ClusterInfo *cluster)
 {
-	int			dbnum;
 	FILE	   *script = NULL;
 	char	   *output_path = "update_extensions.sql";
+	AsyncTask  *task = async_task_create();
 
 	prep_status("Checking for extension updates");
 
-	for (dbnum = 0; dbnum < cluster->dbarr.ndbs; dbnum++)
-	{
-		PGresult   *res;
-		bool		db_used = false;
-		int			ntups;
-		int			rowno;
-		int			i_name;
-		DbInfo	   *active_db = &cluster->dbarr.dbs[dbnum];
-		PGconn	   *conn = connectToServer(cluster, active_db->db_name);
-
-		/* find extensions needing updates */
-		res = executeQueryOrDie(conn,
-								"SELECT name "
-								"FROM pg_available_extensions "
-								"WHERE installed_version != default_version"
-			);
-
-		ntups = PQntuples(res);
-		i_name = PQfnumber(res, "name");
-		for (rowno = 0; rowno < ntups; rowno++)
-		{
-			if (script == NULL && (script = fopen_priv(output_path, "w")) == NULL)
-				pg_fatal("could not open file \"%s\": %m", output_path);
-			if (!db_used)
-			{
-				PQExpBufferData connectbuf;
-
-				initPQExpBuffer(&connectbuf);
-				appendPsqlMetaConnect(&connectbuf, active_db->db_name);
-				fputs(connectbuf.data, script);
-				termPQExpBuffer(&connectbuf);
-				db_used = true;
-			}
-			fprintf(script, "ALTER EXTENSION %s UPDATE;\n",
-					quote_identifier(PQgetvalue(res, rowno, i_name)));
-		}
+	async_task_add_step(task, report_extension_updates_query,
+						report_extension_updates_result, true, &script);
 
-		PQclear(res);
-
-		PQfinish(conn);
-	}
+	async_task_run(task, cluster);
+	async_task_free(task);
 
 	if (script)
 	{
-- 
2.39.3 (Apple Git-146)