github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/cli/check.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package cli
    12  
    13  import (
    14  	"context"
    15  	gosql "database/sql"
    16  	"strings"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/workload"
    19  	"github.com/cockroachdb/errors"
    20  	"github.com/spf13/cobra"
    21  	"github.com/spf13/pflag"
    22  )
    23  
    24  func init() {
    25  	AddSubCmd(func(userFacing bool) *cobra.Command {
    26  		var checkCmd = SetCmdDefaults(&cobra.Command{
    27  			Use:   `check`,
    28  			Short: `check a running cluster's data for consistency`,
    29  		})
    30  		for _, meta := range workload.Registered() {
    31  			gen := meta.New()
    32  			if hooks, ok := gen.(workload.Hookser); !ok || hooks.Hooks().CheckConsistency == nil {
    33  				continue
    34  			}
    35  
    36  			var genFlags *pflag.FlagSet
    37  			if f, ok := gen.(workload.Flagser); ok {
    38  				genFlags = f.Flags().FlagSet
    39  				// Hide irrelevant flags so they don't clutter up the help text, but
    40  				// don't remove them entirely so if someone switches from
    41  				// `./workload run` to `./workload check` they don't have to remove
    42  				// them from the invocation.
    43  				for flagName, meta := range f.Flags().Meta {
    44  					if meta.RuntimeOnly && !meta.CheckConsistencyOnly {
    45  						_ = genFlags.MarkHidden(flagName)
    46  					}
    47  				}
    48  			}
    49  
    50  			genCheckCmd := SetCmdDefaults(&cobra.Command{
    51  				Use:  meta.Name + ` [CRDB URI]`,
    52  				Args: cobra.RangeArgs(0, 1),
    53  			})
    54  			genCheckCmd.Flags().AddFlagSet(genFlags)
    55  			genCheckCmd.Run = CmdHelper(gen, check)
    56  			checkCmd.AddCommand(genCheckCmd)
    57  		}
    58  		return checkCmd
    59  	})
    60  }
    61  
    62  func check(gen workload.Generator, urls []string, dbName string) error {
    63  	ctx := context.Background()
    64  
    65  	var fn func(context.Context, *gosql.DB) error
    66  	if hooks, ok := gen.(workload.Hookser); ok {
    67  		fn = hooks.Hooks().CheckConsistency
    68  	}
    69  	if fn == nil {
    70  		return errors.Errorf(`no consistency checks are defined for %s`, gen.Meta().Name)
    71  	}
    72  
    73  	sqlDB, err := gosql.Open(`cockroach`, strings.Join(urls, ` `))
    74  	if err != nil {
    75  		return err
    76  	}
    77  	defer sqlDB.Close()
    78  	if err := sqlDB.Ping(); err != nil {
    79  		return err
    80  	}
    81  	return fn(ctx, sqlDB)
    82  }