github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/tests/data.go (about)

     1  // Copyright 2017 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 tests
    12  
    13  import (
    14  	"bytes"
    15  	"context"
    16  	gosql "database/sql"
    17  	"fmt"
    18  	"testing"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/kv"
    21  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    22  )
    23  
    24  // CheckKeyCount checks that the number of keys in the provided span matches
    25  // numKeys.
    26  func CheckKeyCount(t *testing.T, kvDB *kv.DB, span roachpb.Span, numKeys int) {
    27  	t.Helper()
    28  	if kvs, err := kvDB.Scan(context.TODO(), span.Key, span.EndKey, 0); err != nil {
    29  		t.Fatal(err)
    30  	} else if l := numKeys; len(kvs) != l {
    31  		t.Fatalf("expected %d key value pairs, but got %d", l, len(kvs))
    32  	}
    33  }
    34  
    35  // CreateKVTable creates a basic table named t.<name> that stores key/value
    36  // pairs with numRows of arbitrary data.
    37  func CreateKVTable(sqlDB *gosql.DB, name string, numRows int) error {
    38  	// Fix the column families so the key counts don't change if the family
    39  	// heuristics are updated.
    40  	schema := fmt.Sprintf(`
    41  		CREATE DATABASE IF NOT EXISTS t;
    42  		CREATE TABLE t.%s (k INT PRIMARY KEY, v INT, FAMILY (k), FAMILY (v));
    43  		CREATE INDEX foo on t.%s (v);`, name, name)
    44  
    45  	if _, err := sqlDB.Exec(schema); err != nil {
    46  		return err
    47  	}
    48  
    49  	// Bulk insert.
    50  	var insert bytes.Buffer
    51  	if _, err := insert.WriteString(
    52  		fmt.Sprintf(`INSERT INTO t.%s VALUES (%d, %d)`, name, 0, numRows-1)); err != nil {
    53  		return err
    54  	}
    55  	for i := 1; i < numRows; i++ {
    56  		if _, err := insert.WriteString(fmt.Sprintf(` ,(%d, %d)`, i, numRows-i)); err != nil {
    57  			return err
    58  		}
    59  	}
    60  	_, err := sqlDB.Exec(insert.String())
    61  	return err
    62  }
    63  
    64  // CreateKVInterleavedTable is like CreateKVTable, but it interleaves table
    65  // t.intlv inside of t.kv and adds rows to both.
    66  func CreateKVInterleavedTable(t *testing.T, sqlDB *gosql.DB, numRows int) {
    67  	// Fix the column families so the key counts don't change if the family
    68  	// heuristics are updated.
    69  	if _, err := sqlDB.Exec(`
    70  CREATE DATABASE t;
    71  SET DATABASE=t;
    72  CREATE TABLE kv (k INT PRIMARY KEY, v INT);
    73  CREATE TABLE intlv (k INT, m INT, n INT, PRIMARY KEY (k, m)) INTERLEAVE IN PARENT kv (k);
    74  CREATE INDEX intlv_idx ON intlv (k, n) INTERLEAVE IN PARENT kv (k);
    75  `); err != nil {
    76  		t.Fatal(err)
    77  	}
    78  
    79  	var insert bytes.Buffer
    80  	if _, err := insert.WriteString(fmt.Sprintf(`INSERT INTO t.kv VALUES (%d, %d)`, 0, numRows-1)); err != nil {
    81  		t.Fatal(err)
    82  	}
    83  	for i := 1; i < numRows; i++ {
    84  		if _, err := insert.WriteString(fmt.Sprintf(` ,(%d, %d)`, i, numRows-i)); err != nil {
    85  			t.Fatal(err)
    86  		}
    87  	}
    88  	if _, err := sqlDB.Exec(insert.String()); err != nil {
    89  		t.Fatal(err)
    90  	}
    91  	insert.Reset()
    92  	if _, err := insert.WriteString(fmt.Sprintf(`INSERT INTO t.intlv VALUES (%d, %d, %d)`, 0, numRows-1, numRows-1)); err != nil {
    93  		t.Fatal(err)
    94  	}
    95  	for i := 1; i < numRows; i++ {
    96  		if _, err := insert.WriteString(fmt.Sprintf(` ,(%d, %d, %d)`, i, numRows-i, numRows-i)); err != nil {
    97  			t.Fatal(err)
    98  		}
    99  	}
   100  	if _, err := sqlDB.Exec(insert.String()); err != nil {
   101  		t.Fatal(err)
   102  	}
   103  }