github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/bench/query.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 bench
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"math/rand"
    17  	"net"
    18  	"net/url"
    19  	"os/exec"
    20  
    21  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    22  )
    23  
    24  // This is the TPC-B(ish) query that pgbench runs.
    25  // We don't use placeholders because pgwire protocol does not
    26  // allow multiple statements in prepared queries.
    27  const tpcbQuery = `BEGIN;
    28  UPDATE pgbench_accounts SET abalance = abalance + %[1]d WHERE aid = %[2]d;
    29  SELECT abalance FROM pgbench_accounts WHERE aid = %[2]d;
    30  UPDATE pgbench_tellers SET tbalance = tbalance + %[1]d WHERE tid = %[3]d;
    31  UPDATE pgbench_branches SET bbalance = bbalance + %[1]d WHERE bid = %[4]d;
    32  INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (%[3]d, %[4]d, %[2]d, %[1]d, CURRENT_TIMESTAMP);
    33  END;` // vars: 1 delta, 2 aid, 3 tid, 4 bid
    34  
    35  // RunOne executes one iteration of the query batch that `pgbench` executes.
    36  // Calls b.Fatalf if it encounters an error.
    37  func RunOne(db sqlutils.DBHandle, r *rand.Rand, accounts int) error {
    38  	account := r.Intn(accounts)
    39  	delta := r.Intn(5000)
    40  	teller := r.Intn(tellers)
    41  	branch := 1
    42  
    43  	q := fmt.Sprintf(tpcbQuery, delta, account, teller, branch)
    44  	_, err := db.ExecContext(context.TODO(), q)
    45  	return err
    46  }
    47  
    48  // ExecPgbench returns a ready-to-run pgbench Cmd, that will run
    49  // against the db specified by `pgURL`.
    50  func ExecPgbench(pgURL url.URL, dbname string, count int) (*exec.Cmd, error) {
    51  	host, port, err := net.SplitHostPort(pgURL.Host)
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  
    56  	args := []string{
    57  		"-n", // disable (pg-specific) vacuum
    58  		"-r", // print stats
    59  		fmt.Sprintf("--transactions=%d", count),
    60  		"-h", host,
    61  		"-p", port,
    62  	}
    63  
    64  	if pgURL.User != nil {
    65  		if user := pgURL.User.Username(); user != "" {
    66  			args = append(args, "-U", user)
    67  		}
    68  	}
    69  	args = append(args, dbname)
    70  
    71  	return exec.Command("pgbench", args...), nil
    72  }