From 7ad5a2aabab2393406f632e2559474af8af1382c Mon Sep 17 00:00:00 2001
From: Shenhao Wang <wangsh.fnst@cn.fujitsu.com>
Date: Fri, 18 Dec 2020 11:36:08 +0800
Subject: [PATCH] Improve pgbench when -f option contain char '@' in filepath

---
 src/bin/pgbench/pgbench.c                    | 43 +++++++++++++++-----
 src/bin/pgbench/t/001_pgbench_with_server.pl |  4 +-
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 3057665bbe..ceebb38409 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -4925,24 +4925,37 @@ read_file_contents(FILE *fd)
 }
 
 /*
- * Given a file name, read it and add its script to the list.
+ * Given a file name, and open it
  * "-" means to read stdin.
  * NB: filename must be storage that won't disappear.
  */
-static void
-process_file(const char *filename, int weight)
+static FILE*
+open_file(const char * filename, bool error_exit)
 {
 	FILE	   *fd;
-	char	   *buf;
-
-	/* Slurp the file contents into "buf" */
 	if (strcmp(filename, "-") == 0)
 		fd = stdin;
-	else if ((fd = fopen(filename, "r")) == NULL)
+	else if ((fd = fopen(filename, "r")) == NULL && error_exit)
 	{
 		pg_log_fatal("could not open file \"%s\": %m", filename);
 		exit(1);
 	}
+	return fd;
+}
+
+/*
+ * Given a FILE ptr, read it and add its script to the list.
+ * if FILE ptr is null, open file with filename
+ * NB: filename must be storage that won't disappear.
+ */
+static void
+process_file(FILE *fd, int weight, const char *filename)
+{
+	char	   *buf;
+
+	/* Slurp the file contents into "buf" */
+	if (fd == NULL)
+		fd = open_file(filename, true);
 
 	buf = read_file_contents(fd);
 
@@ -5633,10 +5646,18 @@ main(int argc, char **argv)
 				internal_script_used = true;
 				break;
 			case 'f':
-				weight = parseScriptWeight(optarg, &script);
-				process_file(script, weight);
-				benchmarking_option_set = true;
-				break;
+				{
+					FILE *fd = open_file(optarg, false);
+					if (fd != NULL)
+						process_file(fd, 1, optarg);
+					else
+					{
+						weight = parseScriptWeight(optarg, &script);
+						process_file(NULL, weight, script);
+					}
+					benchmarking_option_set = true;
+					break;
+				}
 			case 'D':
 				{
 					char	   *p;
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index 61b671d54f..ecb8024ef0 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -233,14 +233,14 @@ pgbench(
 	'-n -t 10 -c 1 -M simple',
 	0,
 	[
-		qr{type: .*/001_pgbench_custom_script_3},
+		qr{type: .*/001_pgbench_\@custom_script_3},
 		qr{processed: 10/10},
 		qr{mode: simple}
 	],
 	[qr{^$}],
 	'pgbench custom script',
 	{
-		'001_pgbench_custom_script_3' => q{-- select only variant
+		'001_pgbench_@custom_script_3' => q{-- select only variant
 \set aid random(1, :scale * 100000)
 BEGIN;
 SELECT abalance::INTEGER AS balance
-- 
2.26.2