github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/bench/cmd/pgbenchsetup/main.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package main
    12  
    13  import (
    14  	gosql "database/sql"
    15  	"flag"
    16  	"fmt"
    17  	"net/url"
    18  	"os"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/bench"
    21  	_ "github.com/lib/pq"
    22  )
    23  
    24  var usage = func() {
    25  	fmt.Fprintln(os.Stderr, "Creates the schema and initial data used by the `pgbench` tool")
    26  	fmt.Fprintf(os.Stderr, "\nUsage: %s <db URL>\n", os.Args[0])
    27  	flag.PrintDefaults()
    28  }
    29  
    30  func main() {
    31  	accounts := flag.Int("accounts", 100000, "number of accounts to create")
    32  
    33  	createDb := flag.Bool("createdb", false, "attempt to create named db, dropping first if exists (must be able to connect to default db to do so).")
    34  
    35  	flag.Parse()
    36  	flag.Usage = usage
    37  	if flag.NArg() != 1 {
    38  		flag.Usage()
    39  		os.Exit(2)
    40  	}
    41  
    42  	var db *gosql.DB
    43  	var err error
    44  
    45  	if *createDb {
    46  		name := ""
    47  		parsed, parseErr := url.Parse(flag.Arg(0))
    48  		if parseErr != nil {
    49  			panic(parseErr)
    50  		} else if len(parsed.Path) < 2 { // first char is '/'
    51  			panic("URL must include db name")
    52  		} else {
    53  			name = parsed.Path[1:]
    54  		}
    55  
    56  		db, err = bench.CreateAndConnect(*parsed, name)
    57  	} else {
    58  		db, err = gosql.Open("postgres", flag.Arg(0))
    59  	}
    60  	if err != nil {
    61  		panic(err)
    62  	}
    63  
    64  	defer db.Close()
    65  
    66  	if err := bench.SetupBenchDB(db, *accounts, false); err != nil {
    67  		panic(err)
    68  	}
    69  }