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

@@ -103,9 +103,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,23 @@ go_library(
 
 go_test(
     name = "randgen_test",
-    srcs = ["mutator_test.go"],
+    srcs = [
+        "main_test.go",
+        "mutator_test.go",
+        "schema_test.go",
+    ],
     embed = [":randgen"],
-    deps = ["//pkg/util/randutil"],
+    deps = [
+        "//pkg/base",
+        "//pkg/security",
+        "//pkg/security/securitytest",
+        "//pkg/server",
+        "//pkg/sql/sem/tree",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/sqlutils",
+        "//pkg/testutils/testcluster",
+        "//pkg/util/leaktest",
+        "//pkg/util/randutil",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/sql/randgen/main_test.go

@@ -0,0 +1,33 @@
+// Copyright 2019 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 (
+	"os"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/security"
+	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+//go:generate ../../util/leaktest/add-leaktest.sh *_test.go
+
+func TestMain(m *testing.M) {
+	security.SetAssetLoader(securitytest.EmbeddedAssets)
+	randutil.SeedForTests()
+	serverutils.InitTestServerFactory(server.TestServerFactory)
+	serverutils.InitTestClusterFactory(testcluster.TestClusterFactory)
+	os.Exit(m.Run())
+}

pkg/sql/randgen/mutator_test.go

@@ -14,10 +14,12 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 )
 
 func TestPostgresMutator(t *testing.T) {
+	defer leaktest.AfterTest(t)()
 	q := `
 		CREATE TABLE t (s STRING FAMILY fam1, b BYTES, FAMILY fam2 (b), PRIMARY KEY (s ASC, b DESC), INDEX (s) STORING (b))
 		    PARTITION BY LIST (s)

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,102 @@ 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 * 100
+	for success < numRows {
+		var valBuilder strings.Builder
+		valBuilder.WriteString("(")
+		comma := ""
+		for j := 0; j < len(colTypes); j++ {
+			valBuilder.WriteString(comma)
+			var d tree.Datum
+			if rand.Intn(6)%3 == 0 {
+				d = randInterestingDatum(rng, colTypes[j])
+			}
+			if colTypes[j].Family() == types.OidFamily {
+				d = tree.NewDOid(tree.DInt(rand.Intn(2)))
+			}
+			if d == nil {
+				d = RandDatum(rng, colTypes[j], nullable[j])
+			}
+			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 {
+				// This could mean PopulateRandTable or RandDatum couldn't handle this table's schema.
+				// consider filing a bug. schema:
+				return fmt.Errorf(`could not populate table schema \n \t %s \n with %d insert attempts`, createStmt.String(), maxTries)
+			}
+		} 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,72 @@
+// 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/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPopulateRandTable(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	ctx := context.Background()
+	s, dbConn, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+
+	rng, _ := randutil.NewTestRand()
+
+	sqlDB := sqlutils.MakeSQLRunner(dbConn)
+	sqlDB.Exec(t, "CREATE DATABASE rand")
+
+	tablePrefix := "table"
+	numTables := 10
+
+	stmts := RandCreateTables(rng, tablePrefix, numTables,
+		PartialIndexMutator,
+		ForeignKeyMutator,
+	)
+
+	var sb strings.Builder
+	for _, stmt := range stmts {
+		sb.WriteString(tree.SerializeForDisplay(stmt))
+		sb.WriteString(";\n")
+	}
+	sqlDB.Exec(t, sb.String())
+
+	// To prevent the test from being flaky, pass the test if PopulateRandTable
+	// works at least once
+	success := false
+	for i := 1; i <= numTables; i++ {
+		tableName := tablePrefix + fmt.Sprint(i)
+		numRows := rng.Intn(100)
+		err := PopulateRandTable(rng, dbConn, tableName, numRows)
+		if err != nil {
+			t.Log(err)
+			continue
+		}
+		res := sqlDB.QueryStr(t, fmt.Sprintf("SELECT count(*) FROM %s", tableName))
+		require.Equal(t, fmt.Sprint(numRows), res[0][0])
+		success = true
+		break
+	}
+	require.Equal(t, true, success)
+}