From 1e908ecdd6add81c96dca489f45f5ae9d0c2a5f1 Mon Sep 17 00:00:00 2001
From: David Zhang <david.zhang@highgo.ca>
Date: Wed, 25 Nov 2020 10:41:20 -0800
Subject: [PATCH] add table access method option to pgbench

---
 doc/src/sgml/ref/pgbench.sgml | 11 +++++++++++
 src/bin/pgbench/pgbench.c     | 20 ++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 7180fedd65..a8f8d24fe1 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -378,6 +378,17 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+      <term><option>--table-am=<replaceable>TABLEAM</replaceable></option></term>
+      <listitem>
+       <para>
+        Create all tables with specified table access method 
+        <replaceable>TABLEAM</replaceable>.
+        If unspecified, default is <literal>heap</literal>.
+       </para>
+      </listitem>
+     </varlistentry>
+
     </variablelist>
    </para>
 
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 3057665bbe..62b1dd3e23 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -189,6 +189,11 @@ int64		latency_limit = 0;
 char	   *tablespace = NULL;
 char	   *index_tablespace = NULL;
 
+/*
+ * table access method selection
+ */
+char	   *tableam = NULL;
+
 /*
  * Number of "pgbench_accounts" partitions.  0 is the default and means no
  * partitioning.
@@ -643,6 +648,7 @@ usage(void)
 		   "  --partitions=NUM         partition pgbench_accounts into NUM parts (default: 0)\n"
 		   "  --tablespace=TABLESPACE  create tables in the specified tablespace\n"
 		   "  --unlogged-tables        create tables as unlogged tables\n"
+		   "  --table-am=TABLEAM       create tables with specified table access method (default: heap)\n"
 		   "\nOptions to select what to run:\n"
 		   "  -b, --builtin=NAME[@W]   add builtin script NAME weighted at W (default: 1)\n"
 		   "                           (use \"-b list\" to list available scripts)\n"
@@ -3759,6 +3765,15 @@ initCreateTables(PGconn *con)
 						  ddl->table,
 						  (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols);
 
+		if (tableam != NULL)
+		{
+			char	   *escape_tableam;
+
+			escape_tableam = PQescapeIdentifier(con, tableam, strlen(tableam));
+			appendPQExpBuffer(&query, " using %s", escape_tableam);
+			PQfreemem(escape_tableam);
+		}
+
 		/* Partition pgbench_accounts table */
 		if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0)
 			appendPQExpBuffer(&query,
@@ -5419,6 +5434,7 @@ main(int argc, char **argv)
 		{"show-script", required_argument, NULL, 10},
 		{"partitions", required_argument, NULL, 11},
 		{"partition-method", required_argument, NULL, 12},
+		{"table-am", required_argument, NULL, 13},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -5792,6 +5808,10 @@ main(int argc, char **argv)
 					exit(1);
 				}
 				break;
+			case 13:			/* table access method*/
+				initialization_option_set = true;
+				tableam = pg_strdup(optarg);
+				break;
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
-- 
2.17.1