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

     1  // Copyright 2019 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 sqlsmith
    12  
    13  import (
    14  	"context"
    15  	gosql "database/sql"
    16  	"math/rand"
    17  	"strings"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/internal/sqlsmith"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    22  	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
    23  	"github.com/cockroachdb/cockroach/pkg/workload"
    24  	"github.com/cockroachdb/cockroach/pkg/workload/histogram"
    25  	"github.com/cockroachdb/errors"
    26  	"github.com/spf13/pflag"
    27  )
    28  
    29  type sqlSmith struct {
    30  	flags     workload.Flags
    31  	connFlags *workload.ConnFlags
    32  
    33  	seed          int64
    34  	tables        int
    35  	errorSettings int
    36  }
    37  
    38  type errorSettingTypes int
    39  
    40  const (
    41  	ignoreExecErrors errorSettingTypes = iota
    42  	returnOnInternalError
    43  	returnOnError
    44  )
    45  
    46  func init() {
    47  	workload.Register(sqlSmithMeta)
    48  }
    49  
    50  var sqlSmithMeta = workload.Meta{
    51  	Name:        `sqlsmith`,
    52  	Description: `sqlsmith is a random SQL query generator`,
    53  	Version:     `1.0.0`,
    54  	New: func() workload.Generator {
    55  		g := &sqlSmith{}
    56  		g.flags.FlagSet = pflag.NewFlagSet(`sqlsmith`, pflag.ContinueOnError)
    57  		g.flags.Int64Var(&g.seed, `seed`, 1, `Key hash seed.`)
    58  		g.flags.IntVar(&g.tables, `tables`, 1, `Number of tables.`)
    59  		g.flags.IntVar(&g.errorSettings, `error-sensitivity`, 0,
    60  			`SQLSmith's sensitivity to errors. 0=ignore all errors. 1=quit on internal errors. 2=quit on any error.`)
    61  		g.connFlags = workload.NewConnFlags(&g.flags)
    62  		return g
    63  	},
    64  }
    65  
    66  // Meta implements the Generator interface.
    67  func (*sqlSmith) Meta() workload.Meta { return sqlSmithMeta }
    68  
    69  // Flags implements the Flagser interface.
    70  func (g *sqlSmith) Flags() workload.Flags { return g.flags }
    71  
    72  // Tables implements the Generator interface.
    73  func (g *sqlSmith) Tables() []workload.Table {
    74  	rng := rand.New(rand.NewSource(g.seed))
    75  	var tables []workload.Table
    76  	for idx := 0; idx < g.tables; idx++ {
    77  		schema := sqlbase.RandCreateTable(rng, "table", idx)
    78  		table := workload.Table{
    79  			Name:   schema.Table.String(),
    80  			Schema: tree.Serialize(schema),
    81  		}
    82  		// workload expects the schema to be missing the CREATE TABLE "name", so
    83  		// drop everything before the first `(`.
    84  		table.Schema = table.Schema[strings.Index(table.Schema, `(`):]
    85  		tables = append(tables, table)
    86  	}
    87  	return tables
    88  }
    89  
    90  func (g *sqlSmith) handleError(err error) error {
    91  	if err != nil {
    92  		switch errorSettingTypes(g.errorSettings) {
    93  		case ignoreExecErrors:
    94  			return nil
    95  		case returnOnInternalError:
    96  			if strings.Contains(err.Error(), "internal error") {
    97  				return err
    98  			}
    99  		case returnOnError:
   100  			return err
   101  		}
   102  	}
   103  	return nil
   104  }
   105  
   106  func (g *sqlSmith) validateErrorSetting() error {
   107  	switch errorSettingTypes(g.errorSettings) {
   108  	case ignoreExecErrors:
   109  	case returnOnInternalError:
   110  	case returnOnError:
   111  	default:
   112  		return errors.Newf("invalid value for error-sensitivity: %d", g.errorSettings)
   113  	}
   114  	return nil
   115  }
   116  
   117  // Ops implements the Opser interface.
   118  func (g *sqlSmith) Ops(urls []string, reg *histogram.Registry) (workload.QueryLoad, error) {
   119  	if err := g.validateErrorSetting(); err != nil {
   120  		return workload.QueryLoad{}, err
   121  	}
   122  	sqlDatabase, err := workload.SanitizeUrls(g, g.connFlags.DBOverride, urls)
   123  	if err != nil {
   124  		return workload.QueryLoad{}, err
   125  	}
   126  	db, err := gosql.Open(`cockroach`, strings.Join(urls, ` `))
   127  	if err != nil {
   128  		return workload.QueryLoad{}, err
   129  	}
   130  	// Allow a maximum of concurrency+1 connections to the database.
   131  	db.SetMaxOpenConns(g.connFlags.Concurrency + 1)
   132  	db.SetMaxIdleConns(g.connFlags.Concurrency + 1)
   133  
   134  	ql := workload.QueryLoad{SQLDatabase: sqlDatabase}
   135  	for i := 0; i < g.connFlags.Concurrency; i++ {
   136  		rng := rand.New(rand.NewSource(g.seed + int64(i)))
   137  		smither, err := sqlsmith.NewSmither(db, rng)
   138  		if err != nil {
   139  			return workload.QueryLoad{}, err
   140  		}
   141  
   142  		hists := reg.GetHandle()
   143  		workerFn := func(ctx context.Context) error {
   144  			start := timeutil.Now()
   145  			query := smither.Generate()
   146  			elapsed := timeutil.Since(start)
   147  			hists.Get(`generate`).Record(elapsed)
   148  
   149  			start = timeutil.Now()
   150  			_, err := db.ExecContext(ctx, query)
   151  			if handledErr := g.handleError(err); handledErr != nil {
   152  				return handledErr
   153  			}
   154  			elapsed = timeutil.Since(start)
   155  
   156  			hists.Get(`exec`).Record(elapsed)
   157  
   158  			return nil
   159  		}
   160  		ql.WorkerFns = append(ql.WorkerFns, workerFn)
   161  	}
   162  	return ql, nil
   163  }