Added connection DB flag to tsbs_load_timescaledb. Flag description i…

…s a bit

on the longer side but it makes a clear point not to confuse it with a
similar flag used for specifying the benchmark database name.

cmd/tsbs_load_timescaledb/creator.go

@@ -18,6 +18,7 @@ type dbCreator struct {
 	tags    string
 	cols    []string
 	connStr string
+	connDB  string
 }
 
 func (d *dbCreator) Init() {
@@ -26,6 +27,10 @@ func (d *dbCreator) Init() {
 	// Needed to connect to user's database in order to drop/create db-name database
 	re := regexp.MustCompile(`(dbname)=\S*\b`)
 	d.connStr = strings.TrimSpace(re.ReplaceAllString(d.connStr, ""))
+
+	if d.connDB != "" {
+		d.connStr = fmt.Sprintf("dbname=%s %s", d.connDB, d.connStr)
+	}
 }
 
 func (d *dbCreator) readDataHeader(br *bufio.Reader) {

cmd/tsbs_load_timescaledb/creator_test.go

@@ -12,6 +12,7 @@ func TestDBCreatorInit(t *testing.T) {
 	cases := []struct {
 		desc    string
 		connStr string
+		connDB  string
 		want    string
 	}{
 		{
@@ -29,10 +30,28 @@ func TestDBCreatorInit(t *testing.T) {
 			connStr: "dbname=test2 host=localhost dbname=test1 user=foo dbname=test3",
 			want:    "host=localhost  user=foo",
 		},
+		{
+			desc:    "add dbname by specifying a connDB",
+			connStr: "host=localhost user=foo",
+			connDB:  "bar",
+			want:    "dbname=bar host=localhost user=foo",
+		},
+		{
+			desc:    "override once dbname by specifying a connDB",
+			connStr: "host=localhost dbname=test1 user=foo",
+			connDB:  "bar",
+			want:    "dbname=bar host=localhost  user=foo",
+		},
+		{
+			desc:    "override all dbnames by specifying a connDB",
+			connStr: "dbname=test2 host=localhost dbname=test1 user=foo dbname=test3",
+			connDB:  "bar",
+			want:    "dbname=bar host=localhost  user=foo",
+		},
 	}
 	for _, c := range cases {
 		br := bufio.NewReader(bytes.NewBufferString(buf))
-		dbc := &dbCreator{br: br, connStr: c.connStr}
+		dbc := &dbCreator{br: br, connStr: c.connStr, connDB: c.connDB}
 		dbc.Init()
 		if got := dbc.connStr; got != c.want {
 			t.Errorf("%s: incorrect connstr: got %s want %s", c.desc, got, c.want)

cmd/tsbs_load_timescaledb/main.go

@@ -30,6 +30,7 @@ var (
 	user            string
 	pass            string
 	port            string
+	connDB          string
 
 	useHypertable bool
 	logBatches    bool
@@ -72,6 +73,9 @@ func init() {
 	flag.StringVar(&port, "port", "5432", "Which port to connect to on the database host")
 	flag.StringVar(&user, "user", "postgres", "User to connect to PostgreSQL as")
 	flag.StringVar(&pass, "pass", "", "Password for user connecting to PostgreSQL (leave blank if not password protected)")
+	flag.StringVar(&connDB, "admin-db-name", user, "Database to connect to in order to create additional benchmark databases.\n"+
+		"By default this is the same as the `user` (i.e., `postgres` if neither is set),\n"+
+		"but sometimes a user does not have its own database.")
 
 	flag.BoolVar(&logBatches, "log-batches", false, "Whether to time individual batches.")
 
@@ -122,6 +126,7 @@ func (b *benchmark) GetDBCreator() load.DBCreator {
 	return &dbCreator{
 		br:      loader.GetBufferedReader(),
 		connStr: getConnectString(),
+		connDB:  connDB,
 	}
 }
 

docs/timescaledb.md

@@ -61,6 +61,12 @@ Whether to actually use TimescaleDB's hypertable for storing data. Set to
 
 User to use to connect to the PostgreSQL server.
 
+#### `--admin-db-name` (type: `string`, default: value of `-user`)
+
+Database to connect to in order to create additional benchmark databases.
+By default this is the same as the `user` (i.e., `postgres` if neither is set),
+but sometimes a user does not have its own database.
+
 ### Tags related
 
 #### `-in-table-partition-tag` (type: `boolean`, default: `false`)