From b4721f4c91017297cfbdeaa458291f7a97b023d7 Mon Sep 17 00:00:00 2001
From: Mahendra Singh Thalor <mahi6run@gmail.com>
Date: Tue, 15 Apr 2025 23:45:44 +0530
Subject: [PATCH] pg_restore: refactor code of dump file extenion

After this refactor, we will find file extension once but earlier we were trying
to get extension for each database in loop.
---
 src/bin/pg_dump/pg_restore.c | 82 +++++++++++++++++++++++++++---------
 1 file changed, 61 insertions(+), 21 deletions(-)

diff --git a/src/bin/pg_dump/pg_restore.c b/src/bin/pg_dump/pg_restore.c
index ff4bb320fc9..e6486957620 100644
--- a/src/bin/pg_dump/pg_restore.c
+++ b/src/bin/pg_dump/pg_restore.c
@@ -71,6 +71,8 @@ static int	get_dbnames_list_to_restore(PGconn *conn,
 										SimpleStringList db_exclude_patterns);
 static int	get_dbname_oid_list_from_mfile(const char *dumpdirpath,
 										   SimpleOidStringList *dbname_oid_list);
+static
+char *get_dump_file_exten(const char *dumpdirpath, Oid dboid, const ArchiveFormat format);
 
 int
 main(int argc, char **argv)
@@ -1109,6 +1111,7 @@ restore_all_databases(PGconn *conn, const char *dumpdirpath,
 	bool		dumpData = opts->dumpData;
 	bool		dumpSchema = opts->dumpSchema;
 	bool		dumpStatistics = opts->dumpSchema;
+	const char *file_exten;
 
 	/* Save db name to reuse it for all the database. */
 	if (opts->cparams.dbname)
@@ -1163,6 +1166,9 @@ restore_all_databases(PGconn *conn, const char *dumpdirpath,
 
 	pg_log_info("need to restore %d databases out of %d databases", num_db_restore, num_total_db);
 
+	/* Now get dump file extention. */
+	file_exten = get_dump_file_exten(dumpdirpath, dbname_oid_list.head->oid, opts->format);
+
 	/*
 	 * Till now, we made a list of databases, those needs to be restored after
 	 * skipping names of exclude-database.  Now we can launch parallel workers
@@ -1172,8 +1178,6 @@ restore_all_databases(PGconn *conn, const char *dumpdirpath,
 		 db_cell; db_cell = db_cell->next)
 	{
 		char		subdirpath[MAXPGPATH];
-		char		subdirdbpath[MAXPGPATH];
-		char		dbfilename[MAXPGPATH];
 		int			n_errors;
 
 		/* ignore dbs marked for skipping */
@@ -1190,25 +1194,8 @@ restore_all_databases(PGconn *conn, const char *dumpdirpath,
 			opts->cparams.override_dbname = NULL;
 		}
 
-		snprintf(subdirdbpath, MAXPGPATH, "%s/databases", dumpdirpath);
-
-		/*
-		 * Look for the database dump file/dir. If there is an {oid}.tar or
-		 * {oid}.dmp file, use it. Otherwise try to use a directory called
-		 * {oid}
-		 */
-		snprintf(dbfilename, MAXPGPATH, "%u.tar", db_cell->oid);
-		if (file_exists_in_directory(subdirdbpath, dbfilename))
-			snprintf(subdirpath, MAXPGPATH, "%s/databases/%u.tar", dumpdirpath, db_cell->oid);
-		else
-		{
-			snprintf(dbfilename, MAXPGPATH, "%u.dmp", db_cell->oid);
-
-			if (file_exists_in_directory(subdirdbpath, dbfilename))
-				snprintf(subdirpath, MAXPGPATH, "%s/databases/%u.dmp", dumpdirpath, db_cell->oid);
-			else
-				snprintf(subdirpath, MAXPGPATH, "%s/databases/%u", dumpdirpath, db_cell->oid);
-		}
+		/* Set particular dump file path. */
+		snprintf(subdirpath, MAXPGPATH, "%s/databases/%u%s", dumpdirpath, db_cell->oid, file_exten);
 
 		pg_log_info("restoring database \"%s\"", db_cell->str);
 
@@ -1384,3 +1371,56 @@ copy_or_print_global_file(const char *outfile, FILE *pfile)
 	if (strcmp(outfile, "-") != 0)
 		fclose(OPF);
 }
+
+/*
+ * get_dump_file_exten
+ *
+ * Look for the database dump file/dir. If there is an {oid}.tar or
+ * {oid}.dmp file, use it. Otherwise try to use a directory called
+ * {oid}
+ */
+static
+char *get_dump_file_exten(const char *dumpdirpath, Oid dboid, const ArchiveFormat format)
+{
+	char	*file_exten = "";
+	char	subdirdbpath[MAXPGPATH];
+	char	dbfilename[MAXPGPATH];
+
+	snprintf(subdirdbpath, MAXPGPATH, "%s/databases", dumpdirpath);
+
+	switch(format)
+	{
+		case archCustom:
+			file_exten = ".dmp";
+			break;
+		case archTar:
+			file_exten = ".tar";
+			break;
+		case archDirectory:
+			file_exten = "";
+			break;
+		case archUnknown: /* based on file exist, try to get file extension name. */
+			snprintf(dbfilename, MAXPGPATH, "%u", dboid);
+			if (file_exists_in_directory(subdirdbpath, dbfilename))
+				file_exten = "";
+			else
+			{
+				snprintf(dbfilename, MAXPGPATH, "%u.dmp", dboid);
+				if (file_exists_in_directory(subdirdbpath, dbfilename))
+					file_exten = ".dmp";
+				else
+				{
+					snprintf(dbfilename, MAXPGPATH, "%u.tar", dboid);
+					if (file_exists_in_directory(subdirdbpath, dbfilename))
+						file_exten = ".tar";
+					else
+						file_exten = "";
+				}
+			}
+			break;
+		default:
+			pg_fatal("unrecognized file format \"%d\"", format);
+	}
+
+	return file_exten;
+}
-- 
2.39.3