github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/utilccl/sampledataccl/bankdata_test.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package sampledataccl
    10  
    11  import (
    12  	"bytes"
    13  	"context"
    14  	"fmt"
    15  	"io"
    16  	"path/filepath"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/base"
    20  	"github.com/cockroachdb/cockroach/pkg/keys"
    21  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    23  	"github.com/cockroachdb/cockroach/pkg/testutils"
    24  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    26  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    27  	"github.com/cockroachdb/cockroach/pkg/workload/bank"
    28  )
    29  
    30  func TestToBackup(t *testing.T) {
    31  	defer leaktest.AfterTest(t)()
    32  
    33  	outerDir, dirCleanupFn := testutils.TempDir(t)
    34  	defer dirCleanupFn()
    35  
    36  	ctx := context.Background()
    37  	args := base.TestServerArgs{
    38  		ExternalIODir: outerDir,
    39  		UseDatabase:   "data",
    40  	}
    41  	s, db, _ := serverutils.StartServer(t, args)
    42  	defer s.Stopper().Stop(ctx)
    43  
    44  	const payloadBytes, ranges = 100, 10
    45  	chunkBytesSizes := []int64{
    46  		0,                // 0 means default ~32MB
    47  		3 * payloadBytes, // a number that will not evently divide the number of rows
    48  	}
    49  
    50  	for _, rows := range []int{1, 10} {
    51  		for _, chunkBytes := range chunkBytesSizes {
    52  			t.Run(fmt.Sprintf("rows=%d/chunk=%d", rows, chunkBytes), func(t *testing.T) {
    53  				dir := fmt.Sprintf("%d-%d", rows, chunkBytes)
    54  				data := bank.FromConfig(rows, rows, payloadBytes, ranges).Tables()[0]
    55  				backup, err := toBackup(t, data, filepath.Join(outerDir, dir), chunkBytes)
    56  				if err != nil {
    57  					t.Fatalf("%+v", err)
    58  				}
    59  
    60  				t.Run("Restore", func(t *testing.T) {
    61  					sqlDB := sqlutils.MakeSQLRunner(db)
    62  					sqlDB.Exec(t, `DROP DATABASE IF EXISTS data CASCADE`)
    63  					sqlDB.Exec(t, `CREATE DATABASE data`)
    64  					sqlDB.Exec(t, `RESTORE data.* FROM $1`, `nodelocal://0/`+dir)
    65  
    66  					var rowCount int
    67  					sqlDB.QueryRow(t, fmt.Sprintf(`SELECT count(*) FROM %s`, data.Name)).Scan(&rowCount)
    68  					if rowCount != rows {
    69  						t.Errorf("got %d rows expected %d", rowCount, rows)
    70  					}
    71  				})
    72  
    73  				t.Run("NextKeyValues", func(t *testing.T) {
    74  					for _, requestedKVs := range []int{2, 3} {
    75  						newTableID := sqlbase.ID(keys.MaxReservedDescID + requestedKVs)
    76  						newTablePrefix := keys.SystemSQLCodec.TablePrefix(uint32(newTableID))
    77  
    78  						keys := make(map[string]struct{}, rows)
    79  						for {
    80  							kvs, span, err := backup.NextKeyValues(requestedKVs, newTableID)
    81  							if err != nil && err != io.EOF {
    82  								t.Fatalf("%+v", err)
    83  							}
    84  							if len(kvs) < requestedKVs && err != io.EOF {
    85  								t.Errorf("got %d kvs requested at least %d", len(kvs), requestedKVs)
    86  							}
    87  							for k := range keys {
    88  								key := roachpb.Key(k)
    89  								if key.Compare(span.Key) >= 0 && key.Compare(span.EndKey) < 0 {
    90  									t.Errorf("previous key %s overlaps this span %s", key, span)
    91  								}
    92  							}
    93  							for _, kv := range kvs {
    94  								key := kv.Key.Key
    95  								if key.Compare(span.Key) < 0 || key.Compare(span.EndKey) >= 0 {
    96  									t.Errorf("key %s is not in %s", kv.Key, span)
    97  								}
    98  								if _, ok := keys[string(key)]; ok {
    99  									t.Errorf("key %s was output twice", key)
   100  								}
   101  								keys[string(key)] = struct{}{}
   102  								if !bytes.HasPrefix(key, newTablePrefix) {
   103  									t.Errorf("key %s is not for table %d", key, newTableID)
   104  								}
   105  							}
   106  							if err == io.EOF {
   107  								break
   108  							}
   109  						}
   110  
   111  						if len(keys) != rows {
   112  							t.Errorf("got %d kvs expected %d", len(keys), rows)
   113  						}
   114  						// Don't reconstruct the backup between runs so we can
   115  						// test the reset.
   116  						backup.ResetKeyValueIteration()
   117  					}
   118  				})
   119  			})
   120  		}
   121  	}
   122  }