cmd/workload: add fixtures subcommand

This commit adds tooling for a central repository of precomputed test
fixtures corresponding to different configurations of
`workload.Generator`s.

To list all available fixtures:

    workload fixtures list

To load a fixture into a CockroachDB cluster:

    workload fixtures load <generator> <generator flags> [crdb url]

To compute and store a fixture (the cluster is used to distribute work):

    workload fixtures load --into_db=<db> <generator> <generator flags> [crdb url]

`workload <generator>` has been moved to `workload run <generator>` so
we can have subcommands on `workload`.

Still TODO is a better integration with `workload run`, which currently
always recreates initial data (and does this via the much slower
workload.Setup). There is also lots of opportunity for parallelism in
the fixture loading and storing, but I wanted to get the basics working
first.

Release note: None

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
+}

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()
+}

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
 	}

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",

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.`)