diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 03e1212..6736d61 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -211,6 +211,65 @@ pgbench options dbname
+
+
+
+
+ Required to invoke initialization with custom steps.
+ custom_init_command specifies custom
+ initialization commands for the initialization. Each command is invoked
+ in the specified order. The supported commands are:
+
+
+
+ t
+
+
+ Create four tables pgbench_accounts,
+ pgbench_branches, pgbench_history,
+ and pgbench_tellers, destroying any existing
+ tables of these names.
+
+
+
+
+ d
+
+
+ Load data to standard tables.
+
+
+
+
+ p
+
+
+ Create primary keys on standard tables.
+
+
+
+
+ f
+
+
+ Create foreign keys constraints between the standard tables.
+
+
+
+
+ v
+
+
+ Invoke vacuum on standard tables.
+
+
+
+
+
+
+
+
+
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index ae78c7b..71aae6b 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -117,6 +117,11 @@ int fillfactor = 100;
int foreign_keys = 0;
/*
+ * no vacuum at all before testing?
+ */
+int is_no_vacuum = 0;
+
+/*
* use unlogged tables?
*/
int unlogged_tables = 0;
@@ -458,6 +463,11 @@ static void pgbench_error(const char *fmt,...) pg_attribute_printf(1, 2);
static void addScript(ParsedScript script);
static void *threadRun(void *arg);
static void setalarm(int seconds);
+static void initCreateTables(PGconn *con);
+static void initLoadData(PGconn *con);
+static void initVacuum(PGconn *con);
+static void initCreatePKeys(PGconn *con);
+static void initCreateFKeys(PGconn *con);
/* callback functions for our flex lexer */
@@ -484,6 +494,8 @@ usage(void)
" create indexes in the specified tablespace\n"
" --tablespace=TABLESPACE create tables in the specified tablespace\n"
" --unlogged-tables create tables as unlogged tables\n"
+ " --custom-initialize=[tdvpf]+\n"
+ " initialize with custom initialization commands\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"
@@ -2566,9 +2578,37 @@ disconnect_all(CState *state, int length)
}
}
-/* create tables and setup data */
+/* Check custom initialization commands */
+static bool
+checkCustomCmds(char *commands)
+{
+ char *cmd;
+
+ for (cmd = commands; *cmd != '\0'; cmd++)
+ {
+ switch (*cmd)
+ {
+ case 't':
+ case 'd':
+ case 'p':
+ case 'f':
+ case 'v':
+ /* valid commands */
+ break;
+ default:
+ fprintf(stderr, "invalid custom initialization script command \"%c\"\n", *cmd);
+ fprintf(stderr, "possible commands are: \"t\", \"d\", \"p\", \"f\", \"v\"\n");
+ return false;
+ }
+ }
+ return true;
+}
+
+/*
+ * Create tables, remove old tables if exist.
+ */
static void
-init(bool is_no_vacuum)
+initCreateTables(PGconn *con)
{
/*
* The scale factor at/beyond which 32-bit integers are insufficient for
@@ -2623,34 +2663,8 @@ init(bool is_no_vacuum)
1
}
};
- static const char *const DDLINDEXes[] = {
- "alter table pgbench_branches add primary key (bid)",
- "alter table pgbench_tellers add primary key (tid)",
- "alter table pgbench_accounts add primary key (aid)"
- };
- static const char *const DDLKEYs[] = {
- "alter table pgbench_tellers add foreign key (bid) references pgbench_branches",
- "alter table pgbench_accounts add foreign key (bid) references pgbench_branches",
- "alter table pgbench_history add foreign key (bid) references pgbench_branches",
- "alter table pgbench_history add foreign key (tid) references pgbench_tellers",
- "alter table pgbench_history add foreign key (aid) references pgbench_accounts"
- };
- PGconn *con;
- PGresult *res;
- char sql[256];
- int i;
- int64 k;
-
- /* used to track elapsed time and estimate of the remaining time */
- instr_time start,
- diff;
- double elapsed_sec,
- remaining_sec;
- int log_interval = 1;
-
- if ((con = doConnect()) == NULL)
- exit(1);
+ int i;
for (i = 0; i < lengthof(DDLs); i++)
{
@@ -2687,6 +2701,25 @@ init(bool is_no_vacuum)
executeStatement(con, buffer);
}
+}
+
+/*
+ * Fill the standard table with some data.
+ */
+static void
+initLoadData(PGconn *con)
+{
+ char sql[256];
+ PGresult *res;
+ int i;
+ int64 k;
+
+ /* used to track elapsed time and estimate of the remaining time */
+ instr_time start,
+ diff;
+ double elapsed_sec,
+ remaining_sec;
+ int log_interval = 1;
executeStatement(con, "begin");
@@ -2792,20 +2825,56 @@ init(bool is_no_vacuum)
exit(1);
}
executeStatement(con, "commit");
+}
- /* vacuum */
- if (!is_no_vacuum)
+/*
+ * Create foreign key constraints between the standard tables
+ */
+static void
+initCreateFKeys(PGconn *con)
+{
+ static const char *const DDLKEYs[] = {
+ "alter table pgbench_tellers add foreign key (bid) references pgbench_branches",
+ "alter table pgbench_accounts add foreign key (bid) references pgbench_branches",
+ "alter table pgbench_history add foreign key (bid) references pgbench_branches",
+ "alter table pgbench_history add foreign key (tid) references pgbench_tellers",
+ "alter table pgbench_history add foreign key (aid) references pgbench_accounts"
+ };
+
+ int i;
+
+ fprintf(stderr, "set foreign keys...\n");
+ for (i = 0; i < lengthof(DDLKEYs); i++)
{
- fprintf(stderr, "vacuum...\n");
- executeStatement(con, "vacuum analyze pgbench_branches");
- executeStatement(con, "vacuum analyze pgbench_tellers");
- executeStatement(con, "vacuum analyze pgbench_accounts");
- executeStatement(con, "vacuum analyze pgbench_history");
+ executeStatement(con, DDLKEYs[i]);
}
+}
+
+/* Invoke vacuum on all tables */
+static void
+initVacuum(PGconn *con)
+{
+ fprintf(stderr, "vacuum...\n");
+ executeStatement(con, "vacuum analyze pgbench_branches");
+ executeStatement(con, "vacuum analyze pgbench_tellers");
+ executeStatement(con, "vacuum analyze pgbench_accounts");
+ executeStatement(con, "vacuum analyze pgbench_history");
+}
+
+/*
+ * Create primary keys on three tables; pgbench_accounts,
+ * pgbench_branches and pgbench_tellers.
+ */
+static void
+initCreatePKeys(PGconn *con)
+{
+ static const char *const DDLINDEXes[] = {
+ "alter table pgbench_branches add primary key (bid)",
+ "alter table pgbench_tellers add primary key (tid)",
+ "alter table pgbench_accounts add primary key (aid)"
+ };
+ int i;
- /*
- * create indexes
- */
fprintf(stderr, "set primary keys...\n");
for (i = 0; i < lengthof(DDLINDEXes); i++)
{
@@ -2826,16 +2895,45 @@ init(bool is_no_vacuum)
executeStatement(con, buffer);
}
+}
- /*
- * create foreign keys
- */
- if (foreign_keys)
+/* invoke each initialization commands */
+static void
+init(char *initialize_cmds)
+{
+ char *cmd = initialize_cmds;
+ PGconn *con;
+
+ if ((con = doConnect()) == NULL)
+ exit(1);
+
+ for (cmd = initialize_cmds; *cmd != '\0'; cmd++)
{
- fprintf(stderr, "set foreign keys...\n");
- for (i = 0; i < lengthof(DDLKEYs); i++)
+ switch (*cmd)
{
- executeStatement(con, DDLKEYs[i]);
+ case 't':
+ initCreateTables(con);
+ break;
+ case 'd':
+ initLoadData(con);
+ break;
+ case 'p':
+ initCreatePKeys(con);
+ break;
+ case 'f':
+ initCreateFKeys(con);
+ break;
+ case 'v':
+ initVacuum(con);
+ break;
+ case ' ':
+ break; /* skip */
+ default:
+ {
+ fprintf(stderr, "invalid custom initialization script command \"%c\"\n", *cmd);
+ PQfinish(con);
+ exit(1);
+ }
}
}
@@ -3615,6 +3713,7 @@ main(int argc, char **argv)
{"fillfactor", required_argument, NULL, 'F'},
{"host", required_argument, NULL, 'h'},
{"initialize", no_argument, NULL, 'i'},
+ {"custom-initialize", required_argument, NULL, 'I'},
{"jobs", required_argument, NULL, 'j'},
{"log", no_argument, NULL, 'l'},
{"latency-limit", required_argument, NULL, 'L'},
@@ -3646,7 +3745,8 @@ main(int argc, char **argv)
int c;
int is_init_mode = 0; /* initialize mode? */
- int is_no_vacuum = 0; /* no vacuum at all before testing? */
+ int is_custom_init = 0;
+ char *initialize_cmds = "tdvp";
int do_vacuum_accounts = 0; /* do vacuum accounts before testing? */
int optindex;
bool scale_given = false;
@@ -3707,7 +3807,7 @@ main(int argc, char **argv)
state = (CState *) pg_malloc(sizeof(CState));
memset(state, 0, sizeof(CState));
- while ((c = getopt_long(argc, argv, "ih:nvp:dqb:SNc:j:Crs:t:T:U:lf:D:F:M:P:R:L:", long_options, &optindex)) != -1)
+ while ((c = getopt_long(argc, argv, "iI:h:nvp:dqb:SNc:j:Crs:t:T:U:lf:D:F:M:P:R:L:", long_options, &optindex)) != -1)
{
char *script;
@@ -3716,6 +3816,15 @@ main(int argc, char **argv)
case 'i':
is_init_mode++;
break;
+ case 'I':
+ initialize_cmds = pg_strdup(optarg);
+
+ if (!checkCustomCmds(initialize_cmds))
+ exit(1);
+
+ is_custom_init++;
+ initialization_option_set = true;
+ break;
case 'h':
pghost = pg_strdup(optarg);
break;
@@ -4061,7 +4170,47 @@ main(int argc, char **argv)
exit(1);
}
- init(is_no_vacuum);
+ if (initialize_cmds[0] == '\0')
+ {
+ fprintf(stderr, "initialize command is empty\n");
+ exit(1);
+ }
+
+ /*
+ * If custom initialization option and other options related to initialization
+ * are specified at the same time, we process them. Even if vacuum initialization
+ * command is specified in custom initialization commands we remove it if --no-vacuum
+ * is specified. Also, if building foreign keys command is not specified in custom
+ * initialization commands but --foreign-key option is specified, building foreign
+ * keys command is added to at end of the commands.
+ */
+ if (is_custom_init && (is_no_vacuum || foreign_keys))
+ {
+ char *cmd;
+ int n_cmds = 0;
+ bool fk_given = false;
+
+ for (cmd = initialize_cmds; *cmd != '\0'; cmd++)
+ {
+ if (*cmd == 'v' && is_no_vacuum)
+ *cmd = ' ';
+ if (*cmd== 'f')
+ fk_given = true;
+ n_cmds++;
+ }
+
+ if (!fk_given && foreign_keys)
+ {
+ /* Reallocate space for 'f' command */
+ n_cmds++;
+ initialize_cmds = (char *) pg_realloc(initialize_cmds,
+ sizeof(char) * n_cmds);
+ initialize_cmds[n_cmds -1] = 'f';
+ initialize_cmds[n_cmds] = '\0';
+ }
+ }
+
+ init(initialize_cmds);
exit(0);
}
else