From 798f970e22034d33146265ed98922c605d0dc237 Mon Sep 17 00:00:00 2001
From: David Zhang <david.zhang@highgo.ca>
Date: Tue, 24 Nov 2020 15:14:42 -0800
Subject: [PATCH] add table access method option to pgbench

---
 src/bin/pgbench/pgbench.c | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 08a5947a9e..24bc6bdbe3 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -188,6 +188,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 using specified access method\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"
@@ -3750,12 +3756,14 @@ initCreateTables(PGconn *con)
 	for (i = 0; i < lengthof(DDLs); i++)
 	{
 		char		opts[256];
+		char		opam[256];
 		char		buffer[256];
 		const struct ddlinfo *ddl = &DDLs[i];
 		const char *cols;
 
 		/* Construct new create table statement. */
 		opts[0] = '\0';
+		opam[0] = '\0';
 
 		/* Partition pgbench_accounts table */
 		if (partition_method != PART_NONE && strcmp(ddl->table, "pgbench_accounts") == 0)
@@ -3776,11 +3784,22 @@ initCreateTables(PGconn *con)
 			PQfreemem(escape_tablespace);
 		}
 
+		if (tableam != NULL)
+		{
+			char	   *escape_tableam;
+
+			escape_tableam = PQescapeIdentifier(con, tableam,
+												   strlen(tableam));
+			snprintf(opam + strlen(opam), sizeof(opam) - strlen(opam),
+					 " using %s", escape_tableam);
+			PQfreemem(escape_tableam);
+		}
+
 		cols = (scale >= SCALE_32BIT_THRESHOLD) ? ddl->bigcols : ddl->smcols;
 
-		snprintf(buffer, sizeof(buffer), "create%s table %s(%s)%s",
+		snprintf(buffer, sizeof(buffer), "create%s table %s(%s)%s%s",
 				 unlogged_tables ? " unlogged" : "",
-				 ddl->table, cols, opts);
+				 ddl->table, cols, opam, opts);
 
 		executeStatement(con, buffer);
 	}
@@ -5422,6 +5441,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}
 	};
 
@@ -5795,6 +5815,10 @@ main(int argc, char **argv)
 					exit(1);
 				}
 				break;
+			case 13:			/* table-am */
+				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