From 8769394bafd04c18a997ea6b84f34366f42e4a2b Mon Sep 17 00:00:00 2001 From: Michael Butler Date: Wed, 19 Jan 2022 19:27:00 -0500 Subject: [PATCH] randgen: add PopulateRandTable PopulateRandTable populates the caller's table with random data. This helper function aims to make it easier for engineers to develop randomized tests that leverage randgen / sqlsmith. Informs #72345 Release note: None --- pkg/internal/sqlsmith/setup.go | 6 +-- pkg/sql/randgen/BUILD.bazel | 2 +- pkg/sql/randgen/schema.go | 98 ++++++++++++++++++++++++++++++++++ pkg/sql/randgen/schema_test.go | 78 +++++++++++++++++++++++++++ 4 files changed, 179 insertions(+), 5 deletions(-) create mode 100644 pkg/sql/randgen/schema_test.go diff --git a/pkg/internal/sqlsmith/setup.go b/pkg/internal/sqlsmith/setup.go index 479b5ab38dab..42040b429e11 100644 --- a/pkg/internal/sqlsmith/setup.go +++ b/pkg/internal/sqlsmith/setup.go @@ -84,7 +84,8 @@ func randTablesN(r *rand.Rand, n int) string { `) // Create the random tables. - stmts := randgen.RandCreateTables(r, "table", n, + tablePrefix := "table" + stmts := randgen.RandCreateTables(r, tablePrefix, n, randgen.StatisticsMutator, randgen.PartialIndexMutator, randgen.ForeignKeyMutator, @@ -103,9 +104,6 @@ func randTablesN(r *rand.Rand, n int) string { sb.WriteString(stmt.String()) sb.WriteString(";\n") } - - // TODO(mjibson): add random INSERTs. - return sb.String() } diff --git a/pkg/sql/randgen/BUILD.bazel b/pkg/sql/randgen/BUILD.bazel index 9ac9f8484fb0..5dc63f5ffaa3 100644 --- a/pkg/sql/randgen/BUILD.bazel +++ b/pkg/sql/randgen/BUILD.bazel @@ -49,7 +49,7 @@ go_library( go_test( name = "randgen_test", - srcs = ["mutator_test.go"], + srcs = ["mutator_test.go","schema_test.go"], embed = [":randgen"], deps = ["//pkg/util/randutil"], ) diff --git a/pkg/sql/randgen/schema.go b/pkg/sql/randgen/schema.go index a5c73399cc49..db09f9e37a30 100644 --- a/pkg/sql/randgen/schema.go +++ b/pkg/sql/randgen/schema.go @@ -20,6 +20,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo" + "github.com/cockroachdb/cockroach/pkg/sql/parser" "github.com/cockroachdb/cockroach/pkg/sql/rowenc" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/sql/types" @@ -186,6 +187,103 @@ func RandCreateTableWithColumnIndexNumberGenerator( return res[0].(*tree.CreateTable) } +func parseCreateStatement(createStmtSQL string) (*tree.CreateTable, error) { + var p parser.Parser + stmts, err := p.Parse(createStmtSQL) + if err != nil { + return nil, err + } + if len(stmts) != 1 { + return nil, fmt.Errorf("parsed CreateStatement string yielded more than one parsed statment") + } + + tableStmt, ok := stmts[0].AST.(*tree.CreateTable) + if !ok { + return nil, fmt.Errorf("AST could not be cast to *tree.CreateTable") + } + return tableStmt, nil +} + +// PopulateRandTable populates the provided table with `numrows` rows of random data. +func PopulateRandTable(rng *rand.Rand, db *gosql.DB, tableName string, numRows int) error { + var ignored, createStmtSQL string + res := db.QueryRow(fmt.Sprintf("SHOW CREATE TABLE %s", tableName)) + err := res.Scan(&ignored, &createStmtSQL) + if err != nil { + return errors.Wrapf(err, "table does not exist in db") + } + createStmt, err := parseCreateStatement(createStmtSQL) + if err != nil { + return errors.Wrapf(err, "unable to parse table's create statment") + } + + defs := createStmt.Defs + colTypes := make([]*types.T, 0) + nullable := make([]bool, 0) + var colNameBuilder strings.Builder + comma := "" + for _, def := range defs { + if col, ok := def.(*tree.ColumnTableDef); ok { + if !col.Computed.Computed { + colTypes = append(colTypes, col.Type.(*types.T)) + if col.Nullable.Nullability == tree.Null { + nullable = append(nullable, true) + } else { + nullable = append(nullable, false) + } + colNameBuilder.WriteString(comma) + colNameBuilder.WriteString(col.Name.String()) + comma = ", " + } + } + } + // TODO (Butler): deal with user defined types + var ( + success int + fail int + ) + maxTries := numRows * 1000 + for success < numRows { + var valBuilder strings.Builder + valBuilder.WriteString("(") + comma := "" + for j := 0; j < len(colTypes); j++ { + valBuilder.WriteString(comma) + d := RandDatum(rng, colTypes[j], nullable[j]) + + if colTypes[j].Family() == types.OidFamily { + d = tree.NewDOid(tree.DInt(rand.Intn(2))) + } + /* + RandDatum generates bad int2's + if colTypes[j] == types.Int2 { + d = tree.NewDInt(2) + } + */ + valBuilder.WriteString(tree.AsStringWithFlags(d, tree.FmtParsable)) + comma = ", " + } + valBuilder.WriteString(")") + insertStmt := fmt.Sprintf("INSERT INTO %s (%s) VALUES %s;", + tableName, + colNameBuilder.String(), + valBuilder.String()) + _, err := db.Exec(insertStmt) + if err != nil { + // inserting into an arbitrary table can be finicky, so allow some room for error! + fail++ + if fail > maxTries { + return fmt.Errorf(`Couldn't populate table with %d insert attempts. + This could mean this function or RandDatum couldn't handle this table's schema. + Consider filing a bug. Schema: %s`, maxTries,createStmt.String()) + } + } else { + success++ + } + } + return nil +} + // GenerateRandInterestingTable takes a gosql.DB connection and creates // a table with all the types in randInterestingDatums and rows of the // interesting datums. diff --git a/pkg/sql/randgen/schema_test.go b/pkg/sql/randgen/schema_test.go new file mode 100644 index 000000000000..5e5afccbc876 --- /dev/null +++ b/pkg/sql/randgen/schema_test.go @@ -0,0 +1,78 @@ +// Copyright 2022 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package randgen + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" + "github.com/cockroachdb/cockroach/pkg/testutils" + "github.com/cockroachdb/cockroach/pkg/testutils/skip" + "github.com/cockroachdb/cockroach/pkg/testutils/sqlutils" + "github.com/cockroachdb/cockroach/pkg/testutils/testcluster" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/cockroachdb/cockroach/pkg/util/log" + "github.com/cockroachdb/cockroach/pkg/util/randutil" + "github.com/stretchr/testify/require" +) + +func PopulateRandTableTest(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + skip.UnderRace(t, "takes >1 min under race") + rng, _ := randutil.NewPseudoRand() + dir, dirCleanupFn := testutils.TempDir(t) + defer dirCleanupFn() + params := base.TestClusterArgs{ + ServerArgs: base.TestServerArgs{ + UseDatabase: "rand", + ExternalIODir: dir, + }, + } + ctx := context.Background() + tc := testcluster.StartTestCluster(t, 1, params) + defer tc.Stopper().Stop(ctx) + sqlDB := sqlutils.MakeSQLRunner(tc.Conns[0]) + sqlDB.Exec(t, "CREATE DATABASE rand") + + stmts := RandCreateTables(rng, "table", 5, + PartialIndexMutator, + ForeignKeyMutator, + ) + + var sb strings.Builder + for _, stmt := range stmts { + sb.WriteString(tree.SerializeForDisplay(stmt)) + sb.WriteString(";\n") + } + sqlDB.Exec(t, sb.String()) + + tables := sqlDB.Query(t, `SELECT name FROM crdb_internal. +tables WHERE database_name = 'rand' AND schema_name = 'public'`) + + for tables.Next() { + numRows := rng.Intn(100) + var tableName string + if err := tables.Scan(&tableName); err != nil { + t.Fatal(err) + } + err := PopulateRandTable(rng, tc.Conns[0], tableName, numRows) + if err != nil { + t.Fatal(err) + } + res := sqlDB.QueryStr(t, fmt.Sprintf("SELECT count(*) FROM %s", tableName)) + require.Equal(t, fmt.Sprint(numRows), res[0][0]) + } +}