-
Notifications
You must be signed in to change notification settings - Fork 3.8k
/
datadriven_test.go
860 lines (797 loc) · 26.8 KB
/
datadriven_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
// Copyright 2022 The Cockroach Authors.
//
// Licensed as a CockroachDB Enterprise file under the Cockroach Community
// License (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
package backupccl
import (
"context"
gosql "database/sql"
"fmt"
"net/url"
"os"
"path/filepath"
"regexp"
"strconv"
"strings"
"testing"
"time"
"github.com/cockroachdb/cockroach/pkg/base"
"github.com/cockroachdb/cockroach/pkg/ccl/backupccl/backuputils"
"github.com/cockroachdb/cockroach/pkg/clusterversion"
"github.com/cockroachdb/cockroach/pkg/jobs"
"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
"github.com/cockroachdb/cockroach/pkg/kv"
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/server"
"github.com/cockroachdb/cockroach/pkg/settings/cluster"
"github.com/cockroachdb/cockroach/pkg/sql"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
"github.com/cockroachdb/cockroach/pkg/testutils"
"github.com/cockroachdb/cockroach/pkg/testutils/jobutils"
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
"github.com/cockroachdb/cockroach/pkg/testutils/skip"
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/datadriven"
"github.com/cockroachdb/errors"
"github.com/lib/pq"
"github.com/stretchr/testify/require"
)
var localityCfgs = map[string]roachpb.Locality{
"us-east-1": {
Tiers: []roachpb.Tier{
{Key: "region", Value: "us-east-1"},
{Key: "availability-zone", Value: "us-east1"},
},
},
"us-west-1": {
Tiers: []roachpb.Tier{
{Key: "region", Value: "us-west-1"},
{Key: "availability-zone", Value: "us-west1"},
},
},
"eu-central-1": {
Tiers: []roachpb.Tier{
{Key: "region", Value: "eu-central-1"},
{Key: "availability-zone", Value: "eu-central-1"},
},
},
"eu-north-1": {
Tiers: []roachpb.Tier{
{Key: "region", Value: "eu-north-1"},
{Key: "availability-zone", Value: "eu-north-1"},
},
},
}
var clusterVersionKeys = map[string]clusterversion.Key{
"Start22_2": clusterversion.Start22_2,
}
type sqlDBKey struct {
server string
user string
}
type datadrivenTestState struct {
servers map[string]serverutils.TestServerInterface
dataDirs map[string]string
sqlDBs map[sqlDBKey]*gosql.DB
jobTags map[string]jobspb.JobID
clusterTimestamps map[string]string
noticeBuffer []string
cleanupFns []func()
vars map[string]string
}
func newDatadrivenTestState() datadrivenTestState {
return datadrivenTestState{
servers: make(map[string]serverutils.TestServerInterface),
dataDirs: make(map[string]string),
sqlDBs: make(map[sqlDBKey]*gosql.DB),
jobTags: make(map[string]jobspb.JobID),
clusterTimestamps: make(map[string]string),
vars: make(map[string]string),
}
}
func (d *datadrivenTestState) cleanup(ctx context.Context, t *testing.T) {
// While the testCluster cleanupFns would close the dbConn and servers, close
// them manually to ensure all queries finish on tests that share these
// resources.
for _, db := range d.sqlDBs {
backuputils.CheckForInvalidDescriptors(t, db)
db.Close()
}
for _, s := range d.servers {
s.Stopper().Stop(ctx)
}
for _, f := range d.cleanupFns {
f()
}
d.noticeBuffer = nil
}
type serverCfg struct {
name string
iodir string
nodes int
splits int
ioConf base.ExternalIODirConfig
localities string
beforeVersion string
testingKnobCfg string
}
func (d *datadrivenTestState) addServer(t *testing.T, cfg serverCfg) error {
var tc serverutils.TestClusterInterface
var cleanup func()
params := base.TestClusterArgs{}
params.ServerArgs.ExternalIODirConfig = cfg.ioConf
params.ServerArgs.Knobs = base.TestingKnobs{
JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
}
settings := cluster.MakeTestingClusterSettings()
if cfg.beforeVersion != "" {
beforeKey, ok := clusterVersionKeys[cfg.beforeVersion]
if !ok {
t.Fatalf("clusterVersion %s does not exist in data driven global map", cfg.beforeVersion)
}
beforeKey--
params.ServerArgs.Knobs.Server = &server.TestingKnobs{
BinaryVersionOverride: clusterversion.ByKey(beforeKey),
DisableAutomaticVersionUpgrade: make(chan struct{})}
settings = cluster.MakeTestingClusterSettingsWithVersions(
clusterversion.TestingBinaryVersion,
clusterversion.ByKey(beforeKey),
false,
)
}
closedts.TargetDuration.Override(context.Background(), &settings.SV, 10*time.Millisecond)
closedts.SideTransportCloseInterval.Override(context.Background(), &settings.SV, 10*time.Millisecond)
sql.TempObjectWaitInterval.Override(context.Background(), &settings.SV, time.Millisecond)
params.ServerArgs.Settings = settings
clusterSize := cfg.nodes
if cfg.localities != "" {
cfgs := strings.Split(cfg.localities, ",")
clusterSize = len(cfgs)
serverArgsPerNode := make(map[int]base.TestServerArgs)
for i, cfg := range cfgs {
param := params.ServerArgs
param.Locality = localityCfgs[cfg]
serverArgsPerNode[i] = param
}
params.ServerArgsPerNode = serverArgsPerNode
}
if cfg.testingKnobCfg != "" {
switch cfg.testingKnobCfg {
default:
t.Fatalf("TestingKnobCfg %s not found", cfg.testingKnobCfg)
}
}
if cfg.iodir == "" {
tc, _, cfg.iodir, cleanup = backupRestoreTestSetupWithParams(t, clusterSize, cfg.splits,
InitManualReplication, params)
} else {
tc, _, cleanup = backupRestoreTestSetupEmptyWithParams(t, clusterSize, cfg.iodir,
InitManualReplication, params)
}
cleanupFn := func() {
cleanup()
}
d.servers[cfg.name] = tc.Server(0)
d.dataDirs[cfg.name] = cfg.iodir
d.cleanupFns = append(d.cleanupFns, cleanupFn)
return nil
}
func (d *datadrivenTestState) getIODir(t *testing.T, server string) string {
dir, ok := d.dataDirs[server]
if !ok {
t.Fatalf("server %s does not exist", server)
}
return dir
}
func (d *datadrivenTestState) getSQLDB(t *testing.T, server string, user string) *gosql.DB {
key := sqlDBKey{server, user}
if db, ok := d.sqlDBs[key]; ok {
return db
}
addr := d.servers[server].ServingSQLAddr()
pgURL, cleanup := sqlutils.PGUrl(t, addr, "TestBackupRestoreDataDriven", url.User(user))
d.cleanupFns = append(d.cleanupFns, cleanup)
base, err := pq.NewConnector(pgURL.String())
if err != nil {
t.Fatal(err)
}
connector := pq.ConnectorWithNoticeHandler(base, func(notice *pq.Error) {
d.noticeBuffer = append(d.noticeBuffer, notice.Severity+": "+notice.Message)
if notice.Detail != "" {
d.noticeBuffer = append(d.noticeBuffer, "DETAIL: "+notice.Detail)
}
if notice.Hint != "" {
d.noticeBuffer = append(d.noticeBuffer, "HINT: "+notice.Hint)
}
})
d.sqlDBs[key] = gosql.OpenDB(connector)
return d.sqlDBs[key]
}
// TestDataDriven is a datadriven test to test standard backup/restore
// interactions involving setting up clusters and running different SQL
// commands. The test files are in testdata/backup-restore. The following
// syntax is provided:
//
// - "new-server name=<name> [args]"
// Create a new server with the input name.
//
// Supported arguments:
//
// - share-io-dir: can be specified to share an IO directory with an existing
// server. This is useful when restoring from a backup taken in another
// server.
//
// - allow-implicit-access: can be specified to set
// `EnableNonAdminImplicitAndArbitraryOutbound` to true
//
// - disable-http: disables use of external HTTP endpoints.
//
// - localities: specifies the localities that will be used when starting up
// the test cluster. The cluster will have len(localities) nodes, with each
// node assigned a locality config corresponding to the locality. Please
// update the `localityCfgs` map when adding new localities.
//
// - nodes: specifies the number of nodes in the test cluster.
//
// - splits: specifies the number of ranges the bank table is split into.
//
// - before-version=<beforeVersion>: creates a mixed version cluster where all
// nodes running the test server binary think the clusterVersion is one
// version before the passed in <beforeVersion> key. See cockroach_versions.go
// for possible values.
//
// - testingKnobCfg: specifies a key to a hardcoded testingKnob configuration
//
// - "upgrade-server version=<version>"
// Upgrade the cluster version of the active server to the passed in
// clusterVersion key. See cockroach_versions.go for possible values.
//
// - "exec-sql [server=<name>] [user=<name>] [args]"
// Executes the input SQL query on the target server. By default, server is
// the last created server.
//
// Supported arguments:
//
// - expect-error-regex=<regex>: expects the query to return an error with a string
// matching the provided regex
//
// - expect-error-ignore: expects the query to return an error, but we will
// ignore it.
//
// - ignore-notice: does not print out the notice that is buffered during
// query execution.
//
// - "query-sql [server=<name>] [user=<name>] [regex=<regex pattern>]"
// Executes the input SQL query and print the results.
//
// - regex: return true if the query result matches the regex pattern and
// false otherwise.
//
// - "reset"
// Clear all state associated with the test.
//
// - "job" [server=<name>] [user=<name>] [args]
// Executes job specific operations.
//
// Supported arguments:
//
// - resume=<tag>: resumes the job referenced by the tag, use in conjunction
// with wait-for-state.
//
// - cancel=<tag>: cancels the job referenced by the tag and waits for it to
// reach a CANCELED state.
//
// - wait-for-state=<succeeded|paused|failed|cancelled> tag=<tag>: wait for
// the job referenced by the tag to reach the specified state.
//
// - "let" [args]
// Assigns the returned value of the SQL query to the provided args as variables.
//
// - "save-cluster-ts" tag=<tag>
// Saves the `SELECT cluster_logical_timestamp()` with the tag. Can be used
// in the future with intstructions such as `aost`.
//
// - "backup" [args]
// Executes backup specific operations.
//
// Supported arguments:
//
// - tag=<tag>: tag the backup job to reference it in the future.
//
// - expect-pausepoint: expects the backup job to end up in a paused state because
// of a pausepoint error.
//
// - "restore" [args]
// Executes restore specific operations.
//
// Supported arguments:
//
// - tag=<tag>: tag the restore job to reference it in the future.
//
// - expect-pausepoint: expects the restore job to end up in a paused state because
// of a pausepoint error.
//
// - aost: expects a tag referencing a previously saved cluster timestamp
// using `save-cluster-ts`. It then runs the restore as of the saved cluster
// timestamp.
//
// - "schema" [args]
// Executes schema change specific operations.
//
// Supported arguments:
//
// - tag=<tag>: tag the schema change job to reference it in the future.
//
// - expect-pausepoint: expects the schema change job to end up in a paused state because
// of a pausepoint error.
//
// - "kv" [args]
// Issues a kv request
//
// Supported arguments:
//
// - type: kv request type. Currently, only DeleteRange is supported
//
// - target: SQL target. Currently, only table names are supported.
//
// - "corrupt-backup" uri=<collectionUri>
// Finds the latest backup in the provided collection uri an flips a bit in one SST in the backup
//
// - "link-backup" server=<server> src-path=<testDataPathRelative> dest-path=<fileIO path relative>
// Creates a symlink from the testdata path to the file IO path, so that we
// can restore precreated backup. src-path and dest-path are comma seperated
// paths that will be joined.
//
// - "sleep ms=TIME"
// Sleep for TIME milliseconds.
func TestDataDriven(t *testing.T) {
defer leaktest.AfterTest(t)()
defer log.Scope(t).Close(t)
skip.UnderRace(t, "takes ~3mins to run")
// This test uses this mock HTTP server to pass the backup files between tenants.
httpAddr, httpServerCleanup := makeInsecureHTTPServer(t)
defer httpServerCleanup()
ctx := context.Background()
datadriven.Walk(t, testutils.TestDataPath(t, "backup-restore"), func(t *testing.T, path string) {
var lastCreatedServer string
ds := newDatadrivenTestState()
defer ds.cleanup(ctx, t)
datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
execWithTagAndPausePoint := func(jobType jobspb.Type) string {
const user = "root"
sqlDB := ds.getSQLDB(t, lastCreatedServer, user)
// First, run the schema change.
_, err := sqlDB.Exec(d.Input)
var jobID jobspb.JobID
{
const qFmt = `SELECT job_id FROM [SHOW JOBS] WHERE job_type = '%s' ORDER BY created DESC LIMIT 1`
errJob := sqlDB.QueryRow(fmt.Sprintf(qFmt, jobType)).Scan(&jobID)
if !errors.Is(errJob, gosql.ErrNoRows) {
require.NoError(t, errJob)
}
require.NotZerof(t, jobID, "job not found for %q: %+v", d.Input, err)
}
// Tag the job.
if d.HasArg("tag") {
var jobTag string
d.ScanArgs(t, "tag", &jobTag)
if _, exists := ds.jobTags[jobTag]; exists {
t.Fatalf("failed to `tag`, job with tag %s already exists", jobTag)
}
ds.jobTags[jobTag] = jobID
}
// Check if we expect a pausepoint error.
if d.HasArg("expect-pausepoint") {
// Check if we are expecting a pausepoint error.
require.NotNilf(t, err, "expected pause point error")
require.Regexp(t, "pause point .* hit$", err.Error())
jobutils.WaitForJobToPause(t, sqlutils.MakeSQLRunner(sqlDB), jobID)
ret := append(ds.noticeBuffer, "job paused at pausepoint")
return strings.Join(ret, "\n")
}
// All other errors are bad.
require.NoError(t, err)
return ""
}
for v := range ds.vars {
d.Input = strings.Replace(d.Input, v, ds.vars[v], -1)
d.Expected = strings.Replace(d.Expected, v, ds.vars[v], -1)
}
switch d.Cmd {
case "skip":
var issue int
d.ScanArgs(t, "issue-num", &issue)
skip.WithIssue(t, issue)
return ""
case "reset":
ds.cleanup(ctx, t)
ds = newDatadrivenTestState()
return ""
case "new-server":
var name, shareDirWith, iodir, localities, beforeVersion, testingKnobCfg string
var splits int
nodes := singleNode
var io base.ExternalIODirConfig
d.ScanArgs(t, "name", &name)
if d.HasArg("share-io-dir") {
d.ScanArgs(t, "share-io-dir", &shareDirWith)
}
if shareDirWith != "" {
iodir = ds.getIODir(t, shareDirWith)
}
if d.HasArg("allow-implicit-access") {
io.EnableNonAdminImplicitAndArbitraryOutbound = true
}
if d.HasArg("disable-http") {
io.DisableHTTP = true
}
if d.HasArg("localities") {
d.ScanArgs(t, "localities", &localities)
}
if d.HasArg("nodes") {
d.ScanArgs(t, "nodes", &nodes)
}
if d.HasArg("splits") {
d.ScanArgs(t, "splits", &splits)
}
if d.HasArg("beforeVersion") {
d.ScanArgs(t, "beforeVersion", &beforeVersion)
}
if d.HasArg("testingKnobCfg") {
d.ScanArgs(t, "testingKnobCfg", &testingKnobCfg)
}
lastCreatedServer = name
cfg := serverCfg{
name: name,
iodir: iodir,
nodes: nodes,
splits: splits,
ioConf: io,
localities: localities,
beforeVersion: beforeVersion,
testingKnobCfg: testingKnobCfg,
}
err := ds.addServer(t, cfg)
if err != nil {
return err.Error()
}
return ""
case "switch-server":
var name string
d.ScanArgs(t, "name", &name)
lastCreatedServer = name
return ""
case "upgrade-server":
server := lastCreatedServer
user := "root"
var version string
if d.HasArg("version") {
d.ScanArgs(t, "version", &version)
}
key, ok := clusterVersionKeys[version]
if !ok {
t.Fatalf("clusterVersion %s does not exist in data driven global map", version)
}
clusterVersion := clusterversion.ByKey(key)
_, err := ds.getSQLDB(t, server, user).Exec("SET CLUSTER SETTING version = $1", clusterVersion.String())
require.NoError(t, err)
return ""
case "exec-sql":
server := lastCreatedServer
user := "root"
if d.HasArg("server") {
d.ScanArgs(t, "server", &server)
}
if d.HasArg("user") {
d.ScanArgs(t, "user", &user)
}
ds.noticeBuffer = nil
d.Input = strings.ReplaceAll(d.Input, "http://COCKROACH_TEST_HTTP_SERVER/", httpAddr)
_, err := ds.getSQLDB(t, server, user).Exec(d.Input)
ret := ds.noticeBuffer
if d.HasArg("ignore-notice") {
ret = nil
}
// Check if we are expecting a pausepoint error.
if d.HasArg("expect-pausepoint") {
require.NotNilf(t, err, "expected pause point error")
require.True(t, strings.Contains(err.Error(), "job requested it be paused"))
// Find job ID of the pausepoint job.
var jobID jobspb.JobID
require.NoError(t,
ds.getSQLDB(t, server, user).QueryRow(
`SELECT job_id FROM [SHOW JOBS] ORDER BY created DESC LIMIT 1`).Scan(&jobID))
fmt.Printf("expecting pausepoint, found job ID %d\n\n\n", jobID)
runner := sqlutils.MakeSQLRunner(ds.getSQLDB(t, server, user))
jobutils.WaitForJobToPause(t, runner, jobID)
ret = append(ds.noticeBuffer, "job paused at pausepoint")
ret = append(ret, "")
return strings.Join(ret, "\n")
}
// Check if we are expecting an error, and want to ignore outputting it.
if d.HasArg("expect-error-ignore") {
require.NotNilf(t, err, "expected error")
ret = append(ret, "ignoring expected error")
return strings.Join(ret, "\n")
}
// Check if we are expecting an error, and want to match it against a
// regex.
if d.HasArg("expect-error-regex") {
require.NotNilf(t, err, "expected error")
var expectErrorRegex string
d.ScanArgs(t, "expect-error-regex", &expectErrorRegex)
require.True(t,
testutils.IsError(err, expectErrorRegex),
"Regex `%s` did not match `%s`",
expectErrorRegex,
err)
ret = append(ret, "regex matches error")
return strings.Join(ret, "\n")
}
// Check for other errors.
if err != nil {
if pqErr := (*pq.Error)(nil); errors.As(err, &pqErr) {
ret = append(ds.noticeBuffer, err.Error())
if pqErr.Detail != "" {
ret = append(ret, "DETAIL: "+pqErr.Detail)
}
if pqErr.Hint != "" {
ret = append(ret, "HINT: "+pqErr.Hint)
}
} else {
t.Errorf("failed to execute stmt %s due to %s", d.Input, err.Error())
}
}
return strings.Join(ret, "\n")
case "query-sql":
server := lastCreatedServer
user := "root"
if d.HasArg("server") {
d.ScanArgs(t, "server", &server)
}
if d.HasArg("user") {
d.ScanArgs(t, "user", &user)
}
rows, err := ds.getSQLDB(t, server, user).Query(d.Input)
if err != nil {
return err.Error()
}
output, err := sqlutils.RowsToDataDrivenOutput(rows)
require.NoError(t, err)
if d.HasArg("regex") {
var pattern string
d.ScanArgs(t, "regex", &pattern)
matched, err := regexp.MatchString(pattern, output)
require.NoError(t, err)
if matched {
return "true"
}
return "false"
}
return output
case "let":
server := lastCreatedServer
user := "root"
if len(d.CmdArgs) == 0 {
t.Fatalf("Must specify at least one variable name.")
}
rows, err := ds.getSQLDB(t, server, user).Query(d.Input)
if err != nil {
return err.Error()
}
output, err := sqlutils.RowsToDataDrivenOutput(rows)
output = strings.TrimSpace(output)
values := strings.Split(output, "\n")
if len(values) != len(d.CmdArgs) {
t.Fatalf("Expecting %d vars, found %d", len(d.CmdArgs), len(values))
}
for i := range values {
key := d.CmdArgs[i].Key
if !strings.HasPrefix(key, "$") {
t.Fatalf("Vars must start with `$`.")
}
ds.vars[key] = values[i]
}
require.NoError(t, err)
return ""
case "sleep":
var msStr string
if d.HasArg("ms") {
d.ScanArgs(t, "ms", &msStr)
} else {
t.Fatalf("must specify sleep time in ms")
}
ms, err := strconv.ParseInt(msStr, 10, 64)
if err != nil {
t.Fatalf("invalid sleep time: %v", err)
}
time.Sleep(time.Duration(ms) * time.Millisecond)
return ""
case "backup":
return execWithTagAndPausePoint(jobspb.TypeBackup)
case "import":
return execWithTagAndPausePoint(jobspb.TypeImport)
case "restore":
if d.HasArg("aost") {
var aost string
d.ScanArgs(t, "aost", &aost)
var ts string
var ok bool
if ts, ok = ds.clusterTimestamps[aost]; !ok {
t.Fatalf("no cluster timestamp found for %s", aost)
}
// Replace the ts tag with the actual timestamp.
d.Input = strings.Replace(d.Input, aost,
fmt.Sprintf("'%s'", ts), 1)
}
return execWithTagAndPausePoint(jobspb.TypeRestore)
case "new-schema-change":
return execWithTagAndPausePoint(jobspb.TypeNewSchemaChange)
case "schema-change":
return execWithTagAndPausePoint(jobspb.TypeSchemaChange)
case "job":
server := lastCreatedServer
const user = "root"
if d.HasArg("cancel") {
var cancelJobTag string
d.ScanArgs(t, "cancel", &cancelJobTag)
var jobID jobspb.JobID
var ok bool
if jobID, ok = ds.jobTags[cancelJobTag]; !ok {
t.Fatalf("could not find job with tag %s", cancelJobTag)
}
runner := sqlutils.MakeSQLRunner(ds.getSQLDB(t, server, user))
runner.Exec(t, `CANCEL JOB $1`, jobID)
jobutils.WaitForJobToCancel(t, runner, jobID)
} else if d.HasArg("resume") {
var resumeJobTag string
d.ScanArgs(t, "resume", &resumeJobTag)
var jobID jobspb.JobID
var ok bool
if jobID, ok = ds.jobTags[resumeJobTag]; !ok {
t.Fatalf("could not find job with tag %s", resumeJobTag)
}
runner := sqlutils.MakeSQLRunner(ds.getSQLDB(t, server, user))
runner.Exec(t, `RESUME JOB $1`, jobID)
} else if d.HasArg("wait-for-state") {
var tag string
d.ScanArgs(t, "tag", &tag)
var jobID jobspb.JobID
var ok bool
if jobID, ok = ds.jobTags[tag]; !ok {
t.Fatalf("could not find job with tag %s", tag)
}
runner := sqlutils.MakeSQLRunner(ds.getSQLDB(t, server, user))
var state string
d.ScanArgs(t, "wait-for-state", &state)
switch state {
case "succeeded":
jobutils.WaitForJobToSucceed(t, runner, jobID)
case "cancelled":
jobutils.WaitForJobToCancel(t, runner, jobID)
case "paused":
jobutils.WaitForJobToPause(t, runner, jobID)
case "failed":
jobutils.WaitForJobToFail(t, runner, jobID)
default:
t.Fatalf("unknown state %s", state)
}
}
return ""
case "kv":
var request string
d.ScanArgs(t, "request", &request)
var target string
d.ScanArgs(t, "target", &target)
handleKVRequest(ctx, t, lastCreatedServer, ds, request, target)
return ""
case "save-cluster-ts":
server := lastCreatedServer
const user = "root"
var timestampTag string
d.ScanArgs(t, "tag", ×tampTag)
if _, ok := ds.clusterTimestamps[timestampTag]; ok {
t.Fatalf("cannot reuse cluster ts tag %s", timestampTag)
}
var ts string
err := ds.getSQLDB(t, server, user).QueryRow(`SELECT cluster_logical_timestamp()`).Scan(&ts)
require.NoError(t, err)
ds.clusterTimestamps[timestampTag] = ts
return ""
case "create-dummy-system-table":
db := ds.servers[lastCreatedServer].DB()
codec := ds.servers[lastCreatedServer].ExecutorConfig().(sql.ExecutorConfig).Codec
dummyTable := systemschema.SettingsTable
err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
id, err := ds.servers[lastCreatedServer].ExecutorConfig().(sql.ExecutorConfig).
DescIDGenerator.GenerateUniqueDescID(ctx)
if err != nil {
return err
}
mut := dummyTable.NewBuilder().BuildCreatedMutable().(*tabledesc.Mutable)
mut.ID = id
mut.Name = fmt.Sprintf("%s_%d", "crdb_internal_copy", id)
tKey := catalogkeys.EncodeNameKey(codec, mut)
b := txn.NewBatch()
b.CPut(tKey, mut.GetID(), nil)
b.CPut(catalogkeys.MakeDescMetadataKey(codec, mut.GetID()), mut.DescriptorProto(), nil)
return txn.Run(ctx, b)
})
require.NoError(t, err)
return ""
case "corrupt-backup":
server := lastCreatedServer
const user = "root"
var uri string
d.ScanArgs(t, "uri", &uri)
parsedURI, err := url.Parse(strings.Replace(uri, "'", "", -1))
require.NoError(t, err)
var filePath string
filePathQuery := fmt.Sprintf("SELECT path FROM [SHOW BACKUP FILES FROM LATEST IN %s] LIMIT 1", uri)
err = ds.getSQLDB(t, server, user).QueryRow(filePathQuery).Scan(&filePath)
require.NoError(t, err)
fullPath := filepath.Join(ds.getIODir(t, server), parsedURI.Path, filePath)
print(fullPath)
data, err := os.ReadFile(fullPath)
require.NoError(t, err)
data[20] ^= 1
if err := os.WriteFile(fullPath, data, 0644 /* perm */); err != nil {
t.Fatal(err)
}
return ""
case "link-backup":
server := lastCreatedServer
sourceRelativePath := ""
destRelativePath := ""
ioDir := ds.getIODir(t, server)
d.ScanArgs(t, "server", &server)
d.ScanArgs(t, "src-path", &sourceRelativePath)
d.ScanArgs(t, "dest-path", &destRelativePath)
splitSrcPath := strings.Split(sourceRelativePath, ",")
sourcePath, err := filepath.Abs(testutils.TestDataPath(t, splitSrcPath...))
require.NoError(t, err)
splitDestPath := strings.Split(destRelativePath, ",")
destPath := filepath.Join(ioDir, filepath.Join(splitDestPath...))
require.NoError(t, err)
require.NoError(t, os.Symlink(sourcePath, destPath))
return ""
default:
return fmt.Sprintf("unknown command: %s", d.Cmd)
}
})
})
}
func handleKVRequest(
ctx context.Context, t *testing.T, server string, ds datadrivenTestState, request, target string,
) {
user := "root"
if request == "DeleteRange" {
var tableID uint32
err := ds.getSQLDB(t, server, user).QueryRow(`SELECT id FROM system.namespace WHERE name = $1`,
target).Scan(&tableID)
require.NoError(t, err)
bankSpan := makeTableSpan(tableID)
dr := roachpb.DeleteRangeRequest{
// Bogus span to make it a valid request.
RequestHeader: roachpb.RequestHeader{
Key: bankSpan.Key,
EndKey: bankSpan.EndKey,
},
UseRangeTombstone: true,
}
if _, err := kv.SendWrapped(ctx, ds.servers[server].DistSenderI().(*kvcoord.DistSender), &dr); err != nil {
t.Fatal(err)
}
} else {
t.Fatalf("Unknown kv request")
}
}