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 cockroachdb#72345

Release note: None

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

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

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.

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