feat: turn off foreign key checks on updates that have foreign key co…

…lumns being set to non-literal values

Signed-off-by: Manan Gupta <[email protected]>

go/vt/vtgate/planbuilder/operators/ast_to_op.go

@@ -214,7 +214,10 @@ func createOpFromStmt(ctx *plancontext.PlanningContext, stmt sqlparser.Statement
 	// we should augment the semantic analysis to also tell us whether the given query has any cross shard parent foreign keys to validate.
 	// If there are, then we have to run the query with FOREIGN_KEY_CHECKS off because we can't be sure if the DML will succeed on MySQL with the checks on.
 	// So, we should set VerifyAllFKs to true. i.e. we should add `|| ctx.SemTable.RequireForeignKeyChecksOff()` to the below condition.
-	ctx.VerifyAllFKs = verifyAllFKs
+	if verifyAllFKs {
+		// If ctx.VerifyAllFKs is already true we don't want to turn it off.
+		ctx.VerifyAllFKs = verifyAllFKs
+	}
 
 	// From all the parent foreign keys involved, we should remove the one that we need to ignore.
 	err = ctx.SemTable.RemoveParentForeignKey(fkToIgnore)

go/vt/vtgate/planbuilder/operators/update.go

@@ -176,6 +176,11 @@ func createUpdateOperator(ctx *plancontext.PlanningContext, updStmt *sqlparser.U
 		return nil, vterrors.VT12001("multi shard UPDATE with LIMIT")
 	}
 
+	if ctx.SemTable.FKChecksOff {
+		// We have to run the query with FKChecksOff.
+		updStmt.Comments = updStmt.Comments.Prepend("/*+ SET_VAR(foreign_key_checks=OFF) */").Parsed()
+	}
+
 	route := &Route{
 		Source: &Update{
 			QTable:                       qt,
@@ -213,25 +218,6 @@ func buildFkOperator(ctx *plancontext.PlanningContext, updOp ops.Operator, updCl
 	return createFKVerifyOp(ctx, op, updClone, parentFks, restrictChildFks)
 }
 
-func hasNonLiteral(updExprs sqlparser.UpdateExprs, parentFks []vindexes.ParentFKInfo, childFks []vindexes.ChildFKInfo) bool {
-	for _, updateExpr := range updExprs {
-		if sqlparser.IsLiteral(updateExpr.Expr) {
-			continue
-		}
-		for _, parentFk := range parentFks {
-			if parentFk.ChildColumns.FindColumn(updateExpr.Name.Name) >= 0 {
-				return true
-			}
-		}
-		for _, childFk := range childFks {
-			if childFk.ParentColumns.FindColumn(updateExpr.Name.Name) >= 0 {
-				return true
-			}
-		}
-	}
-	return false
-}
-
 // splitChildFks splits the child foreign keys into restrict and cascade list as restrict is handled through Verify operator and cascade is handled through Cascade operator.
 func splitChildFks(fks []vindexes.ChildFKInfo) (restrictChildFks, cascadeChildFks []vindexes.ChildFKInfo) {
 	for _, fk := range fks {
@@ -265,7 +251,7 @@ func createFKCascadeOp(ctx *plancontext.PlanningContext, parentOp ops.Operator,
 			return nil, vterrors.VT13001("ON UPDATE RESTRICT foreign keys should already be filtered")
 		}
 
-		nonLiteralUpdate := hasNonLiteral(updStmt.Exprs, nil, []vindexes.ChildFKInfo{fk})
+		nonLiteralUpdate := semantics.HasNonLiteral(updStmt.Exprs, nil, []vindexes.ChildFKInfo{fk})
 
 		// We need to select all the parent columns for the foreign key constraint, to use in the update of the child table.
 		var selectOffsets []int
@@ -506,7 +492,7 @@ func createFKVerifyOp(ctx *plancontext.PlanningContext, childOp ops.Operator, up
 	}
 
 	// We only support simple expressions in update queries for foreign key verification.
-	if hasNonLiteral(updStmt.Exprs, parentFks, restrictChildFks) {
+	if semantics.HasNonLiteral(updStmt.Exprs, parentFks, restrictChildFks) {
 		return nil, vterrors.VT12001("update expression with non-literal values with foreign key constraints")
 	}
 

go/vt/vtgate/planbuilder/plancontext/planning_context.go

@@ -78,6 +78,7 @@ func CreatePlanningContext(stmt sqlparser.Statement,
 		SkipPredicates:    map[sqlparser.Expr]any{},
 		PlannerVersion:    version,
 		ReservedArguments: map[sqlparser.Expr]string{},
+		VerifyAllFKs:      semTable.FKChecksOff,
 	}, nil
 }
 

go/vt/vtgate/semantics/analyzer.go

@@ -107,7 +107,7 @@ func (a *analyzer) newSemTable(statement sqlparser.Statement, coll collations.ID
 		columns[union] = info.exprs
 	}
 
-	childFks, parentFks, childFkToUpdExprs, err := a.getInvolvedForeignKeys(statement)
+	childFks, parentFks, childFkToUpdExprs, fkChecksOff, err := a.getInvolvedForeignKeys(statement)
 	if err != nil {
 		return nil, err
 	}
@@ -130,6 +130,7 @@ func (a *analyzer) newSemTable(statement sqlparser.Statement, coll collations.ID
 		childForeignKeysInvolved:  childFks,
 		parentForeignKeysInvolved: parentFks,
 		ChildFkToUpdExprs:         childFkToUpdExprs,
+		FKChecksOff:               fkChecksOff,
 	}, nil
 }
 
@@ -321,14 +322,14 @@ func (a *analyzer) noteQuerySignature(node sqlparser.SQLNode) {
 }
 
 // getInvolvedForeignKeys gets the foreign keys that might require taking care off when executing the given statement.
-func (a *analyzer) getInvolvedForeignKeys(statement sqlparser.Statement) (map[TableSet][]vindexes.ChildFKInfo, map[TableSet][]vindexes.ParentFKInfo, map[string]sqlparser.UpdateExprs, error) {
+func (a *analyzer) getInvolvedForeignKeys(statement sqlparser.Statement) (map[TableSet][]vindexes.ChildFKInfo, map[TableSet][]vindexes.ParentFKInfo, map[string]sqlparser.UpdateExprs, bool, error) {
 	// There are only the DML statements that require any foreign keys handling.
 	switch stmt := statement.(type) {
 	case *sqlparser.Delete:
 		// For DELETE statements, none of the parent foreign keys require handling.
 		// So we collect all the child foreign keys.
 		allChildFks, _, err := a.getAllManagedForeignKeys()
-		return allChildFks, nil, nil, err
+		return allChildFks, nil, nil, false, err
 	case *sqlparser.Insert:
 		// For INSERT statements, we have 3 different cases:
 		// 1. REPLACE statement: REPLACE statements are essentially DELETEs and INSERTs rolled into one.
@@ -338,29 +339,68 @@ func (a *analyzer) getInvolvedForeignKeys(statement sqlparser.Statement) (map[Ta
 		// 3. INSERT with ON DUPLICATE KEY UPDATE: This might trigger an update on the columns specified in the ON DUPLICATE KEY UPDATE clause.
 		allChildFks, allParentFKs, err := a.getAllManagedForeignKeys()
 		if err != nil {
-			return nil, nil, nil, err
+			return nil, nil, nil, false, err
 		}
 		if stmt.Action == sqlparser.ReplaceAct {
-			return allChildFks, allParentFKs, nil, nil
+			return allChildFks, allParentFKs, nil, false, nil
 		}
 		if len(stmt.OnDup) == 0 {
-			return nil, allParentFKs, nil, nil
+			return nil, allParentFKs, nil, false, nil
 		}
 		// If only a certain set of columns are being updated, then there might be some child foreign keys that don't need any consideration since their columns aren't being updated.
 		// So, we filter these child foreign keys out. We can't filter any parent foreign keys because the statement will INSERT a row too, which requires validating all the parent foreign keys.
 		updatedChildFks, _, childFkToUpdExprs := a.filterForeignKeysUsingUpdateExpressions(allChildFks, nil, sqlparser.UpdateExprs(stmt.OnDup))
-		return updatedChildFks, allParentFKs, childFkToUpdExprs, nil
+		return updatedChildFks, allParentFKs, childFkToUpdExprs, false, nil
 	case *sqlparser.Update:
 		// For UPDATE queries we get all the parent and child foreign keys, but we can filter some of them out if the columns that they consist off aren't being updated or are set to NULLs.
 		allChildFks, allParentFks, err := a.getAllManagedForeignKeys()
 		if err != nil {
-			return nil, nil, nil, err
+			return nil, nil, nil, false, err
 		}
 		childFks, parentFks, childFkToUpdExprs := a.filterForeignKeysUsingUpdateExpressions(allChildFks, allParentFks, stmt.Exprs)
-		return childFks, parentFks, childFkToUpdExprs, nil
+		fkChecksOff := false
+		if HasNonLiteral(stmt.Exprs, collectParentFksFromMap(parentFks), collectChildFksFromMap(childFks)) {
+			fkChecksOff = true
+		}
+		return childFks, parentFks, childFkToUpdExprs, fkChecksOff, nil
 	default:
-		return nil, nil, nil, nil
+		return nil, nil, nil, false, nil
+	}
+}
+
+func collectParentFksFromMap(parentFkMap map[TableSet][]vindexes.ParentFKInfo) []vindexes.ParentFKInfo {
+	var parentFks []vindexes.ParentFKInfo
+	for _, fkInfos := range parentFkMap {
+		parentFks = append(parentFks, fkInfos...)
+	}
+	return parentFks
+}
+
+func collectChildFksFromMap(childFkMap map[TableSet][]vindexes.ChildFKInfo) []vindexes.ChildFKInfo {
+	var childFks []vindexes.ChildFKInfo
+	for _, fkInfos := range childFkMap {
+		childFks = append(childFks, fkInfos...)
+	}
+	return childFks
+}
+
+func HasNonLiteral(updExprs sqlparser.UpdateExprs, parentFks []vindexes.ParentFKInfo, childFks []vindexes.ChildFKInfo) bool {
+	for _, updateExpr := range updExprs {
+		if sqlparser.IsLiteral(updateExpr.Expr) {
+			continue
+		}
+		for _, parentFk := range parentFks {
+			if parentFk.ChildColumns.FindColumn(updateExpr.Name.Name) >= 0 {
+				return true
+			}
+		}
+		for _, childFk := range childFks {
+			if childFk.ParentColumns.FindColumn(updateExpr.Name.Name) >= 0 {
+				return true
+			}
+		}
 	}
+	return false
 }
 
 // filterForeignKeysUsingUpdateExpressions filters the child and parent foreign key constraints that don't require any validations/cascades given the updated expressions.

go/vt/vtgate/semantics/analyzer_test.go

@@ -1701,7 +1701,7 @@ func TestFilterForeignKeysUsingUpdateExpressions(t *testing.T) {
 			},
 			si: &FakeSI{
 				KsForeignKeyMode: map[string]vschemapb.Keyspace_ForeignKeyMode{
-					"ks": vschemapb.Keyspace_FK_MANAGED,
+					"ks": vschemapb.Keyspace_managed,
 				},
 			},
 		},
@@ -2000,7 +2000,7 @@ func TestGetInvolvedForeignKeys(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			childFks, parentFks, _, err := tt.analyzer.getInvolvedForeignKeys(tt.stmt)
+			childFks, parentFks, _, _, err := tt.analyzer.getInvolvedForeignKeys(tt.stmt)
 			if tt.expectedErr != "" {
 				require.EqualError(t, err, tt.expectedErr)
 				return

go/vt/vtgate/semantics/semantic_state.go

@@ -132,6 +132,7 @@ type (
 		childForeignKeysInvolved  map[TableSet][]vindexes.ChildFKInfo
 		parentForeignKeysInvolved map[TableSet][]vindexes.ParentFKInfo
 		ChildFkToUpdExprs         map[string]sqlparser.UpdateExprs
+		FKChecksOff               bool
 	}
 
 	columnName struct {