github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/bench/setup.go

// Copyright 2016 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package bench

import (
	"bytes"
	"context"
	gosql "database/sql"
	"fmt"
	"net/url"
	"os/exec"

	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
)

const schema = `
DROP TABLE IF EXISTS pgbench_accounts;
DROP TABLE IF EXISTS pgbench_branches;
DROP TABLE IF EXISTS pgbench_tellers;
DROP TABLE IF EXISTS pgbench_history;

CREATE TABLE pgbench_accounts (
    aid integer NOT NULL PRIMARY KEY,
    bid integer,
    abalance integer,
    filler character(84)
);

CREATE TABLE pgbench_branches (
    bid integer NOT NULL PRIMARY KEY,
    bbalance integer,
    filler character(88)
);

CREATE TABLE pgbench_tellers (
    tid integer NOT NULL PRIMARY KEY,
    bid integer,
    tbalance integer,
    filler character(84)
);

CREATE TABLE pgbench_history (
    tid integer,
    bid integer,
    aid integer,
    delta integer,
    mtime timestamp,
    filler character(22)
);
`

// CreateAndConnect connects and creates the requested DB (dropping
// if exists) then returns a new connection to the created DB.
func CreateAndConnect(pgURL url.URL, name string) (*gosql.DB, error) {
	{
		pgURL.Path = ""
		db, err := gosql.Open("postgres", pgURL.String())
		if err != nil {
			return nil, err
		}
		defer db.Close()

		if _, err := db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", name)); err != nil {
			return nil, err
		}

		if _, err := db.Exec(fmt.Sprintf(`CREATE DATABASE %s`, name)); err != nil {
			return nil, err
		}
	}

	pgURL.Path = name

	db, err := gosql.Open("postgres", pgURL.String())
	if err != nil {
		return nil, err
	}
	return db, nil
}

// SetupExec creates and fills a DB and prepares a `pgbench` command
// to be run against it.
func SetupExec(pgURL url.URL, name string, accounts, transactions int) (*exec.Cmd, error) {
	db, err := CreateAndConnect(pgURL, name)
	if err != nil {
		return nil, err
	}
	defer db.Close()

	if err := SetupBenchDB(db, accounts, true /*quiet*/); err != nil {
		return nil, err
	}

	return ExecPgbench(pgURL, name, transactions)
}

// SetupBenchDB sets up a db with the schema and initial data used by `pgbench`.
// The `-i` flag to `pgbench` is usually used to do this when testing postgres
// but the statements it generates use postgres-specific flags that cockroach does
// not support. The queries this script runs are based on a dump of a db created
// by `pgbench -i`, but sticking to the compatible subset that both cockroach and
// postgres support.
func SetupBenchDB(db sqlutils.DBHandle, accounts int, quiet bool) error {
	ctx := context.TODO()
	if _, err := db.ExecContext(ctx, schema); err != nil {
		return err
	}
	return populateDB(ctx, db, accounts, quiet)
}

const tellers = 10

func populateDB(ctx context.Context, db sqlutils.DBHandle, accounts int, quiet bool) error {
	branches := `INSERT INTO pgbench_branches (bid, bbalance, filler) VALUES (1, 7354, NULL)`
	if r, err := db.ExecContext(ctx, branches); err != nil {
		return err
	} else if x, err := r.RowsAffected(); err != nil {
		return err
	} else if !quiet {
		fmt.Printf("Inserted %d branch records\n", x)
	}

	// Various magic numbers came from `pg_dump` of a `pgbench` created DB.
	tellers := `INSERT INTO pgbench_tellers VALUES (1, 1, 0, NULL),
	(2, 1, 955, NULL),
	(3, 1, -3338, NULL),
	(4, 1, -24, NULL),
	(5, 1, 2287, NULL),
	(6, 1, 4129, NULL),
	(7, 1, 0, NULL),
	(8, 1, 0, NULL),
	(9, 1, 0, NULL),
	(10, 1, 3345, NULL)
	`
	if r, err := db.ExecContext(ctx, tellers); err != nil {
		return err
	} else if x, err := r.RowsAffected(); err != nil {
		return err
	} else if !quiet {
		fmt.Printf("Inserted %d teller records\n", x)
	}

	// Split account inserts into batches to avoid giant query.
	done := 0
	for {
		batch := 5000
		remaining := accounts - done
		if remaining < 1 {
			break
		}
		if batch > remaining {
			batch = remaining
		}
		var placeholders bytes.Buffer
		for i := 0; i < batch; i++ {
			if i > 0 {
				placeholders.WriteString(", ")
			}
			fmt.Fprintf(&placeholders, "(%d, 1, 0, '                                                                                    ')", done+i)
		}
		stmt := fmt.Sprintf(`INSERT INTO pgbench_accounts VALUES %s`, placeholders.String())
		if r, err := db.ExecContext(ctx, stmt); err != nil {
			return err
		} else if x, err := r.RowsAffected(); err != nil {
			return err
		} else if !quiet {
			fmt.Printf("Inserted %d account records\n", x)
		}
		done += batch
	}

	history := `
INSERT INTO pgbench_history VALUES
(5, 1, 36833, 407, CURRENT_TIMESTAMP, NULL),
(3, 1, 43082, -3338, CURRENT_TIMESTAMP, NULL),
(2, 1, 49129, 2872, CURRENT_TIMESTAMP, NULL),
(6, 1, 81223, 1064, CURRENT_TIMESTAMP, NULL),
(6, 1, 28316, 3065, CURRENT_TIMESTAMP, NULL),
(4, 1, 10146, -24, CURRENT_TIMESTAMP, NULL),
(10, 1, 12019, 2265, CURRENT_TIMESTAMP, NULL),
(2, 1, 46717, -1917, CURRENT_TIMESTAMP, NULL),
(5, 1, 68648, 1880, CURRENT_TIMESTAMP, NULL),
(10, 1, 46989, 1080, CURRENT_TIMESTAMP, NULL);`

	if r, err := db.ExecContext(ctx, history); err != nil {
		return err
	} else if x, err := r.RowsAffected(); err != nil {
		return err
	} else if !quiet {
		fmt.Printf("Inserted %d history records\n", x)
	}

	return nil
}