Merge #42298

42298: sql/mutations: begin work on multi-statement mutations r=mjibson a=mjibson



Co-authored-by: Matt Jibson <matt.jibson@gmail.com>

pkg/internal/sqlsmith/setup.go

@@ -15,6 +15,7 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/mutations"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 )
 
@@ -56,7 +57,7 @@ func stringSetup(s string) Setup {
 }
 
 func randTables(r *rand.Rand) string {
-	stmts := sqlbase.RandCreateTables(r, "table", r.Intn(5)+1)
+	stmts := sqlbase.RandCreateTables(r, "table", r.Intn(5)+1, mutations.ForeignKeyMutator)
 
 	var sb strings.Builder
 	for _, stmt := range stmts {

pkg/sql/logictest/logic.go

@@ -19,6 +19,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"math/rand"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -56,6 +57,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/lib/pq"
@@ -864,6 +866,7 @@ type logicQuery struct {
 type logicTest struct {
 	rootT    *testing.T
 	subtestT *testing.T
+	rng      *rand.Rand
 	cfg      testClusterConfig
 	// the number of nodes in the cluster.
 	cluster serverutils.TestClusterInterface
@@ -1902,7 +1905,7 @@ func (t *logicTest) execStatement(stmt logicStatement) (bool, error) {
 	if *showSQL {
 		t.outf("%s;", stmt.sql)
 	}
-	execSQL, changed := mutations.ForAllStatements(stmt.sql, mutations.ColumnFamilyMutator)
+	execSQL, changed := mutations.ForAllStatements(t.rng, stmt.sql, mutations.ColumnFamilyMutator)
 	for _, c := range changed {
 		t.outf("rewrote: %s;", c)
 	}
@@ -2312,10 +2315,12 @@ func RunLogicTest(t *testing.T, globs ...string) {
 							t.Parallel() // SAFE FOR TESTING (this comments satisfies the linter)
 						}
 					}
+					rng, _ := randutil.NewPseudoRand()
 					lt := logicTest{
 						rootT:           t,
 						verbose:         verbose,
 						perErrorSummary: make(map[string][]string),
+						rng:             rng,
 					}
 					if *printErrorSummary {
 						defer lt.printErrorSummary()

pkg/sql/logictest/parallel_test.go

@@ -39,6 +39,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/gogo/protobuf/proto"
 	yaml "gopkg.in/yaml.v2"
@@ -68,13 +69,15 @@ func (t *parallelTest) processTestFile(path string, nodeIdx int, db *gosql.DB, c
 	}
 
 	// Set up a dummy logicTest structure to use that code.
+	rng, _ := randutil.NewPseudoRand()
 	l := &logicTest{
 		rootT:   t.T,
 		cluster: t.cluster,
 		nodeIdx: nodeIdx,
 		db:      db,
 		user:    security.RootUser,
 		verbose: testing.Verbose() || log.V(1),
+		rng:     rng,
 	}
 	if err := l.processTestFile(path, testClusterConfig{}); err != nil {
 		log.Errorf(context.Background(), "error processing %s: %s", path, err)

pkg/sql/mutations/mutations.go

@@ -19,12 +19,16 @@ import (
 )
 
 // StatementMutator defines a func that can change a statement.
-type StatementMutator func(stmt tree.Statement) (changed bool)
+type StatementMutator func(rng *rand.Rand, stmt tree.Statement) (changed bool)
+
+// MultiStatementMutation defines a func that returns additional statements,
+// but must not change any of the statements passed.
+type MultiStatementMutation func(rng *rand.Rand, stmts []tree.Statement) (additional []tree.Statement)
 
 // ForAllStatements executes stmtMutator on all SQL statements of input. The
 // list of changed statements is returned.
 func ForAllStatements(
-	input string, stmtMutator StatementMutator,
+	rng *rand.Rand, input string, stmtMutator StatementMutator,
 ) (output string, changed []string) {
 	parsed, err := parser.Parse(input)
 	if err != nil {
@@ -33,7 +37,7 @@ func ForAllStatements(
 
 	var sb strings.Builder
 	for _, p := range parsed {
-		stmtChanged := stmtMutator(p.AST)
+		stmtChanged := stmtMutator(rng, p.AST)
 		if !stmtChanged {
 			sb.WriteString(p.SQL)
 		} else {
@@ -46,9 +50,150 @@ func ForAllStatements(
 	return sb.String(), changed
 }
 
+// ForeignKeyMutator adds ALTER TABLE ADD FOREIGN KEY statements.
+func ForeignKeyMutator(rng *rand.Rand, stmts []tree.Statement) (additional []tree.Statement) {
+	// Find columns in the tables.
+	cols := map[tree.TableName][]*tree.ColumnTableDef{}
+	byName := map[tree.TableName]*tree.CreateTable{}
+	var tables []*tree.CreateTable
+	for _, stmt := range stmts {
+		table, ok := stmt.(*tree.CreateTable)
+		if !ok {
+			continue
+		}
+		tables = append(tables, table)
+		byName[table.Table] = table
+		for _, def := range table.Defs {
+			switch def := def.(type) {
+			case *tree.ColumnTableDef:
+				cols[table.Table] = append(cols[table.Table], def)
+			}
+		}
+	}
+
+	toNames := func(cols []*tree.ColumnTableDef) tree.NameList {
+		names := make(tree.NameList, len(cols))
+		for i, c := range cols {
+			names[i] = c.Name
+		}
+		return names
+	}
+
+	// We cannot mutate the table definitions themselves because 1) we
+	// don't know the order of dependencies (i.e., table 1 could reference
+	// table 4 which doesn't exist yet) and relatedly 2) we don't prevent
+	// circular dependencies. Instead, add new ALTER TABLE commands to the
+	// end of a list of statements.
+
+	// Keep track of referencing columns since we have a limitation that a
+	// column can only be used by one FK.
+	usedCols := map[tree.TableName]map[tree.Name]bool{}
+
+	// Create some FKs.
+	for rng.Intn(2) == 0 {
+		// Choose a random table.
+		table := tables[rng.Intn(len(tables))]
+		if _, ok := usedCols[table.Table]; !ok {
+			usedCols[table.Table] = map[tree.Name]bool{}
+		}
+		// Choose a random column subset.
+		var fkCols []*tree.ColumnTableDef
+		for _, c := range cols[table.Table] {
+			if usedCols[table.Table][c.Name] {
+				continue
+			}
+			fkCols = append(fkCols, c)
+		}
+		if len(fkCols) == 0 {
+			continue
+		}
+		rng.Shuffle(len(fkCols), func(i, j int) {
+			fkCols[i], fkCols[j] = fkCols[j], fkCols[i]
+		})
+		// Pick some randomly short prefix. I'm sure there's a closed
+		// form solution to this with a single call to rng.Intn but I'm
+		// not sure what to search for.
+		i := 1
+		for len(fkCols) > i && rng.Intn(2) == 0 {
+			i++
+		}
+		fkCols = fkCols[:i]
+
+		// Check if a table has the needed column types.
+	LoopTable:
+		for refTable, refCols := range cols {
+			// Prevent self references.
+			// TODO(mjibson): Circular references are not
+			// prevented, but it would be nice to detect and
+			// prevent them.
+			if refTable == table.Table || len(refCols) < len(fkCols) {
+				continue
+			}
+
+			// We found a table with enough columns. Check if it
+			// has some columns that are needed types. In order
+			// to not use columns multiple times, keep track of
+			// available columns.
+			availCols := append([]*tree.ColumnTableDef(nil), refCols...)
+			var usingCols []*tree.ColumnTableDef
+			for len(availCols) > 0 && len(usingCols) < len(fkCols) {
+				fkCol := fkCols[len(usingCols)]
+				found := false
+				for refI, refCol := range availCols {
+					if fkCol.Type.Equivalent(refCol.Type) {
+						usingCols = append(usingCols, refCol)
+						availCols = append(availCols[:refI], availCols[refI+1:]...)
+						found = true
+						break
+					}
+				}
+				if !found {
+					continue LoopTable
+				}
+			}
+			// If we didn't find enough columns, try another table.
+			if len(usingCols) != len(fkCols) {
+				continue
+			}
+
+			// Found a suitable table.
+			// TODO(mjibson): prevent the creation of unneeded
+			// unique indexes. One may already exist with the
+			// correct prefix.
+			ref := byName[refTable]
+			refColumns := make(tree.IndexElemList, len(usingCols))
+			for i, c := range usingCols {
+				refColumns[i].Column = c.Name
+			}
+			for _, c := range fkCols {
+				usedCols[table.Table][c.Name] = true
+			}
+			ref.Defs = append(ref.Defs, &tree.UniqueConstraintTableDef{
+				IndexTableDef: tree.IndexTableDef{
+					Columns: refColumns,
+				},
+			})
+
+			additional = append(additional, &tree.AlterTable{
+				Table: table.Table.ToUnresolvedObjectName(),
+				Cmds: tree.AlterTableCmds{&tree.AlterTableAddConstraint{
+					ConstraintDef: &tree.ForeignKeyConstraintTableDef{
+						Table:    ref.Table,
+						FromCols: toNames(fkCols),
+						ToCols:   toNames(usingCols),
+					},
+				}},
+			})
+			break
+		}
+	}
+
+	return additional
+}
+
 // ColumnFamilyMutator modifies a CREATE TABLE statement without any FAMILY
 // definitions to have random FAMILY definitions.
-func ColumnFamilyMutator(stmt tree.Statement) (changed bool) {
+func ColumnFamilyMutator(rng *rand.Rand, stmt tree.Statement) (changed bool) {
 	ast, ok := stmt.(*tree.CreateTable)
 	if !ok {
 		return false
@@ -105,7 +250,7 @@ func ColumnFamilyMutator(stmt tree.Statement) (changed bool) {
 		}
 		columns = columns[:n]
 	}
-	rand.Shuffle(len(columns), func(i, j int) {
+	rng.Shuffle(len(columns), func(i, j int) {
 		columns[i], columns[j] = columns[j], columns[i]
 	})
 	fd := &tree.FamilyTableDef{}
@@ -119,7 +264,7 @@ func ColumnFamilyMutator(stmt tree.Statement) (changed bool) {
 		fd.Columns = append(fd.Columns, columns[0])
 		columns = columns[1:]
 		// 50% chance to make a new column family.
-		if rand.Intn(2) != 0 {
+		if rng.Intn(2) != 0 {
 			ast.Defs = append(ast.Defs, fd)
 			fd = &tree.FamilyTableDef{}
 		}