func backupRestoreTestSetup(
t testing.TB, numAccounts int,
) (
ctx context.Context,
tempDir string,
tc *testcluster.TestCluster,
kvDB *client.DB,
sqlDB *sqlutils.SQLRunner,
cleanup func(),
) {
ctx = context.Background()
dir, dirCleanupFn := testutils.TempDir(t, 1)
// Use ReplicationManual so we can force full replication, which is needed
// to later move the leases around.
tc = testcluster.StartTestCluster(t, backupRestoreClusterSize, base.TestClusterArgs{
ReplicationMode: base.ReplicationManual,
})
sqlDB = sqlutils.MakeSQLRunner(t, tc.Conns[0])
kvDB = tc.Server(0).KVClient().(*client.DB)
sqlDB.Exec(bankCreateDatabase)
sqlDB.Exec(bankCreateTable)
for _, insert := range bankDataInsertStmts(numAccounts) {
sqlDB.Exec(insert)
}
for _, split := range bankSplitStmts(numAccounts, backupRestoreDefaultRanges) {
sqlDB.Exec(split)
}
targets := make([]testcluster.ReplicationTarget, backupRestoreClusterSize-1)
for i := 1; i < backupRestoreClusterSize; i++ {
targets[i-1] = tc.Target(i)
}
txn := client.NewTxn(ctx, *kvDB)
rangeDescs, err := sql.AllRangeDescriptors(txn)
if err != nil {
t.Fatal(err)
}
for _, r := range rangeDescs {
if _, err := tc.AddReplicas(r.StartKey.AsRawKey(), targets...); err != nil {
t.Fatal(err)
}
}
cleanupFn := func() {
tc.Stopper().Stop()
dirCleanupFn()
}
return ctx, dir, tc, kvDB, sqlDB, cleanupFn
}
func rebalanceLeases(t testing.TB, tc *testcluster.TestCluster) {
kvDB := tc.Server(0).KVClient().(*client.DB)
txn := client.NewTxn(context.Background(), *kvDB)
rangeDescs, err := sql.AllRangeDescriptors(txn)
if err != nil {
t.Fatal(err)
}
for _, r := range rangeDescs {
target := tc.Target(int(r.RangeID) % tc.NumServers())
if err := tc.TransferRangeLease(r, target); err != nil {
t.Fatal(err)
}
}
}