diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 03e1212..add6b88 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -25,6 +25,12 @@
pgbench
+
+ option
+ dbname
+
+
+ pgbench
option
dbname
@@ -177,6 +183,8 @@ pgbench options dbname
Perform no vacuuming after initialization.
+ This option cannot be set with
+ at the same time.
@@ -211,10 +219,70 @@ 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.
+
+
+
+
+
+
+
+
+
Create foreign key constraints between the standard tables.
+ This option cannot be set with
+ at the same time.
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index ae78c7b..22d5ad9 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={t|d|p|f|v] [...]}\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,39 @@ disconnect_all(CState *state, int length)
}
}
-/* create tables and setup data */
+/* Check custom initialization commands */
+static bool
+checkCustomCmds(char *initialize_cmds)
+{
+ char *cmd;
+
+ for (cmd = initialize_cmds; *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 +2665,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 +2703,95 @@ init(bool is_no_vacuum)
executeStatement(con, buffer);
}
+}
+
+/* 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 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++)
+ {
+ executeStatement(con, DDLKEYs[i]);
+ }
+}
+
+/*
+ * 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;
+
+ fprintf(stderr, "set primary keys...\n");
+ for (i = 0; i < lengthof(DDLINDEXes); i++)
+ {
+ char buffer[256];
+
+ strlcpy(buffer, DDLINDEXes[i], sizeof(buffer));
+
+ if (index_tablespace != NULL)
+ {
+ char *escape_tablespace;
+
+ escape_tablespace = PQescapeIdentifier(con, index_tablespace,
+ strlen(index_tablespace));
+ snprintf(buffer + strlen(buffer), sizeof(buffer) - strlen(buffer),
+ " using index tablespace %s", escape_tablespace);
+ PQfreemem(escape_tablespace);
+ }
+
+ 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,50 +2897,45 @@ init(bool is_no_vacuum)
exit(1);
}
executeStatement(con, "commit");
+}
- /* vacuum */
- if (!is_no_vacuum)
- {
- 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 indexes
- */
- fprintf(stderr, "set primary keys...\n");
- for (i = 0; i < lengthof(DDLINDEXes); i++)
- {
- char buffer[256];
-
- strlcpy(buffer, DDLINDEXes[i], sizeof(buffer));
-
- if (index_tablespace != NULL)
- {
- char *escape_tablespace;
-
- escape_tablespace = PQescapeIdentifier(con, index_tablespace,
- strlen(index_tablespace));
- snprintf(buffer + strlen(buffer), sizeof(buffer) - strlen(buffer),
- " using index tablespace %s", escape_tablespace);
- PQfreemem(escape_tablespace);
- }
+/* invoke each initialization commands */
+static void
+init(char *initialize_cmds, bool is_custom)
+{
+ char *cmd = initialize_cmds;
+ PGconn *con;
- executeStatement(con, buffer);
- }
+ if ((con = doConnect()) == NULL)
+ exit(1);
- /*
- * create foreign keys
- */
- if (foreign_keys)
+ 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':
+ if (is_custom || foreign_keys)
+ initCreateFKeys(con);
+ break;
+ case 'v':
+ if (is_custom || !is_no_vacuum)
+ initVacuum(con);
+ break;
+ default:
+ {
+ fprintf(stderr, "invalid custom initialization script command \"%c\"\n", *cmd);
+ PQfinish(con);
+ exit(1);
+ }
}
}
@@ -3641,12 +3741,14 @@ main(int argc, char **argv)
{"aggregate-interval", required_argument, NULL, 5},
{"progress-timestamp", no_argument, NULL, 6},
{"log-prefix", required_argument, NULL, 7},
+ {"custom-initialize", required_argument, NULL, 8},
{NULL, 0, NULL, 0}
};
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 = "tdpv";
int do_vacuum_accounts = 0; /* do vacuum accounts before testing? */
int optindex;
bool scale_given = false;
@@ -3944,8 +4046,11 @@ main(int argc, char **argv)
break;
case 0:
/* This covers long options which take no argument. */
- if (foreign_keys || unlogged_tables)
+ if (unlogged_tables || foreign_keys)
initialization_option_set = true;
+
+ if (foreign_keys)
+ initialize_cmds = "tdpvf";
break;
case 2: /* tablespace */
initialization_option_set = true;
@@ -3982,6 +4087,18 @@ main(int argc, char **argv)
benchmarking_option_set = true;
logfile_prefix = pg_strdup(optarg);
break;
+ case 8:
+ initialize_cmds = pg_strdup(optarg);
+
+ if (!checkCustomCmds(initialize_cmds))
+ {
+ fprintf(stderr, "invalid syntax of custom initialization \"%s\"\n",
+ initialize_cmds);
+ exit(1);
+ }
+ is_custom_init++;
+ initialization_option_set = true;
+ break;
default:
fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
exit(1);
@@ -4061,7 +4178,19 @@ main(int argc, char **argv)
exit(1);
}
- init(is_no_vacuum);
+ if (is_custom_init && is_no_vacuum)
+ {
+ fprintf(stderr, "cannot specify no vacuum (-n) and custom initialization mode at the same time\n");
+ exit(1);
+ }
+
+ if (is_custom_init && foreign_keys)
+ {
+ fprintf(stderr, "cannot specify to create foreign keys (--foreign-keys) and custom initialization mode at the same time\n");
+ exit(1);
+ }
+
+ init(initialize_cmds, is_custom_init);
exit(0);
}
else