From c64173842a163abbdd6210af5f9643b3301d8de0 Mon Sep 17 00:00:00 2001
From: David Zhang <david.zhang@highgo.ca>
Date: Thu, 26 Nov 2020 18:08:49 -0800
Subject: [PATCH] add table access method as an option to pgbench

---
 doc/src/sgml/ref/pgbench.sgml                |  9 ++++++
 src/bin/pgbench/pgbench.c                    | 23 ++++++++++++++-
 src/bin/pgbench/t/001_pgbench_with_server.pl | 30 +++++++++++++++++++-
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 7180fedd65..3af5230f85 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>TABLEACCESSMETHOD</replaceable></option></term>
+      <listitem>
+       <para>
+        Create tables using the 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 3057665bbe..c3939b6bb4 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=TABLEACCESSMETHOD\n"
+		   "                           create tables with using 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"
@@ -3759,6 +3763,18 @@ initCreateTables(PGconn *con)
 						  ddl->table,
 						  (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols);
 
+		if (partition_method == PART_NONE && table_access_method != NULL)
+		{
+			char	   *escaped_table_access_method;
+
+			escaped_table_access_method =
+					PQescapeIdentifier(con,
+					table_access_method, strlen(table_access_method));
+			appendPQExpBuffer(&query, " using %s",
+					escaped_table_access_method);
+			PQfreemem(escaped_table_access_method);
+		}
+
 		/* Partition pgbench_accounts table */
 		if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0)
 			appendPQExpBuffer(&query,
@@ -5419,6 +5435,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}
 	};
 
@@ -5792,6 +5809,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..4b780fd0e0 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -114,7 +114,7 @@ pgbench(
 
 # Again, with all possible options
 pgbench(
-	'--initialize --init-steps=dtpvg --scale=1 --unlogged-tables --fillfactor=98 --foreign-keys --quiet --tablespace=regress_pgbench_tap_1_ts --index-tablespace=regress_pgbench_tap_1_ts --partitions=2 --partition-method=hash',
+	'--initialize --init-steps=dtpvg --scale=1 --unlogged-tables --fillfactor=98 --foreign-keys --quiet --tablespace=regress_pgbench_tap_1_ts --index-tablespace=regress_pgbench_tap_1_ts --partitions=2 --partition-method=hash --table-access-method=heap',
 	0,
 	[qr{^$}i],
 	[
@@ -129,6 +129,22 @@ pgbench(
 	],
 	'pgbench scale 1 initialization');
 
+# Again, with all possible options except partition
+pgbench(
+	'--initialize --init-steps=dtpvg --scale=1 --unlogged-tables --fillfactor=98 --foreign-keys --quiet --tablespace=regress_pgbench_tap_1_ts --index-tablespace=regress_pgbench_tap_1_ts --table-access-method=heap',
+	0,
+	[qr{^$}i],
+	[
+		qr{dropping old tables},
+		qr{creating tables},
+		qr{vacuuming},
+		qr{creating primary keys},
+		qr{creating foreign keys},
+		qr{(?!vacuuming)},    # no vacuum
+		qr{done in \d+\.\d\d s }
+	],
+	'pgbench scale 1 initialization');
+
 # Test interaction of --init-steps with legacy step-selection options
 pgbench(
 	'--initialize --init-steps=dtpvGvv --no-vacuum --foreign-keys --unlogged-tables --partitions=3',
@@ -146,6 +162,18 @@ pgbench(
 	],
 	'pgbench --init-steps');
 
+# Test interaction of --table-access
+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');
+
 # Run all builtin scripts, for a few transactions each
 pgbench(
 	'--transactions=5 -Dfoo=bla --client=2 --protocol=simple --builtin=t'
-- 
2.17.1

