Skip to content

Commit

Permalink
Merge pull request #21373 from danhhz/workload_fixture
Browse files Browse the repository at this point in the history
 cmd/workload: add fixtures subcommand
  • Loading branch information
danhhz authored Jan 22, 2018
2 parents 0a9bee4 + 7762d23 commit c14c3ba
Showing 10 changed files with 1,016 additions and 384 deletions.
179 changes: 179 additions & 0 deletions pkg/cmd/workload/fixtures.go
Original file line number Diff line number Diff line change
@@ -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
}
Loading

0 comments on commit c14c3ba

Please sign in to comment.