Merge remote-tracking branch 'upstream/main' into am-random-2

Signed-off-by: Andres Taylor <andres@planetscale.com>

go/test/endtoend/vtgate/queries/random/query_gen.go

@@ -20,9 +20,9 @@ import (
 	"fmt"
 	"math/rand"
 
-	"golang.org/x/exp/slices"
+	"vitess.io/vitess/go/slice"
 
-	"vitess.io/vitess/go/slices2"
+	"golang.org/x/exp/slices"
 
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -298,7 +298,7 @@ func (sg *selectGenerator) randomSelect() {
 	// TODO: only_full_group_by related errors in Vitess
 	var exprGenerators []sqlparser.ExprGenerator
 	if canAggregate && testFailingQueries {
-		exprGenerators = slices2.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })
+		exprGenerators = slice.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })
 		// add scalar subqueries
 		if sg.r.Intn(10) < 1 {
 			exprGenerators = append(exprGenerators, sg)
@@ -470,7 +470,7 @@ func (sg *selectGenerator) aliasGroupingColumns(grouping []column) []column {
 
 // returns the aggregation columns as three types: sqlparser.SelectExprs, []column
 func (sg *selectGenerator) createAggregations(tables []tableT) (aggregates []column) {
-	exprGenerators := slices2.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })
+	exprGenerators := slice.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })
 	// add scalar subqueries
 	// TODO: subqueries fail
 	if sg.r.Intn(10) < 1 && testFailingQueries {
@@ -511,7 +511,7 @@ func (sg *selectGenerator) createOrderBy() {
 
 // returns 0-2 random expressions based on tables
 func (sg *selectGenerator) createWherePredicates(tables []tableT) {
-	exprGenerators := slices2.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })
+	exprGenerators := slice.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })
 	// add scalar subqueries
 	// TODO: subqueries fail
 	if sg.r.Intn(10) < 1 && testFailingQueries {
@@ -524,7 +524,7 @@ func (sg *selectGenerator) createWherePredicates(tables []tableT) {
 
 // creates predicates for the having clause comparing a column to a random expression
 func (sg *selectGenerator) createHavingPredicates(grouping []column) {
-	exprGenerators := slices2.Map(grouping, func(c column) sqlparser.ExprGenerator { return &c })
+	exprGenerators := slice.Map(grouping, func(c column) sqlparser.ExprGenerator { return &c })
 	// add scalar subqueries
 	// TODO: subqueries fail
 	if sg.r.Intn(10) < 1 && testFailingQueries {
@@ -613,7 +613,7 @@ func (sg *selectGenerator) matchNumCols(tables []tableT, newTable tableT, canAgg
 	if sg.genConfig.NumCols > len(sg.sel.SelectExprs) {
 		diff := sg.genConfig.NumCols - len(sg.sel.SelectExprs)
 		exprs := sg.createRandomExprs(diff, diff,
-			slices2.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })...)
+			slice.Map(tables, func(t tableT) sqlparser.ExprGenerator { return &t })...)
 
 		for i, expr := range exprs {
 			col := sg.randomlyAlias(expr, fmt.Sprintf("crandom%d", i+1))

go/test/endtoend/vtgate/queries/random/random_expr_test.go

@@ -53,7 +53,7 @@ func TestRandomExprWithTables(t *testing.T) {
 		seed := time.Now().UnixNano()
 		r := rand.New(rand.NewSource(seed))
 		genConfig := sqlparser.NewExprGeneratorConfig(sqlparser.CanAggregate, "", 0, false)
-		g := sqlparser.NewGenerator(r, 3, slices2.Map(schemaTables, func(t tableT) sqlparser.ExprGenerator { return &t })...)
+		g := sqlparser.NewGenerator(r, 3, slice.Map(schemaTables, func(t tableT) sqlparser.ExprGenerator { return &t })...)
 		g.Expression(genConfig)
 	}
 }

go/test/endtoend/vtgate/queries/random/random_test.go

@@ -317,7 +317,7 @@ func TestRandom(t *testing.T) {
 					break
 				}
 				// EOF
-				if sqlError, ok := vtErr.(*mysql.SQLError); ok && strings.Contains(sqlError.Message, "EOF") {
+				if sqlError, ok := vtErr.(*sqlerror.SQLError); ok && strings.Contains(sqlError.Message, "EOF") {
 					break
 				}
 			}

go/test/vschemawrapper/vschema_wrapper.go

@@ -123,8 +123,12 @@ func (vw *VSchemaWrapper) ConnCollation() collations.ID {
 func (vw *VSchemaWrapper) PlannerWarning(_ string) {
 }
 
-func (vw *VSchemaWrapper) ForeignKeyMode() string {
-	return "allow"
+func (vw *VSchemaWrapper) ForeignKeyMode(keyspace string) (vschemapb.Keyspace_ForeignKeyMode, error) {
+	defaultFkMode := vschemapb.Keyspace_FK_UNMANAGED
+	if vw.V.Keyspaces[keyspace] != nil && vw.V.Keyspaces[keyspace].ForeignKeyMode != vschemapb.Keyspace_FK_DEFAULT {
+		return vw.V.Keyspaces[keyspace].ForeignKeyMode, nil
+	}
+	return defaultFkMode, nil
 }
 
 func (vw *VSchemaWrapper) AllKeyspace() ([]*vindexes.Keyspace, error) {

go/vt/vtgate/planbuilder/plan_test.go

@@ -99,10 +99,10 @@ func TestPlan(t *testing.T) {
 
 // TestForeignKeyPlanning tests the planning of foreign keys in a managed mode by Vitess.
 func TestForeignKeyPlanning(t *testing.T) {
-	vschemaWrapper := &vschemaWrapper{
-		v: loadSchema(t, "vschemas/schema.json", true),
+	vschemaWrapper := &vschemawrapper.VSchemaWrapper{
+		V: loadSchema(t, "vschemas/schema.json", true),
 		// Set the keyspace with foreign keys enabled as the default.
-		keyspace: &vindexes.Keyspace{
+		Keyspace: &vindexes.Keyspace{
 			Name:    "user_fk_allow",
 			Sharded: true,
 		},
@@ -116,7 +116,9 @@ func TestSystemTables57(t *testing.T) {
 	// first we move everything to use 5.7 logic
 	oldVer := servenv.MySQLServerVersion()
 	servenv.SetMySQLServerVersionForTest("5.7")
-	defer servenv.SetMySQLServerVersionForTest("")
+	defer func() {
+		servenv.SetMySQLServerVersionForTest(oldVer)
+	}()
 	vschemaWrapper := &vschemawrapper.VSchemaWrapper{V: loadSchema(t, "vschemas/schema.json", true)}
 	testOutputTempDir := makeTestOutput(t)
 	testFile(t, "info_schema57_cases.json", testOutputTempDir, vschemaWrapper, false)
@@ -208,7 +210,9 @@ func TestOneWith57Version(t *testing.T) {
 	// first we move everything to use 5.7 logic
 	oldVer := servenv.MySQLServerVersion()
 	servenv.SetMySQLServerVersionForTest("5.7")
-	defer servenv.SetMySQLServerVersionForTest("")
+	defer func() {
+		servenv.SetMySQLServerVersionForTest(oldVer)
+	}()
 	vschema := &vschemawrapper.VSchemaWrapper{V: loadSchema(t, "vschemas/schema.json", true)}
 
 	testFile(t, "onecase.json", "", vschema, false)
@@ -446,6 +450,17 @@ func loadSchema(t testing.TB, filename string, setCollation bool) *vindexes.VSch
 	return vschema
 }
 
+// createFkDefinition is a helper function to create a Foreign key definition struct from the columns used in it provided as list of strings.
+func createFkDefinition(childCols []string, parentTableName string, parentCols []string) *sqlparser.ForeignKeyDefinition {
+	return &sqlparser.ForeignKeyDefinition{
+		Source: sqlparser.MakeColumns(childCols...),
+		ReferenceDefinition: &sqlparser.ReferenceDefinition{
+			ReferencedTable:   sqlparser.NewTableName(parentTableName),
+			ReferencedColumns: sqlparser.MakeColumns(parentCols...),
+		},
+	}
+}
+
 type (
 	planTest struct {
 		Comment string          `json:"comment,omitempty"`