diff --git a/pkg/cmd/workload/fixtures.go b/pkg/cmd/workload/fixtures.go new file mode 100644 index 000000000000..3760e9c3dfc9 --- /dev/null +++ b/pkg/cmd/workload/fixtures.go @@ -0,0 +1,179 @@ +// Copyright 2018 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. See the AUTHORS file +// for names of contributors. + +package main + +import ( + "context" + gosql "database/sql" + "fmt" + + "cloud.google.com/go/storage" + "github.com/pkg/errors" + "github.com/spf13/cobra" + "google.golang.org/api/option" + + "github.com/cockroachdb/cockroach/pkg/testutils/workload" + _ "github.com/cockroachdb/cockroach/pkg/testutils/workload/all" + "github.com/cockroachdb/cockroach/pkg/util/log" +) + +var useast1bFixtures = workload.FixtureStore{ + // TODO(dan): Keep fixtures in more than one region to better support + // geo-distributed clusters. + GCSBucket: `cockroach-fixtures`, + GCSPrefix: `workload`, +} + +var fixturesCmd = &cobra.Command{Use: `fixtures`} +var fixturesListCmd = &cobra.Command{ + Use: `list`, + Short: `List all fixtures stored on GCS`, + RunE: fixturesList, +} +var fixturesStoreCmd = &cobra.Command{ + Use: `store`, + Short: `Regenerate and store a fixture on GCS`, +} +var fixturesLoadCmd = &cobra.Command{ + Use: `load`, + Short: `Load a fixture into a running cluster. ` + + `An enterprise license is required.`, +} + +var fixturesLoadDB = fixturesLoadCmd.PersistentFlags().String( + `into_db`, `workload`, `SQL database to load fixture into`) + +const storageError = `failed to create google cloud client ` + + `(You may need to setup the GCS application default credentials: ` + + `'gcloud auth application-default login --project=cockroach-shared')` + +// getStorage returns a GCS client using "application default" credentials. The +// caller is responsible for closing it. +func getStorage(ctx context.Context) (*storage.Client, error) { + // TODO(dan): Right now, we don't need all the complexity of + // storageccl.ExportStorage, but if we start supporting more than just GCS, + // this should probably be switched to it. + g, err := storage.NewClient(ctx, option.WithScopes(storage.ScopeReadWrite)) + return g, errors.Wrap(err, storageError) +} + +func init() { + for _, meta := range workload.Registered() { + gen := meta.New() + genFlags := gen.Flags() + + genStoreCmd := &cobra.Command{ + Use: meta.Name + ` [CRDB URI]`, + Args: cobra.RangeArgs(0, 1), + } + genStoreCmd.Flags().AddFlagSet(genFlags) + genStoreCmd.RunE = func(cmd *cobra.Command, args []string) error { + crdb := crdbDefaultURI + if len(args) > 0 { + crdb = args[0] + } + return fixturesStore(cmd, gen, crdb) + } + fixturesStoreCmd.AddCommand(genStoreCmd) + + genLoadCmd := &cobra.Command{ + Use: meta.Name + ` [CRDB URI]`, + Args: cobra.RangeArgs(0, 1), + } + genLoadCmd.Flags().AddFlagSet(genFlags) + genLoadCmd.RunE = func(cmd *cobra.Command, args []string) error { + crdb := crdbDefaultURI + if len(args) > 0 { + crdb = args[0] + } + return fixturesLoad(cmd, gen, crdb) + } + fixturesLoadCmd.AddCommand(genLoadCmd) + } + fixturesCmd.AddCommand(fixturesListCmd) + fixturesCmd.AddCommand(fixturesStoreCmd) + fixturesCmd.AddCommand(fixturesLoadCmd) + rootCmd.AddCommand(fixturesCmd) +} + +func fixturesList(cmd *cobra.Command, _ []string) error { + ctx := context.Background() + gcs, err := getStorage(ctx) + if err != nil { + return err + } + defer func() { _ = gcs.Close() }() + fixtures, err := workload.ListFixtures(ctx, gcs, useast1bFixtures) + if err != nil { + return err + } + for _, fixture := range fixtures { + fmt.Println(fixture) + } + return nil +} + +func fixturesStore(cmd *cobra.Command, gen workload.Generator, crdbURI string) error { + ctx := context.Background() + gcs, err := getStorage(ctx) + if err != nil { + return err + } + defer func() { _ = gcs.Close() }() + + sqlDB, err := gosql.Open(`postgres`, crdbURI) + if err != nil { + return err + } + + fixture, err := workload.MakeFixture(ctx, sqlDB, gcs, useast1bFixtures, gen) + if err != nil { + return err + } + for _, table := range fixture.Tables { + log.Infof(ctx, `stored %s`, table.BackupURI) + } + return nil +} + +func fixturesLoad(cmd *cobra.Command, gen workload.Generator, crdbURI string) error { + ctx := context.Background() + gcs, err := getStorage(ctx) + if err != nil { + return err + } + defer func() { _ = gcs.Close() }() + + sqlDB, err := gosql.Open(`postgres`, crdbURI) + if err != nil { + return err + } + if _, err := sqlDB.Exec(`CREATE DATABASE IF NOT EXISTS ` + *fixturesLoadDB); err != nil { + return err + } + + fixture, err := workload.GetFixture(ctx, gcs, useast1bFixtures, gen) + if err != nil { + return errors.Wrap(err, `finding fixture`) + } + if err := workload.RestoreFixture(ctx, sqlDB, fixture, *fixturesLoadDB); err != nil { + return errors.Wrap(err, `restoring fixture`) + } + for _, table := range fixture.Tables { + log.Infof(ctx, `loaded %s`, table.BackupURI) + } + return nil +} diff --git a/pkg/cmd/workload/main.go b/pkg/cmd/workload/main.go new file mode 100644 index 000000000000..2851f45f2455 --- /dev/null +++ b/pkg/cmd/workload/main.go @@ -0,0 +1,30 @@ +// Copyright 2018 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. See the AUTHORS file +// for names of contributors. + +package main + +import ( + "github.com/spf13/cobra" + + _ "github.com/cockroachdb/cockroach/pkg/testutils/workload/all" +) + +var rootCmd = &cobra.Command{ + Use: `workload`, +} + +func main() { + _ = rootCmd.Execute() +} diff --git a/pkg/cmd/workload/run.go b/pkg/cmd/workload/run.go index 639542e969a8..26e03c77a846 100644 --- a/pkg/cmd/workload/run.go +++ b/pkg/cmd/workload/run.go @@ -45,8 +45,11 @@ import ( "github.com/cockroachdb/cockroach/pkg/util/timeutil" ) -var rootCmd = &cobra.Command{ - Use: `workload`, +const crdbDefaultURI = `postgres://root@localhost:26257/test?sslmode=disable` + +var runCmd = &cobra.Command{ + Use: `run`, + Short: `Run a workload's operations against a cluster`, } var rootFlags = pflag.NewFlagSet(`workload`, pflag.ContinueOnError) @@ -65,19 +68,9 @@ var histFile = rootFlags.String( "hist-file", "", "Write histogram data to file for HdrHistogram Plotter, or stdout if - is specified.") -func main() { - _ = rootCmd.Execute() -} - func init() { for _, meta := range workload.Registered() { gen := meta.New() - // Avoid use of `genFn` in the closures below, because it is a) not needed and - // b) prevents regression of a bug in which a missing loop-local copy in - // the RunE closure below caused the last registered command to be the one - // run, whichever generator was specified. - meta.New = nil - genFlags := gen.Flags() genCmd := &cobra.Command{Use: meta.Name, Short: meta.Description} @@ -99,11 +92,12 @@ func init() { if err := gen.Configure(flags); err != nil { return err } - return runRoot(gen, args) + return runRun(gen, args) } - rootCmd.AddCommand(genCmd) + runCmd.AddCommand(genCmd) } + rootCmd.AddCommand(runCmd) } // numOps keeps a global count of successful operations. @@ -220,10 +214,10 @@ func setupDatabase(dbURLs []string) (*gosql.DB, error) { } } -func runRoot(gen workload.Generator, args []string) error { +func runRun(gen workload.Generator, args []string) error { ctx := context.Background() - dbURLs := []string{"postgres://root@localhost:26257/test?sslmode=disable"} + dbURLs := []string{crdbDefaultURI} if len(args) >= 1 { dbURLs = args } diff --git a/pkg/testutils/lint/lint_test.go b/pkg/testutils/lint/lint_test.go index 0e138806bf5d..d89363c7f5d4 100644 --- a/pkg/testutils/lint/lint_test.go +++ b/pkg/testutils/lint/lint_test.go @@ -157,6 +157,7 @@ func TestLint(t *testing.T) { ":!ccl/storageccl/export_storage_test.go", ":!cmd", ":!testutils/lint", + ":!testutils/workload/testonlyccl/fixture_test.go", ":!util/envutil/env.go", ":!util/log/clog.go", ":!util/color/color.go", diff --git a/pkg/testutils/workload/bank/bank.go b/pkg/testutils/workload/bank/bank.go index ba3ae29e1bd7..25a92b12b128 100644 --- a/pkg/testutils/workload/bank/bank.go +++ b/pkg/testutils/workload/bank/bank.go @@ -59,6 +59,7 @@ func init() { var bankMeta = workload.Meta{ Name: `bank`, Description: `Bank models a set of accounts with currency balances`, + Version: `1.0.0`, New: func() workload.Generator { g := &bank{flags: pflag.NewFlagSet(`kv`, pflag.ContinueOnError)} g.flags.Int64Var(&g.seed, `seed`, 1, `Key hash seed.`) diff --git a/pkg/testutils/workload/fixture.go b/pkg/testutils/workload/fixture.go new file mode 100644 index 000000000000..43d0b7d342b1 --- /dev/null +++ b/pkg/testutils/workload/fixture.go @@ -0,0 +1,242 @@ +// Copyright 2018 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. + +package workload + +import ( + "bytes" + "context" + gosql "database/sql" + "encoding/csv" + "fmt" + "net/url" + "path/filepath" + + "cloud.google.com/go/storage" + "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/util/retry" + "github.com/pkg/errors" + "github.com/spf13/pflag" + "google.golang.org/api/iterator" +) + +const ( + fixtureGCSURIScheme = `gs` +) + +// FixtureStore describes a storage place for fixtures. +type FixtureStore struct { + // GCSBucket is a Google Cloud Storage bucket. + GCSBucket string + + // GCSPrefix is a prefix to prepend to each Google Cloud Storage object + // path. + GCSPrefix string +} + +func (s FixtureStore) objectPathToURI(folder string) string { + return (&url.URL{ + Scheme: fixtureGCSURIScheme, + Host: s.GCSBucket, + Path: folder, + }).String() +} + +// Fixture describes pre-computed data for a Generator, allowing quick +// initialization of large clusters. +type Fixture struct { + Store FixtureStore + Generator Generator + Tables []FixtureTable +} + +// FixtureTable describes pre-computed data for a single table in a Generator, +// allowing quick initializaiton of large clusters. +type FixtureTable struct { + TableName string + BackupURI string +} + +// serializeOptions deterministically represents the configuration of a +// Generator as a string. +func serializeOptions(gen Generator) string { + // NB: VisitAll visits in a deterministic (alphabetical) order. + var buf bytes.Buffer + gen.Flags().VisitAll(func(f *pflag.Flag) { + if buf.Len() > 0 { + buf.WriteString(`,`) + } + fmt.Fprintf(&buf, `%s=%s`, url.PathEscape(f.Name), url.PathEscape(f.Value.String())) + }) + return buf.String() +} + +func generatorToGCSFolder(store FixtureStore, gen Generator) string { + meta := gen.Meta() + return filepath.Join( + store.GCSPrefix, + meta.Name, + fmt.Sprintf(`version=%s,%s`, meta.Version, serializeOptions(gen)), + ) +} + +// GetFixture returns a handle for pre-computed Generator data stored on GCS. It +// is expected that the generator will have had Configure called on it. +func GetFixture( + ctx context.Context, gcs *storage.Client, store FixtureStore, gen Generator, +) (Fixture, error) { + b := gcs.Bucket(store.GCSBucket) + + fixtureFolder := generatorToGCSFolder(store, gen) + _, err := b.Objects(ctx, &storage.Query{Prefix: fixtureFolder, Delimiter: `/`}).Next() + if err == iterator.Done { + return Fixture{}, errors.Errorf(`fixture not found: %s`, fixtureFolder) + } else if err != nil { + return Fixture{}, err + } + + fixture := Fixture{Store: store, Generator: gen} + for _, table := range gen.Tables() { + tableFolder := filepath.Join(fixtureFolder, table.Name) + _, err := b.Objects(ctx, &storage.Query{Prefix: tableFolder, Delimiter: `/`}).Next() + if err == iterator.Done { + return Fixture{}, errors.Errorf(`fixture table not found: %s`, tableFolder) + } else if err != nil { + return Fixture{}, err + } + fixture.Tables = append(fixture.Tables, FixtureTable{ + TableName: table.Name, + BackupURI: store.objectPathToURI(tableFolder), + }) + } + return fixture, nil +} + +// writeCSVs creates a file on GCS in the specified folder that contains the +// data for the given table. The GCS object path to the written file is +// returned. +func writeCSVs( + ctx context.Context, gcs *storage.Client, table Table, store FixtureStore, folder string, +) (string, error) { + // TODO(dan): For large tables, break this up into multiple CSVs. + csvPath := filepath.Join(folder, table.Name+`.csv`) + fmt.Println(csvPath) + const maxAttempts = 3 + err := retry.WithMaxAttempts(ctx, base.DefaultRetryOptions(), maxAttempts, func() error { + w := gcs.Bucket(store.GCSBucket).Object(csvPath).NewWriter(ctx) + csvW := csv.NewWriter(w) + for rowIdx := 0; rowIdx < table.InitialRowCount; rowIdx++ { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + if err := csvW.Write(table.InitialRowFn(rowIdx)); err != nil { + return err + } + } + csvW.Flush() + if err := csvW.Error(); err != nil { + return err + } + return w.Close() + }) + return csvPath, err +} + +// MakeFixture regenerates a fixture, storing it to GCS. It is expected that the +// generator will have had Configure called on it. +// +// There's some ideal world in which we can generate backups (and thus fixtures) +// directly from a Generator, but for now, we use `IMPORT ... CSV DATA`. First a +// CSV file with the table data is written to GCS. `IMPORT ... CSV DATA` works +// by turning a set of CSV files for a single table into a backup (in a cloud +// storage location specified by the `temp=` option), then restoring that backup +// into a cluster. The `transform_only` option gives us only the first half +// (which is all we want for fixture generation). +func MakeFixture( + ctx context.Context, sqlDB *gosql.DB, gcs *storage.Client, store FixtureStore, gen Generator, +) (Fixture, error) { + fixtureFolder := generatorToGCSFolder(store, gen) + + // TODO(dan): Break up large tables into many CSV files and use the + // `distributed` option of `IMPORT` to do the backup generation work in + // parallel. Afterward, experiment with parallelizing the per-table work + // (IMPORT is pretty good at using cluster resources, so this is not as + // obvious of a win as it might seem). + for _, table := range gen.Tables() { + tableFolder := filepath.Join(fixtureFolder, table.Name) + csvPath, err := writeCSVs(ctx, gcs, table, store, tableFolder) + if err != nil { + return Fixture{}, err + } + backupURI := store.objectPathToURI(tableFolder) + csvURI := store.objectPathToURI(csvPath) + + if _, err := sqlDB.ExecContext( + ctx, `SET CLUSTER SETTING experimental.importcsv.enabled = true`, + ); err != nil { + return Fixture{}, err + } + importStmt := fmt.Sprintf( + `IMPORT TABLE %s %s CSV DATA ($1) WITH transform_only, temp=$2`, + table.Name, table.Schema, + ) + if _, err := sqlDB.ExecContext(ctx, importStmt, csvURI, backupURI); err != nil { + return Fixture{}, errors.Wrapf(err, `creating backup for table %s`, table.Name) + } + } + // TODO(dan): Clean up the CSVs. + return GetFixture(ctx, gcs, store, gen) +} + +// RestoreFixture loads a fixture into a CockroachDB cluster. An enterprise +// license is required to have been set in the cluster. +func RestoreFixture(ctx context.Context, sqlDB *gosql.DB, fixture Fixture, database string) error { + for _, table := range fixture.Tables { + // The IMPORT ... CSV DATA command generates a backup with the table in + // database `csv`. + importStmt := fmt.Sprintf(`RESTORE csv.%s FROM $1 WITH into_db=$2`, table.TableName) + if _, err := sqlDB.ExecContext(ctx, importStmt, table.BackupURI, database); err != nil { + return err + } + } + return nil +} + +// ListFixtures returns the object paths to all fixtures stored in a FixtureStore. +func ListFixtures(ctx context.Context, gcs *storage.Client, store FixtureStore) ([]string, error) { + b := gcs.Bucket(store.GCSBucket) + + var fixtures []string + gensPrefix := store.GCSPrefix + `/` + for genIter := b.Objects(ctx, &storage.Query{Prefix: gensPrefix, Delimiter: `/`}); ; { + gen, err := genIter.Next() + if err == iterator.Done { + break + } else if err != nil { + return nil, err + } + for genConfigIter := b.Objects(ctx, &storage.Query{Prefix: gen.Prefix, Delimiter: `/`}); ; { + genConfig, err := genConfigIter.Next() + if err == iterator.Done { + break + } else if err != nil { + return nil, err + } + fixtures = append(fixtures, genConfig.Prefix) + } + } + return fixtures, nil +} diff --git a/pkg/testutils/workload/kv/kv.go b/pkg/testutils/workload/kv/kv.go index d10434d7c004..4a553f565a68 100644 --- a/pkg/testutils/workload/kv/kv.go +++ b/pkg/testutils/workload/kv/kv.go @@ -56,6 +56,7 @@ var kvMeta = workload.Meta{ Name: `kv`, Description: `KV reads and writes to keys spread (by default, uniformly` + ` at random) across the cluster`, + Version: `1.0.0`, New: func() workload.Generator { g := &kv{flags: pflag.NewFlagSet(`kv`, pflag.ContinueOnError)} g.flags.IntVar(&g.batchSize, `batch`, 1, `Number of blocks to insert in a single SQL statement`) diff --git a/pkg/testutils/workload/testonlyccl/fixture_test.go b/pkg/testutils/workload/testonlyccl/fixture_test.go new file mode 100644 index 000000000000..929a703b2ffe --- /dev/null +++ b/pkg/testutils/workload/testonlyccl/fixture_test.go @@ -0,0 +1,138 @@ +// Copyright 2018 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. + +package testonlyccl + +import ( + "context" + "fmt" + "net/http" + "os" + "strconv" + "strings" + "testing" + + "cloud.google.com/go/storage" + _ "github.com/cockroachdb/cockroach/pkg/ccl/sqlccl" + "github.com/cockroachdb/cockroach/pkg/testutils" + "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" + "github.com/cockroachdb/cockroach/pkg/testutils/sqlutils" + "github.com/cockroachdb/cockroach/pkg/testutils/workload" + "github.com/spf13/pflag" + "golang.org/x/oauth2/google" + "google.golang.org/api/option" + + "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/cockroachdb/cockroach/pkg/util/timeutil" +) + +const fixtureTestGenRows = 10 + +type fixtureTestGen struct { + flags *pflag.FlagSet + val string +} + +func makeTestWorkload() workload.Generator { + g := &fixtureTestGen{flags: pflag.NewFlagSet(`fx`, pflag.ContinueOnError)} + g.flags.StringVar(&g.val, `val`, `default`, `The value for each row`) + return g +} + +func (fixtureTestGen) Meta() workload.Meta { return workload.Meta{Name: `fixture`} } +func (g fixtureTestGen) Flags() *pflag.FlagSet { return g.flags } +func (g fixtureTestGen) Configure(flags []string) error { return g.flags.Parse(flags) } +func (fixtureTestGen) Ops() []workload.Operation { return nil } +func (g fixtureTestGen) Tables() []workload.Table { + return []workload.Table{{ + Name: `fx`, + Schema: `(key INT PRIMARY KEY, value INT)`, + InitialRowCount: fixtureTestGenRows, + InitialRowFn: func(rowIdx int) []string { return []string{strconv.Itoa(rowIdx), g.val} }, + }} +} + +func TestFixture(t *testing.T) { + defer leaktest.AfterTest(t)() + ctx := context.Background() + + gcsBucket := os.Getenv(`GS_BUCKET`) + gcsKey := os.Getenv(`GS_JSONKEY`) + if gcsBucket == "" || gcsKey == "" { + t.Skip("GS_BUCKET and GS_JSONKEY env vars must be set") + } + + // This prevents leaking an http conn goroutine. + http.DefaultTransport.(*http.Transport).DisableKeepAlives = true + source, err := google.JWTConfigFromJSON([]byte(gcsKey), storage.ScopeReadWrite) + if err != nil { + t.Fatalf(`%+v`, err) + } + gcs, err := storage.NewClient(ctx, + option.WithScopes(storage.ScopeReadWrite), + option.WithTokenSource(source.TokenSource(ctx))) + if err != nil { + t.Fatalf(`%+v`, err) + } + defer func() { _ = gcs.Close() }() + + s, db, _ := serverutils.StartServer(t, base.TestServerArgs{}) + defer s.Stopper().Stop(ctx) + sqlDB := sqlutils.MakeSQLRunner(db) + sqlDB.Exec(t, `SET CLUSTER SETTING cloudstorage.gs.default.key = $1`, gcsKey) + + gen := makeTestWorkload() + flag := fmt.Sprintf(`val=%d`, timeutil.Now().UnixNano()) + if err := gen.Configure([]string{`--` + flag}); err != nil { + t.Fatalf(`%+v`, err) + } + + store := workload.FixtureStore{ + GCSBucket: gcsBucket, + GCSPrefix: fmt.Sprintf(`TestFixture-%d`, timeutil.Now().UnixNano()), + } + + if _, err := workload.GetFixture(ctx, gcs, store, gen); !testutils.IsError(err, `fixture not found`) { + t.Fatalf(`expected "fixture not found" error but got: %+v`, err) + } + + fixtures, err := workload.ListFixtures(ctx, gcs, store) + if err != nil { + t.Fatalf(`%+v`, err) + } + if len(fixtures) != 0 { + t.Errorf(`expected no fixtures but got: %+v`, fixtures) + } + + fixture, err := workload.MakeFixture(ctx, sqlDB.DB, gcs, store, gen) + if err != nil { + t.Fatalf(`%+v`, err) + } + + fixtures, err = workload.ListFixtures(ctx, gcs, store) + if err != nil { + t.Fatalf(`%+v`, err) + } + if len(fixtures) != 1 || !strings.Contains(fixtures[0], flag) { + t.Errorf(`expected exactly one %s fixture but got: %+v`, flag, fixtures) + } + + sqlDB.Exec(t, `CREATE DATABASE test`) + if err := workload.RestoreFixture(ctx, sqlDB.DB, fixture, `test`); err != nil { + t.Fatalf(`%+v`, err) + } + sqlDB.CheckQueryResults(t, + `SELECT COUNT(*) FROM test.fx`, [][]string{{strconv.Itoa(fixtureTestGenRows)}}) +} diff --git a/pkg/testutils/workload/testonlyccl/main_test.go b/pkg/testutils/workload/testonlyccl/main_test.go new file mode 100644 index 000000000000..2258ce04537c --- /dev/null +++ b/pkg/testutils/workload/testonlyccl/main_test.go @@ -0,0 +1,42 @@ +// Copyright 2018 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. + +// Package testonlyccl holds TESTS ONLY that require ccl code to pass. +package testonlyccl + +import ( + "os" + "testing" + + "github.com/cockroachdb/cockroach/pkg/security" + "github.com/cockroachdb/cockroach/pkg/security/securitytest" + "github.com/cockroachdb/cockroach/pkg/server" + "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" + "github.com/cockroachdb/cockroach/pkg/testutils/testcluster" + "github.com/cockroachdb/cockroach/pkg/util/randutil" + + _ "github.com/cockroachdb/cockroach/pkg/ccl" // ccl init hooks + "github.com/cockroachdb/cockroach/pkg/ccl/utilccl" +) + +func TestMain(m *testing.M) { + defer utilccl.TestingEnableEnterprise()() + security.SetAssetLoader(securitytest.EmbeddedAssets) + randutil.SeedForTests() + serverutils.InitTestServerFactory(server.TestServerFactory) + serverutils.InitTestClusterFactory(testcluster.TestClusterFactory) + os.Exit(m.Run()) +} + +//go:generate ../../../util/leaktest/add-leaktest.sh *_test.go diff --git a/pkg/testutils/workload/workload.go b/pkg/testutils/workload/workload.go index a9fa6d4b94d7..2a6cd3de53d2 100644 --- a/pkg/testutils/workload/workload.go +++ b/pkg/testutils/workload/workload.go @@ -66,7 +66,11 @@ type Meta struct { Name string // Description is a short description of this generator. Description string - New func() Generator + // Version is a semantic version for this generator. It should be bumped + // whenever InitialRowFn or InitialRowCount change for any of the tables. + Version string + // New returns an unconfigured instance of this generator. + New func() Generator } // Table represents a single table in a Generator. Included is a name, schema,