CockroachDB bulk insert test case

main_test.go

package main

import (
	"bytes"
	"database/sql"
	"flag"
	"html/template"
	"strconv"
	"strings"
	"testing"
	"time"

	_ "github.com/lib/pq"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	// defaultDSN is the data source name for a CockroachDB instance running
	// locally, started with:
	//     cockroach start --insecure --listen-addr=localhost
	// You will need to create the testdb database:
	//     echo "CREATE DATABASE testdb;" | psql -h localhost -p 26257 -U root
	defaultDSN = "postgresql://root@localhost:26257/testdb?sslmode=disable"
)

var (
	dsn = flag.String("dsn", defaultDSN, "data source name")
)

func TestCockroachDBBulkInsert(t *testing.T) {
	batchSize := 1000
	numBatches := 10
	valueSize := 16 * 1024

	// Open the connection to the database.
	db, err := sql.Open("postgres", *dsn)
	require.NoError(t, err)
	defer db.Close()

	// Check the connection.
	require.NoError(t, db.Ping())

	// Create a unique table name for this test.
	tableName := "test_pairs_" + strconv.Itoa(int(time.Now().UnixNano()))

	data := struct {
		TableName string
		BatchSize int
	}{
		TableName: tableName,
		BatchSize: batchSize,
	}
	funcs := template.FuncMap{
		"makeInsertValues": makeInsertValues,
	}

	// Create the table and ensure it is cleaned up afterwards.
	createTableQuery, err := executeTemplate("create_table_query", `
		CREATE TABLE {{ .TableName }} (
			key BYTEA PRIMARY KEY,
			value BYTEA NOT NULL
		);
	`, data, nil)
	require.NoError(t, err)
	dropTableQuery, err := executeTemplate("drop_table_query", `
		DROP TABLE {{ .TableName }};
	`, data, nil)
	require.NoError(t, err)
	_, err = db.Exec(createTableQuery)
	require.NoError(t, err)
	defer func() {
		_, err := db.Exec(dropTableQuery)
		assert.NoError(t, err)
	}()

	// Create a bulk insert query.
	insertQuery, err := executeTemplate("insert_query", `
		INSERT INTO {{ .TableName }} (key, value)
		VALUES {{ makeInsertValues 2 ", " .BatchSize ", "}}
		ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value;
	`, data, funcs)
	require.NoError(t, err)

	// Insert numBatches batches of batchSize pairs with unique keys and values of valueSize.
	value := make([]byte, valueSize)
	for i := 0; i < numBatches; i++ {
		args := make([]interface{}, 0, 2*batchSize)
		for j := 0; j < batchSize; j++ {
			key := []byte(strconv.Itoa(j))
			args = append(args, key, value)
		}
		_, err := db.Exec(insertQuery, args...)
		require.NoError(t, err)
	}
}

func executeTemplate(name, text string, data interface{}, funcs template.FuncMap) (string, error) {
	t, err := template.New(name).Funcs(funcs).Parse(text)
	if err != nil {
		return "", err
	}
	b := &bytes.Buffer{}
	if err := t.Execute(b, data); err != nil {
		return "", err
	}
	return b.String(), nil
}

func makeInsertValues(cols int, colSep string, rows int, rowSep string) string {
	b := &strings.Builder{}
	for r := 0; r < rows; r++ {
		if r != 0 {
			b.WriteString(rowSep)
		}
		b.WriteByte('(')
		for c := 0; c < cols; c++ {
			if c != 0 {
				b.WriteString(colSep)
			}
			b.WriteByte('$')
			b.WriteString(strconv.Itoa(r*cols + c + 1))
		}
		b.WriteByte(')')
	}
	return b.String()
}