v25_0004-update-AX-handle-for-each-database-for-cleanup.patch

From 8450086079480f3784d26c6c85973b84cfa8b484 Mon Sep 17 00:00:00 2001
From: Mahendra Singh Thalor <mahi6run@gmail.com>
Date: Tue, 1 Apr 2025 11:08:38 +0530
Subject: [PATCH 4/4] update AX handle for each database for cleanup.

---
 src/bin/pg_dump/parallel.c           | 10 ++++++++++
 src/bin/pg_dump/pg_backup_archiver.h |  1 +
 src/bin/pg_dump/pg_restore.c         | 10 +++++-----
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/bin/pg_dump/parallel.c b/src/bin/pg_dump/parallel.c
index 086adcdc502..5974d6706fd 100644
--- a/src/bin/pg_dump/parallel.c
+++ b/src/bin/pg_dump/parallel.c
@@ -333,6 +333,16 @@ on_exit_close_archive(Archive *AHX)
 	on_exit_nicely(archive_close_connection, &shutdown_info);
 }
 
+/*
+ * When pg_restore restores multiple databases, then update already added entry
+ * into array for cleanup.
+ */
+void
+replace_on_exit_close_archive(Archive *AHX)
+{
+	shutdown_info.AHX = AHX;
+}
+
 /*
  * on_exit_nicely handler for shutting down database connections and
  * worker processes cleanly.
diff --git a/src/bin/pg_dump/pg_backup_archiver.h b/src/bin/pg_dump/pg_backup_archiver.h
index a2064f471ed..ed0238cca47 100644
--- a/src/bin/pg_dump/pg_backup_archiver.h
+++ b/src/bin/pg_dump/pg_backup_archiver.h
@@ -386,6 +386,7 @@ struct _tocEntry
 
 extern int	parallel_restore(ArchiveHandle *AH, TocEntry *te);
 extern void on_exit_close_archive(Archive *AHX);
+extern void replace_on_exit_close_archive(Archive *AHX);
 
 extern void warn_or_exit_horribly(ArchiveHandle *AH, const char *fmt,...) pg_attribute_printf(2, 3);
 
diff --git a/src/bin/pg_dump/pg_restore.c b/src/bin/pg_dump/pg_restore.c
index aa8887a4eb0..a142b744222 100644
--- a/src/bin/pg_dump/pg_restore.c
+++ b/src/bin/pg_dump/pg_restore.c
@@ -603,14 +603,14 @@ restore_one_database(const char *inputFileSpec, RestoreOptions *opts,
 	 * We don't have a connection yet but that doesn't matter. The connection
 	 * is initialized to NULL and if we terminate through exit_nicely() while
 	 * it's still NULL, the cleanup function will just be a no-op. If we are
-	 * restoring multiple databases, then save index of exit_nicely so that we
-	 * can use same slot for all the databases as we already closed the
-	 * previous archive by CloseArchive.
+	 * restoring multiple databases, then only update AX handle for cleanup as
+	 * previous entry was already in array and we had closed previous
+	 * connection so we can use same slot from array.
 	 */
 	if (!append_data || num == 0)
 		on_exit_close_archive(AH);
-	//else
-		//replace_on_exit_close_archive(AH);
+	else
+		replace_on_exit_close_archive(AH);
 
 	/* Let the archiver know how noisy to be */
 	AH->verbose = opts->verbose;
-- 
2.39.3