From 517fcb61c7065b7bccf6105a7bb89007fe2fcb08 Mon Sep 17 00:00:00 2001
From: David Zhang <david.zhang@highgo.ca>
Date: Fri, 15 Jan 2021 10:32:53 -0800
Subject: [PATCH] add table access method as an option to pgbench

---
 doc/src/sgml/ref/pgbench.sgml                |  9 +++++++
 src/bin/pgbench/pgbench.c                    | 25 +++++++++++++++++++-
 src/bin/pgbench/t/001_pgbench_with_server.pl | 22 +++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index faa7c26b0a..d8018388e6 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -359,6 +359,15 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+      <term><option>--table-access-method=<replaceable>TABLEAM</replaceable></option></term>
+      <listitem>
+       <para>
+        Create tables with specified table access method, rather than the default.
+       </para>
+      </listitem>
+     </varlistentry>
+     
      <varlistentry>
       <term><option>--tablespace=<replaceable>tablespace</replaceable></option></term>
       <listitem>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index f7da3e1f62..3f8a119811 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -184,10 +184,11 @@ double		throttle_delay = 0;
 int64		latency_limit = 0;
 
 /*
- * tablespace selection
+ * tablespace and table access method selection
  */
 char	   *tablespace = NULL;
 char	   *index_tablespace = NULL;
+char	   *table_access_method = NULL;
 
 /*
  * Number of "pgbench_accounts" partitions.  0 is the default and means no
@@ -641,6 +642,9 @@ usage(void)
 		   "  --partition-method=(range|hash)\n"
 		   "                           partition pgbench_accounts with this method (default: range)\n"
 		   "  --partitions=NUM         partition pgbench_accounts into NUM parts (default: 0)\n"
+		   "  --table-access-method=TABLEAM\n"
+		   "                           create tables with specified table access method,\n"
+		   "                           rather than the default.\n"
 		   "  --tablespace=TABLESPACE  create tables in the specified tablespace\n"
 		   "  --unlogged-tables        create tables as unlogged tables\n"
 		   "\nOptions to select what to run:\n"
@@ -3761,6 +3765,20 @@ initCreateTables(PGconn *con)
 
 	fprintf(stderr, "creating tables...\n");
 
+	if (table_access_method != NULL)
+	{
+		char	   *escaped_table_access_method;
+
+		initPQExpBuffer(&query);
+		escaped_table_access_method = PQescapeIdentifier(con,
+				table_access_method, strlen(table_access_method));
+		appendPQExpBuffer(&query, "set default_table_access_method to %s;\n",
+				escaped_table_access_method);
+		PQfreemem(escaped_table_access_method);
+		executeStatement(con, query.data);
+		termPQExpBuffer(&query);
+	}
+
 	initPQExpBuffer(&query);
 
 	for (i = 0; i < lengthof(DDLs); i++)
@@ -5433,6 +5451,7 @@ main(int argc, char **argv)
 		{"show-script", required_argument, NULL, 10},
 		{"partitions", required_argument, NULL, 11},
 		{"partition-method", required_argument, NULL, 12},
+		{"table-access-method", required_argument, NULL, 13},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -5806,6 +5825,10 @@ main(int argc, char **argv)
 					exit(1);
 				}
 				break;
+			case 13:			/* table access method*/
+				initialization_option_set = true;
+				table_access_method = pg_strdup(optarg);
+				break;
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index 61b671d54f..077e862530 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -100,6 +100,16 @@ pgbench(
 	[qr{Perhaps you need to do initialization}],
 	'run without init');
 
+pgbench(
+        '-i --table-access-method=no-such-table-am',
+        1,
+        [qr{^$}],
+        [
+                qr{invalid value for parameter "default_table_access_method"},
+                qr{DETAIL:  Table access method "no-such-table-am" does not exist}
+        ],
+        'no such table access method');
+
 # Initialize pgbench tables scale 1
 pgbench(
 	'-i', 0,
@@ -146,6 +156,18 @@ pgbench(
 	],
 	'pgbench --init-steps');
 
+# Test interaction of --table-access-method
+pgbench(
+	'--initialize --table-access-method=heap',
+	0,
+	[qr{^$}],
+	[
+		qr{dropping old tables},
+		qr{creating tables},
+		qr{done in \d+\.\d\d s }
+	],
+	'pgbench --table-access-method');
+
 # Run all builtin scripts, for a few transactions each
 pgbench(
 	'--transactions=5 -Dfoo=bla --client=2 --protocol=simple --builtin=t'
-- 
2.17.1