Make Foreign_key_checks a Vitess Aware variable

go/test/endtoend/vtgate/foreignkey/fk_fuzz_test.go

@@ -33,6 +33,7 @@ import (
 	"vitess.io/vitess/go/test/endtoend/cluster"
 	"vitess.io/vitess/go/test/endtoend/utils"
 	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/sqlparser"
 )
 
 type QueryFormat string
@@ -53,6 +54,7 @@ type fuzzer struct {
 	updateShare  int
 	concurrency  int
 	queryFormat  QueryFormat
+	fkState      sqlparser.FkChecksState
 
 	// shouldStop is an internal state variable, that tells the fuzzer
 	// whether it should stop or not.
@@ -72,7 +74,7 @@ type debugInfo struct {
 }
 
 // newFuzzer creates a new fuzzer struct.
-func newFuzzer(concurrency int, maxValForId int, maxValForCol int, insertShare int, deleteShare int, updateShare int, queryFormat QueryFormat) *fuzzer {
+func newFuzzer(concurrency int, maxValForId int, maxValForCol int, insertShare int, deleteShare int, updateShare int, queryFormat QueryFormat, fkState sqlparser.FkChecksState) *fuzzer {
 	fz := &fuzzer{
 		concurrency:  concurrency,
 		maxValForId:  maxValForId,
@@ -81,6 +83,7 @@ func newFuzzer(concurrency int, maxValForId int, maxValForCol int, insertShare i
 		deleteShare:  deleteShare,
 		updateShare:  updateShare,
 		queryFormat:  queryFormat,
+		fkState:      fkState,
 		wg:           sync.WaitGroup{},
 	}
 	// Initially the fuzzer thread is stopped.
@@ -128,17 +131,18 @@ func (fz *fuzzer) generateInsertDMLQuery() string {
 	tableId := rand.Intn(len(fkTables))
 	idValue := 1 + rand.Intn(fz.maxValForId)
 	tableName := fkTables[tableId]
+	setVarFkChecksVal := fz.getSetVarFkChecksVal()
 	if tableName == "fk_t20" {
 		colValue := rand.Intn(1 + fz.maxValForCol)
 		col2Value := rand.Intn(1 + fz.maxValForCol)
-		return fmt.Sprintf("insert into %v (id, col, col2) values (%v, %v, %v)", tableName, idValue, convertColValueToString(colValue), convertColValueToString(col2Value))
+		return fmt.Sprintf("insert %vinto %v (id, col, col2) values (%v, %v, %v)", setVarFkChecksVal, tableName, idValue, convertColValueToString(colValue), convertColValueToString(col2Value))
 	} else if isMultiColFkTable(tableName) {
 		colaValue := rand.Intn(1 + fz.maxValForCol)
 		colbValue := rand.Intn(1 + fz.maxValForCol)
-		return fmt.Sprintf("insert into %v (id, cola, colb) values (%v, %v, %v)", tableName, idValue, convertColValueToString(colaValue), convertColValueToString(colbValue))
+		return fmt.Sprintf("insert %vinto %v (id, cola, colb) values (%v, %v, %v)", setVarFkChecksVal, tableName, idValue, convertColValueToString(colaValue), convertColValueToString(colbValue))
 	} else {
 		colValue := rand.Intn(1 + fz.maxValForCol)
-		return fmt.Sprintf("insert into %v (id, col) values (%v, %v)", tableName, idValue, convertColValueToString(colValue))
+		return fmt.Sprintf("insert %vinto %v (id, col) values (%v, %v)", setVarFkChecksVal, tableName, idValue, convertColValueToString(colValue))
 	}
 }
 
@@ -155,25 +159,27 @@ func (fz *fuzzer) generateUpdateDMLQuery() string {
 	tableId := rand.Intn(len(fkTables))
 	idValue := 1 + rand.Intn(fz.maxValForId)
 	tableName := fkTables[tableId]
+	setVarFkChecksVal := fz.getSetVarFkChecksVal()
 	if tableName == "fk_t20" {
 		colValue := rand.Intn(1 + fz.maxValForCol)
 		col2Value := rand.Intn(1 + fz.maxValForCol)
-		return fmt.Sprintf("update %v set col = %v, col2 = %v where id = %v", tableName, convertColValueToString(colValue), convertColValueToString(col2Value), idValue)
+		return fmt.Sprintf("update %v%v set col = %v, col2 = %v where id = %v", setVarFkChecksVal, tableName, convertColValueToString(colValue), convertColValueToString(col2Value), idValue)
 	} else if isMultiColFkTable(tableName) {
 		colaValue := rand.Intn(1 + fz.maxValForCol)
 		colbValue := rand.Intn(1 + fz.maxValForCol)
-		return fmt.Sprintf("update %v set cola = %v, colb = %v where id = %v", tableName, convertColValueToString(colaValue), convertColValueToString(colbValue), idValue)
+		return fmt.Sprintf("update %v%v set cola = %v, colb = %v where id = %v", setVarFkChecksVal, tableName, convertColValueToString(colaValue), convertColValueToString(colbValue), idValue)
 	} else {
 		colValue := rand.Intn(1 + fz.maxValForCol)
-		return fmt.Sprintf("update %v set col = %v where id = %v", tableName, convertColValueToString(colValue), idValue)
+		return fmt.Sprintf("update %v%v set col = %v where id = %v", setVarFkChecksVal, tableName, convertColValueToString(colValue), idValue)
 	}
 }
 
 // generateDeleteDMLQuery generates a DELETE query from the parameters for the fuzzer.
 func (fz *fuzzer) generateDeleteDMLQuery() string {
 	tableId := rand.Intn(len(fkTables))
 	idValue := 1 + rand.Intn(fz.maxValForId)
-	query := fmt.Sprintf("delete from %v where id = %v", fkTables[tableId], idValue)
+	setVarFkChecksVal := fz.getSetVarFkChecksVal()
+	query := fmt.Sprintf("delete %vfrom %v where id = %v", setVarFkChecksVal, fkTables[tableId], idValue)
 	return query
 }
 
@@ -199,6 +205,9 @@ func (fz *fuzzer) runFuzzerThread(t *testing.T, sharded bool, fuzzerThreadId int
 	// Create a MySQL Compare that connects to both Vitess and MySQL and runs the queries against both.
 	mcmp, err := utils.NewMySQLCompare(t, vtParams, mysqlParams)
 	require.NoError(t, err)
+	if fz.fkState != sqlparser.FkChecksUnspecified {
+		mcmp.Exec(fmt.Sprintf("SET FOREIGN_KEY_CHECKS=%v", fz.fkState.String()))
+	}
 	var vitessDb, mysqlDb *sql.DB
 	if fz.queryFormat == PreparedStatementPacket {
 		// Open another connection to Vitess using the go-sql-driver so that we can send prepared statements as COM_STMT_PREPARE packets.
@@ -456,6 +465,21 @@ func (fz *fuzzer) generateParameterizedDeleteQuery() (query string, params []any
 	return fmt.Sprintf("delete from %v where id = ?", fkTables[tableId]), []any{idValue}
 }
 
+// getSetVarFkChecksVal generates an optimizer hint to randomly set the foreign key checks to on or off or leave them unaltered.
+func (fz *fuzzer) getSetVarFkChecksVal() string {
+	if fz.concurrency != 1 {
+		return ""
+	}
+	val := rand.Intn(3)
+	if val == 0 {
+		return ""
+	}
+	if val == 1 {
+		return "/*+ SET_VAR(foreign_key_checks=On) */ "
+	}
+	return "/*+ SET_VAR(foreign_key_checks=Off) */ "
+}
+
 // TestFkFuzzTest is a fuzzer test that works by querying the database concurrently.
 // We have a pre-written set of query templates that we will use, but the data in the queries will
 // be randomly generated. The intent is that we hammer the database as a real-world application would
@@ -603,49 +627,54 @@ func TestFkFuzzTest(t *testing.T) {
 			updateShare:    50,
 		},
 	}
-
-	for _, tt := range testcases {
-		for _, testSharded := range []bool{false, true} {
-			for _, queryFormat := range []QueryFormat{SQLQueries, PreparedStatmentQueries, PreparedStatementPacket} {
-				t.Run(getTestName(tt.name, testSharded)+fmt.Sprintf(" QueryFormat - %v", queryFormat), func(t *testing.T) {
-					mcmp, closer := start(t)
-					defer closer()
-					// Set the correct keyspace to use from VtGates.
-					if testSharded {
-						t.Skip("Skip test since we don't have sharded foreign key support yet")
-						_ = utils.Exec(t, mcmp.VtConn, "use `ks`")
-					} else {
-						_ = utils.Exec(t, mcmp.VtConn, "use `uks`")
+	for _, fkState := range []sqlparser.FkChecksState{sqlparser.FkChecksUnspecified, sqlparser.FkChecksOn, sqlparser.FkChecksOff} {
+		for _, tt := range testcases {
+			for _, testSharded := range []bool{false, true} {
+				for _, queryFormat := range []QueryFormat{SQLQueries, PreparedStatmentQueries, PreparedStatementPacket} {
+					if fkState != sqlparser.FkChecksUnspecified && (queryFormat != SQLQueries || tt.concurrency != 1) {
+						continue
 					}
-					// Ensure that the Vitess database is originally empty
-					ensureDatabaseState(t, mcmp.VtConn, true)
-					ensureDatabaseState(t, mcmp.MySQLConn, true)
+					t.Run(getTestName(tt.name, testSharded)+fmt.Sprintf(" FkState - %v QueryFormat - %v", fkState.String(), queryFormat), func(t *testing.T) {
+						mcmp, closer := start(t)
+						defer closer()
+						// Set the correct keyspace to use from VtGates.
+						if testSharded {
+							t.Skip("Skip test since we don't have sharded foreign key support yet")
+							_ = utils.Exec(t, mcmp.VtConn, "use `ks`")
+						} else {
+							_ = utils.Exec(t, mcmp.VtConn, "use `uks`")
+						}
+
+						// Ensure that the Vitess database is originally empty
+						ensureDatabaseState(t, mcmp.VtConn, true)
+						ensureDatabaseState(t, mcmp.MySQLConn, true)
 
-					// Create the fuzzer.
-					fz := newFuzzer(tt.concurrency, tt.maxValForId, tt.maxValForCol, tt.insertShare, tt.deleteShare, tt.updateShare, queryFormat)
+						// Create the fuzzer.
+						fz := newFuzzer(tt.concurrency, tt.maxValForId, tt.maxValForCol, tt.insertShare, tt.deleteShare, tt.updateShare, queryFormat, fkState)
 
-					// Start the fuzzer.
-					fz.start(t, testSharded)
+						// Start the fuzzer.
+						fz.start(t, testSharded)
 
-					// Wait for the timeForTesting so that the threads continue to run.
-					time.Sleep(tt.timeForTesting)
+						// Wait for the timeForTesting so that the threads continue to run.
+						time.Sleep(tt.timeForTesting)
 
-					fz.stop()
+						fz.stop()
 
-					// We encountered an error while running the fuzzer. Let's print out the information!
-					if fz.firstFailureInfo != nil {
-						log.Errorf("Failing query - %v", fz.firstFailureInfo.queryToFail)
-						for idx, table := range fkTables {
-							log.Errorf("MySQL data for %v -\n%v", table, fz.firstFailureInfo.mysqlState[idx].Rows)
-							log.Errorf("Vitess data for %v -\n%v", table, fz.firstFailureInfo.vitessState[idx].Rows)
+						// We encountered an error while running the fuzzer. Let's print out the information!
+						if fz.firstFailureInfo != nil {
+							log.Errorf("Failing query - %v", fz.firstFailureInfo.queryToFail)
+							for idx, table := range fkTables {
+								log.Errorf("MySQL data for %v -\n%v", table, fz.firstFailureInfo.mysqlState[idx].Rows)
+								log.Errorf("Vitess data for %v -\n%v", table, fz.firstFailureInfo.vitessState[idx].Rows)
+							}
 						}
-					}
 
-					// ensure Vitess database has some data. This ensures not all the commands failed.
-					ensureDatabaseState(t, mcmp.VtConn, false)
-					// Verify the consistency of the data.
-					verifyDataIsCorrect(t, mcmp, tt.concurrency)
-				})
+						// ensure Vitess database has some data. This ensures not all the commands failed.
+						ensureDatabaseState(t, mcmp.VtConn, false)
+						// Verify the consistency of the data.
+						verifyDataIsCorrect(t, mcmp, tt.concurrency)
+					})
+				}
 			}
 		}
 	}
@@ -701,7 +730,7 @@ func verifyDataIsCorrect(t *testing.T, mcmp utils.MySQLCompare, concurrency int)
 			require.NotNil(t, primaryTab)
 			require.NotNil(t, replicaTab)
 			checkReplicationHealthy(t, replicaTab)
-			cluster.WaitForReplicationPos(t, primaryTab, replicaTab, true, 60.0)
+			cluster.WaitForReplicationPos(t, primaryTab, replicaTab, true, 1*time.Minute)
 			primaryConn, err := utils.GetMySQLConn(primaryTab, fmt.Sprintf("vt_%v", keyspace.Name))
 			require.NoError(t, err)
 			replicaConn, err := utils.GetMySQLConn(replicaTab, fmt.Sprintf("vt_%v", keyspace.Name))

go/test/vschemawrapper/vschema_wrapper.go

@@ -40,14 +40,15 @@ import (
 var _ plancontext.VSchema = (*VSchemaWrapper)(nil)
 
 type VSchemaWrapper struct {
-	V             *vindexes.VSchema
-	Keyspace      *vindexes.Keyspace
-	TabletType_   topodatapb.TabletType
-	Dest          key.Destination
-	SysVarEnabled bool
-	Version       plancontext.PlannerVersion
-	EnableViews   bool
-	TestBuilder   func(query string, vschema plancontext.VSchema, keyspace string) (*engine.Plan, error)
+	V                     *vindexes.VSchema
+	Keyspace              *vindexes.Keyspace
+	TabletType_           topodatapb.TabletType
+	Dest                  key.Destination
+	SysVarEnabled         bool
+	ForeignKeyChecksState sqlparser.FkChecksState
+	Version               plancontext.PlannerVersion
+	EnableViews           bool
+	TestBuilder           func(query string, vschema plancontext.VSchema, keyspace string) (*engine.Plan, error)
 }
 
 func (vw *VSchemaWrapper) GetPrepareData(stmtName string) *vtgatepb.PrepareData {
@@ -140,6 +141,10 @@ func (vw *VSchemaWrapper) KeyspaceError(keyspace string) error {
 	return nil
 }
 
+func (vw *VSchemaWrapper) GetForeignKeyChecksState() sqlparser.FkChecksState {
+	return vw.ForeignKeyChecksState
+}
+
 func (vw *VSchemaWrapper) AllKeyspace() ([]*vindexes.Keyspace, error) {
 	if vw.Keyspace == nil {
 		return nil, vterrors.VT13001("keyspace not available")

go/vt/schemadiff/semantics.go

@@ -64,6 +64,10 @@ func (si *declarativeSchemaInformation) KeyspaceError(keyspace string) error {
 	return nil
 }
 
+func (si *declarativeSchemaInformation) GetForeignKeyChecksState() sqlparser.FkChecksState {
+	return sqlparser.FkChecksUnspecified
+}
+
 // addTable adds a fake table with an empty column list
 func (si *declarativeSchemaInformation) addTable(tableName string) {
 	tbl := &vindexes.Table{

go/vt/sqlparser/ast_funcs.go

@@ -357,6 +357,16 @@ func (node *ParsedComments) AddQueryHint(queryHint string) (Comments, error) {
 	return newComments, nil
 }
 
+func (s FkChecksState) String() string {
+	switch s {
+	case FkChecksOff:
+		return "Off"
+	case FkChecksOn:
+		return "On"
+	}
+	return ""
+}
+
 // ParseParams parses the vindex parameter list, pulling out the special-case
 // "owner" parameter
 func (node *VindexSpec) ParseParams() (string, map[string]string) {

go/vt/sqlparser/ast_rewriting.go

@@ -51,6 +51,7 @@ func PrepareAST(
 	selectLimit int,
 	setVarComment string,
 	sysVars map[string]string,
+	fkChecksState FkChecksState,
 	views VSchemaViews,
 ) (*RewriteASTResult, error) {
 	if parameterize {
@@ -59,7 +60,7 @@ func PrepareAST(
 			return nil, err
 		}
 	}
-	return RewriteAST(in, keyspace, selectLimit, setVarComment, sysVars, views)
+	return RewriteAST(in, keyspace, selectLimit, setVarComment, sysVars, fkChecksState, views)
 }
 
 // RewriteAST rewrites the whole AST, replacing function calls and adding column aliases to queries.
@@ -70,9 +71,10 @@ func RewriteAST(
 	selectLimit int,
 	setVarComment string,
 	sysVars map[string]string,
+	fkChecksState FkChecksState,
 	views VSchemaViews,
 ) (*RewriteASTResult, error) {
-	er := newASTRewriter(keyspace, selectLimit, setVarComment, sysVars, views)
+	er := newASTRewriter(keyspace, selectLimit, setVarComment, sysVars, fkChecksState, views)
 	er.shouldRewriteDatabaseFunc = shouldRewriteDatabaseFunc(in)
 	result := SafeRewrite(in, er.rewriteDown, er.rewriteUp)
 	if er.err != nil {
@@ -121,16 +123,18 @@ type astRewriter struct {
 	keyspace      string
 	selectLimit   int
 	setVarComment string
+	fkChecksState FkChecksState
 	sysVars       map[string]string
 	views         VSchemaViews
 }
 
-func newASTRewriter(keyspace string, selectLimit int, setVarComment string, sysVars map[string]string, views VSchemaViews) *astRewriter {
+func newASTRewriter(keyspace string, selectLimit int, setVarComment string, sysVars map[string]string, fkChecksState FkChecksState, views VSchemaViews) *astRewriter {
 	return &astRewriter{
 		bindVars:      &BindVarNeeds{},
 		keyspace:      keyspace,
 		selectLimit:   selectLimit,
 		setVarComment: setVarComment,
+		fkChecksState: fkChecksState,
 		sysVars:       sysVars,
 		views:         views,
 	}
@@ -154,7 +158,7 @@ const (
 )
 
 func (er *astRewriter) rewriteAliasedExpr(node *AliasedExpr) (*BindVarNeeds, error) {
-	inner := newASTRewriter(er.keyspace, er.selectLimit, er.setVarComment, er.sysVars, er.views)
+	inner := newASTRewriter(er.keyspace, er.selectLimit, er.setVarComment, er.sysVars, FkChecksUnspecified, er.views)
 	inner.shouldRewriteDatabaseFunc = er.shouldRewriteDatabaseFunc
 	tmp := SafeRewrite(node.Expr, inner.rewriteDown, inner.rewriteUp)
 	newExpr, ok := tmp.(Expr)
@@ -177,13 +181,19 @@ func (er *astRewriter) rewriteDown(node SQLNode, _ SQLNode) bool {
 
 func (er *astRewriter) rewriteUp(cursor *Cursor) bool {
 	// Add SET_VAR comment to this node if it supports it and is needed
-	if supportOptimizerHint, supportsOptimizerHint := cursor.Node().(SupportOptimizerHint); supportsOptimizerHint && er.setVarComment != "" {
-		newComments, err := supportOptimizerHint.GetParsedComments().AddQueryHint(er.setVarComment)
-		if err != nil {
-			er.err = err
-			return false
+	if supportOptimizerHint, supportsOptimizerHint := cursor.Node().(SupportOptimizerHint); supportsOptimizerHint {
+		if er.setVarComment != "" {
+			newComments, err := supportOptimizerHint.GetParsedComments().AddQueryHint(er.setVarComment)
+			if err != nil {
+				er.err = err
+				return false
+			}
+			supportOptimizerHint.SetComments(newComments)
+		}
+		if er.fkChecksState != FkChecksUnspecified {
+			newComments := supportOptimizerHint.GetParsedComments().SetMySQLSetVarValue(sysvars.ForeignKeyChecks.Name, er.fkChecksState.String())
+			supportOptimizerHint.SetComments(newComments)
 		}
-		supportOptimizerHint.SetComments(newComments)
 	}
 
 	switch node := cursor.Node().(type) {