sqlsmith: add support for computed columns

pkg/sql/mutations/mutations.go

@@ -313,6 +313,8 @@ func statisticsMutator(
 	return stmts, changed
 }
 
+// foreignKeyMutator is a MultiStatementMutation implementation which adds
+// foreign key references between existing columns.
 func foreignKeyMutator(
 	rng *rand.Rand, stmts []tree.Statement,
 ) (mutated []tree.Statement, changed bool) {
@@ -369,6 +371,10 @@ func foreignKeyMutator(
 		// Choose a random column subset.
 		var fkCols []*tree.ColumnTableDef
 		for _, c := range cols[table.Table] {
+			if c.Computed.Computed {
+				// We don't support FK references from computed columns (#46672).
+				continue
+			}
 			if usedCols[table.Table][c.Name] {
 				continue
 			}
@@ -429,6 +435,10 @@ func foreignKeyMutator(
 				fkCol := fkCols[len(usingCols)]
 				found := false
 				for refI, refCol := range availCols {
+					if refCol.Computed.Virtual {
+						// We don't support FK references to virtual columns (#51296).
+						continue
+					}
 					fkColType := tree.MustBeStaticallyKnownType(fkCol.Type)
 					refColType := tree.MustBeStaticallyKnownType(refCol.Type)
 					if fkColType.Equivalent(refColType) && colinfo.ColumnTypeIsIndexable(refColType) {

pkg/sql/rowenc/testutils.go

@@ -1146,6 +1146,15 @@ func RandCreateTableWithInterleave(
 			}
 		}
 	}
+
+	// colIdx generates numbers that are incorporated into column names.
+	colIdx := func(ordinal int) int {
+		if generateColumnIndexNumber != nil {
+			return int(generateColumnIndexNumber())
+		}
+		return ordinal
+	}
+
 	var interleaveDef *tree.InterleaveDef
 	if interleaveIntoPK != nil && len(interleaveIntoPK.Columns) > 0 {
 		// Make the interleave prefix, which has to be exactly the columns in the
@@ -1165,11 +1174,7 @@ func RandCreateTableWithInterleave(
 			// Loop until we generate an indexable column type.
 			var extraCol *tree.ColumnTableDef
 			for {
-				colIdx := i + prefixLength
-				if generateColumnIndexNumber != nil {
-					colIdx = int(generateColumnIndexNumber())
-				}
-				extraCol = randColumnTableDef(rng, tableIdx, colIdx)
+				extraCol = randColumnTableDef(rng, tableIdx, colIdx(i+prefixLength))
 				extraColType := tree.MustBeStaticallyKnownType(extraCol.Type)
 				if colinfo.ColumnTypeIsIndexable(extraColType) {
 					break
@@ -1205,12 +1210,10 @@ func RandCreateTableWithInterleave(
 		}
 	} else {
 		// Make new defs from scratch.
-		for i := 0; i < nColumns; i++ {
-			colIdx := i
-			if generateColumnIndexNumber != nil {
-				colIdx = int(generateColumnIndexNumber())
-			}
-			columnDef := randColumnTableDef(rng, tableIdx, colIdx)
+		nComputedColumns := randutil.RandIntInRange(rng, 0, (nColumns+1)/2)
+		nNormalColumns := nColumns - nComputedColumns
+		for i := 0; i < nNormalColumns; i++ {
+			columnDef := randColumnTableDef(rng, tableIdx, colIdx(i))
 			columnDefs = append(columnDefs, columnDef)
 			defs = append(defs, columnDef)
 		}
@@ -1236,6 +1239,14 @@ func RandCreateTableWithInterleave(
 				}
 			}
 		}
+
+		// Make defs for computed columns.
+		normalColDefs := columnDefs
+		for i := nNormalColumns; i < nColumns; i++ {
+			columnDef := randComputedColumnTableDef(rng, normalColDefs, tableIdx, colIdx(i))
+			columnDefs = append(columnDefs, columnDef)
+			defs = append(defs, columnDef)
+		}
 	}
 
 	// Make indexes.
@@ -1374,11 +1385,15 @@ func IndexStoringMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
 		}
 		return colMap
 	}
-	generateStoringCols := func(rng *rand.Rand, tableCols []tree.Name, indexCols map[tree.Name]struct{}) []tree.Name {
+	generateStoringCols := func(rng *rand.Rand, tableInfo tableInfo, indexCols map[tree.Name]struct{}) []tree.Name {
 		var storingCols []tree.Name
-		for _, col := range tableCols {
-			_, ok := indexCols[col]
-			if ok {
+		for colOrdinal, col := range tableInfo.columnNames {
+			if _, ok := indexCols[col]; ok {
+				// Skip PK columns and columns already in the index.
+				continue
+			}
+			if tableInfo.columnsTableDefs[colOrdinal].Computed.Virtual {
+				// Virtual columns can't be stored.
 				continue
 			}
 			if rng.Intn(2) == 0 {
@@ -1403,7 +1418,7 @@ func IndexStoringMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
 				for _, elem := range ast.Columns {
 					indexCols[elem.Column] = struct{}{}
 				}
-				ast.Storing = generateStoringCols(rng, tableInfo.columnNames, indexCols)
+				ast.Storing = generateStoringCols(rng, tableInfo, indexCols)
 				changed = true
 			}
 		case *tree.CreateTable:
@@ -1430,7 +1445,7 @@ func IndexStoringMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
 					for _, elem := range idx.Columns {
 						indexCols[elem.Column] = struct{}{}
 					}
-					idx.Storing = generateStoringCols(rng, tableInfo.columnNames, indexCols)
+					idx.Storing = generateStoringCols(rng, tableInfo, indexCols)
 					changed = true
 				}
 			}
@@ -1493,8 +1508,8 @@ func PartialIndexMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
 	return stmts, changed
 }
 
-// randColumnTableDef produces a random ColumnTableDef, with a random type and
-// nullability.
+// randColumnTableDef produces a random ColumnTableDef for a non-computed
+// column, with a random type and nullability.
 func randColumnTableDef(rand *rand.Rand, tableIdx int, colIdx int) *tree.ColumnTableDef {
 	columnDef := &tree.ColumnTableDef{
 		// We make a unique name for all columns by prefixing them with the table
@@ -1506,6 +1521,111 @@ func randColumnTableDef(rand *rand.Rand, tableIdx int, colIdx int) *tree.ColumnT
 	return columnDef
 }
 
+// randComputedColumnTableDef produces a random ColumnTableDef for a computed
+// column (either STORED or VIRTUAL). The computed expressions refer to columns
+// in normalColDefs.
+func randComputedColumnTableDef(
+	rng *rand.Rand, normalColDefs []*tree.ColumnTableDef, tableIdx int, colIdx int,
+) *tree.ColumnTableDef {
+	newDef := randColumnTableDef(rng, tableIdx, colIdx)
+	newDef.Computed.Computed = true
+	newDef.Computed.Virtual = (rng.Intn(2) == 0)
+
+	if rng.Intn(2) == 0 {
+		// Try to find a set of numeric columns with the same type; the computed
+		// expression will be of the form "a+b+c".
+		var cols []*tree.ColumnTableDef
+		var fam types.Family
+		for _, idx := range rng.Perm(len(normalColDefs)) {
+			x := normalColDefs[idx]
+			xFam := x.Type.(*types.T).Family()
+
+			if len(cols) == 0 {
+				switch xFam {
+				case types.IntFamily, types.FloatFamily, types.DecimalFamily:
+					fam = xFam
+					cols = append(cols, x)
+				}
+			} else if fam == xFam {
+				cols = append(cols, x)
+				if len(cols) > 1 && rng.Intn(2) == 0 {
+					break
+				}
+			}
+		}
+		if len(cols) > 1 {
+			var expr tree.Expr
+			expr = tree.NewUnresolvedName(string(cols[0].Name))
+			for _, x := range cols[1:] {
+				expr = &tree.BinaryExpr{
+					Operator: tree.Plus,
+					Left:     expr,
+					Right:    tree.NewUnresolvedName(string(x.Name)),
+				}
+			}
+			newDef.Type = cols[0].Type
+			newDef.Computed.Expr = expr
+			return newDef
+		}
+	}
+
+	// Pick a single column and create a computed column that depends on it.
+	// The expression is as follows:
+	//  - for numeric types (int, float, decimal), the expression is "x+1";
+	//  - for string type, the expression is "lower(x)";
+	//  - for types that can be cast to string in computed columns, the expression
+	//    is "lower(x::string)";
+	//  - otherwise, the expression is "IF(x IS NULL, 'foo', 'bar')".
+	x := normalColDefs[randutil.RandIntInRange(rng, 0, len(normalColDefs))]
+	xTyp := x.Type.(*types.T)
+
+	switch xTyp.Family() {
+	case types.IntFamily, types.FloatFamily, types.DecimalFamily:
+		newDef.Type = xTyp
+		newDef.Computed.Expr = &tree.BinaryExpr{
+			Operator: tree.Plus,
+			Left:     tree.NewUnresolvedName(string(x.Name)),
+			Right:    RandDatum(rng, xTyp, false /* nullOk */),
+		}
+
+	case types.StringFamily:
+		newDef.Type = types.String
+		newDef.Computed.Expr = &tree.FuncExpr{
+			Func:  tree.WrapFunction("lower"),
+			Exprs: tree.Exprs{tree.NewUnresolvedName(string(x.Name))},
+		}
+
+	default:
+		volatility, ok := tree.LookupCastVolatility(xTyp, types.String)
+		if ok && volatility <= tree.VolatilityImmutable {
+			// We can cast to string; use lower(x::string)
+			newDef.Type = types.String
+			newDef.Computed.Expr = &tree.FuncExpr{
+				Func: tree.WrapFunction("lower"),
+				Exprs: tree.Exprs{
+					&tree.CastExpr{
+						Expr: tree.NewUnresolvedName(string(x.Name)),
+						Type: types.String,
+					},
+				},
+			}
+		} else {
+			// We cannot cast this type to string in a computed column expression.
+			// Use IF(x IS NULL, 'foo', 'bar').
+			newDef.Type = types.String
+			newDef.Computed.Expr = &tree.IfExpr{
+				Cond: &tree.IsNullExpr{
+					Expr: tree.NewUnresolvedName(string(x.Name)),
+				},
+				True: RandDatum(rng, types.String, true /* nullOK */),
+				Else: RandDatum(rng, types.String, true /* nullOK */),
+			}
+		}
+	}
+
+	return newDef
+}
+
 // randIndexTableDefFromCols creates an IndexTableDef with a random subset of
 // the given columns and a random direction.
 func randIndexTableDefFromCols(

pkg/sql/tests/random_schema_test.go

@@ -19,6 +19,7 @@ import (
 
 	"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/tests"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
@@ -46,23 +47,27 @@ func TestCreateRandomSchema(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	toStr := func(c tree.Statement) string {
+		return tree.AsStringWithFlags(c, tree.FmtParsable)
+	}
+
 	rng := rand.New(rand.NewSource(timeutil.Now().UnixNano()))
 	for i := 0; i < 100; i++ {
-		tab := rowenc.RandCreateTable(rng, "table", i)
+		createTable := rowenc.RandCreateTable(rng, "table", i)
 		setDb(t, db, "test")
-		_, err := db.Exec(tab.String())
+		_, err := db.Exec(toStr(createTable))
 		if err != nil {
-			t.Fatal(tab, err)
+			t.Fatal(createTable, err)
 		}
 
 		var tabName, tabStmt, secondTabStmt string
 		if err := db.QueryRow(fmt.Sprintf("SHOW CREATE TABLE %s",
-			tab.Table.String())).Scan(&tabName, &tabStmt); err != nil {
+			createTable.Table.String())).Scan(&tabName, &tabStmt); err != nil {
 			t.Fatal(err)
 		}
 
-		if tabName != tab.Table.String() {
-			t.Fatalf("found table name %s, expected %s", tabName, tab.Table.String())
+		if tabName != createTable.Table.String() {
+			t.Fatalf("found table name %s, expected %s", tabName, createTable.Table.String())
 		}
 
 		// Reparse the show create table statement that's stored in the database.
@@ -75,7 +80,7 @@ func TestCreateRandomSchema(t *testing.T) {
 		// Now run the SHOW CREATE TABLE statement we found on a new db and verify
 		// that both tables are the same.
 
-		_, err = db.Exec(parsed.AST.String())
+		_, err = db.Exec(tree.AsStringWithFlags(parsed.AST, tree.FmtParsable))
 		if err != nil {
 			t.Fatal(parsed.AST, err)
 		}
@@ -85,21 +90,21 @@ func TestCreateRandomSchema(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		if tabName != tab.Table.String() {
-			t.Fatalf("found table name %s, expected %s", tabName, tab.Table.String())
+		if tabName != createTable.Table.String() {
+			t.Fatalf("found table name %s, expected %s", tabName, createTable.Table.String())
 		}
 		// Reparse the show create table statement that's stored in the database.
 		secondParsed, err := parser.ParseOne(secondTabStmt)
 		if err != nil {
 			t.Fatalf("error parsing show create table: %s", err)
 		}
-		if parsed.AST.String() != secondParsed.AST.String() {
+		if toStr(parsed.AST) != toStr(secondParsed.AST) {
 			t.Fatalf("for input statement\n%s\nfound first output\n%q\nbut second output\n%q",
-				tab.String(), parsed.AST.String(), secondParsed.AST.String())
+				toStr(createTable), toStr(parsed.AST), toStr(secondParsed.AST))
 		}
 		if tabStmt != secondTabStmt {
 			t.Fatalf("for input statement\n%s\nfound first output\n%q\nbut second output\n%q",
-				tab.String(), tabStmt, secondTabStmt)
+				toStr(createTable), tabStmt, secondTabStmt)
 		}
 	}
 }