add foreign key as part of set statement when reserved connection is …

…needed (#14696)

Signed-off-by: Harshit Gangal <[email protected]>

go/test/endtoend/vtgate/foreignkey/fk_test.go

@@ -1207,3 +1207,31 @@ func TestReplaceWithFK(t *testing.T) {
 	utils.AssertMatches(t, conn, `select * from u_t1`, `[[INT64(1) INT64(1)] [INT64(2) INT64(2)]]`)
 	utils.AssertMatches(t, conn, `select * from u_t2`, `[[INT64(1) NULL] [INT64(2) NULL]]`)
 }
+
+// TestDDLFk tests that table is created with fk constraint when foreign_key_checks is off.
+func TestDDLFk(t *testing.T) {
+	mcmp, closer := start(t)
+	defer closer()
+
+	utils.Exec(t, mcmp.VtConn, `use uks`)
+
+	createTableDDLTemp1 := `
+create table temp1(id bigint auto_increment primary key, col varchar(20) not null,
+foreign key (col) references temp2(col))
+`
+	mcmp.Exec(`set foreign_key_checks = off`)
+	// should be able to create `temp1` table without a `temp2`
+	mcmp.Exec(createTableDDLTemp1)
+
+	createTableDDLTemp2 := `
+create table temp2(id bigint auto_increment primary key, col varchar(20) not null, key (col))
+`
+	// now create `temp2`
+	mcmp.Exec(createTableDDLTemp2)
+
+	// inserting some data with fk constraints on.
+	mcmp.Exec(`set foreign_key_checks = on`)
+	mcmp.Exec(`insert into temp2(col) values('a'), ('b'), ('c') `)
+	mcmp.Exec(`insert into temp1(col) values('a') `)
+	mcmp.ExecAllowAndCompareError(`insert into temp1(col) values('d') `)
+}

go/vt/sqlparser/ast_rewriting.go

@@ -191,7 +191,7 @@ func (er *astRewriter) rewriteUp(cursor *Cursor) bool {
 			supportOptimizerHint.SetComments(newComments)
 		}
 		if er.fkChecksState != nil {
-			newComments := supportOptimizerHint.GetParsedComments().SetMySQLSetVarValue(sysvars.ForeignKeyChecks.Name, FkChecksStateString(er.fkChecksState))
+			newComments := supportOptimizerHint.GetParsedComments().SetMySQLSetVarValue(sysvars.ForeignKeyChecks, FkChecksStateString(er.fkChecksState))
 			supportOptimizerHint.SetComments(newComments)
 		}
 	}

go/vt/sqlparser/ast_rewriting_test.go

@@ -365,7 +365,7 @@ func TestRewrites(in *testing.T) {
 			assert.Equal(tc.rowCount, result.NeedsFuncResult(RowCountName), "should need row count")
 			assert.Equal(tc.udv, len(result.NeedUserDefinedVariables), "count of user defined variables")
 			assert.Equal(tc.autocommit, result.NeedsSysVar(sysvars.Autocommit.Name), "should need :__vtautocommit")
-			assert.Equal(tc.foreignKeyChecks, result.NeedsSysVar(sysvars.ForeignKeyChecks.Name), "should need :__vtforeignKeyChecks")
+			assert.Equal(tc.foreignKeyChecks, result.NeedsSysVar(sysvars.ForeignKeyChecks), "should need :__vtforeignKeyChecks")
 			assert.Equal(tc.clientFoundRows, result.NeedsSysVar(sysvars.ClientFoundRows.Name), "should need :__vtclientFoundRows")
 			assert.Equal(tc.skipQueryPlanCache, result.NeedsSysVar(sysvars.SkipQueryPlanCache.Name), "should need :__vtskipQueryPlanCache")
 			assert.Equal(tc.sqlSelectLimit, result.NeedsSysVar(sysvars.SQLSelectLimit.Name), "should need :__vtsqlSelectLimit")

go/vt/sqlparser/comments.go

@@ -558,7 +558,7 @@ func AllowScatterDirective(stmt Statement) bool {
 func ForeignKeyChecksState(stmt Statement) *bool {
 	cmt, ok := stmt.(Commented)
 	if ok {
-		fkChecksVal := cmt.GetParsedComments().GetMySQLSetVarValue(sysvars.ForeignKeyChecks.Name)
+		fkChecksVal := cmt.GetParsedComments().GetMySQLSetVarValue(sysvars.ForeignKeyChecks)
 		// If the value of the `foreign_key_checks` optimizer hint is something that doesn't make sense,
 		// then MySQL just ignores it and treats it like the case, where it is unspecified. We are choosing
 		// to have the same behaviour here. If the value doesn't match any of the acceptable values, we return nil,

go/vt/sqlparser/comments_test.go

@@ -564,7 +564,7 @@ func TestGetMySQLSetVarValue(t *testing.T) {
 		{
 			name:      "SET_VAR clause in the middle",
 			comments:  []string{"/*+ NO_RANGE_OPTIMIZATION(t3 PRIMARY, f2_idx) SET_VAR(foreign_key_checks=OFF) NO_ICP(t1, t2) */"},
-			valToFind: sysvars.ForeignKeyChecks.Name,
+			valToFind: sysvars.ForeignKeyChecks,
 			want:      "OFF",
 		},
 		{
@@ -582,19 +582,19 @@ func TestGetMySQLSetVarValue(t *testing.T) {
 		{
 			name:      "Multiple SET_VAR clauses",
 			comments:  []string{"/*+ SET_VAR(sort_buffer_size = 16M) */", "/*+ SET_VAR(optimizer_switch = 'mrr_cost_b(ased=of\"f') */", "/*+ SET_VAR( foReiGn_key_checks = On) */"},
-			valToFind: sysvars.ForeignKeyChecks.Name,
+			valToFind: sysvars.ForeignKeyChecks,
 			want:      "",
 		},
 		{
 			name:      "Verify casing",
 			comments:  []string{"/*+ SET_VAR(optimizer_switch = 'mrr_cost_b(ased=of\"f') SET_VAR( foReiGn_key_checks = On) */"},
-			valToFind: sysvars.ForeignKeyChecks.Name,
+			valToFind: sysvars.ForeignKeyChecks,
 			want:      "On",
 		},
 		{
 			name:      "Leading comment is a normal comment",
 			comments:  []string{"/* This is a normal comment */", "/*+ MAX_EXECUTION_TIME(1000) SET_VAR( foreign_key_checks = 1) */"},
-			valToFind: sysvars.ForeignKeyChecks.Name,
+			valToFind: sysvars.ForeignKeyChecks,
 			want:      "1",
 		},
 	}
@@ -619,7 +619,7 @@ func TestSetMySQLSetVarValue(t *testing.T) {
 		{
 			name:           "SET_VAR clause in the middle",
 			comments:       []string{"/*+ NO_RANGE_OPTIMIZATION(t3 PRIMARY, f2_idx) SET_VAR(foreign_key_checks=OFF) NO_ICP(t1, t2) */"},
-			key:            sysvars.ForeignKeyChecks.Name,
+			key:            sysvars.ForeignKeyChecks,
 			value:          "On",
 			commentsWanted: []string{"/*+ NO_RANGE_OPTIMIZATION(t3 PRIMARY, f2_idx) SET_VAR(foreign_key_checks=On) NO_ICP(t1, t2) */"},
 		},
@@ -640,21 +640,21 @@ func TestSetMySQLSetVarValue(t *testing.T) {
 		{
 			name:           "Multiple SET_VAR clauses",
 			comments:       []string{"/*+ SET_VAR(sort_buffer_size = 16M) */", "/*+ SET_VAR(optimizer_switch = 'mrr_cost_b(ased=of\"f') */", "/*+ SET_VAR( foReiGn_key_checks = On) */"},
-			key:            sysvars.ForeignKeyChecks.Name,
+			key:            sysvars.ForeignKeyChecks,
 			value:          "1",
 			commentsWanted: []string{"/*+ SET_VAR(sort_buffer_size = 16M) SET_VAR(foreign_key_checks=1) */", "/*+ SET_VAR(optimizer_switch = 'mrr_cost_b(ased=of\"f') */", "/*+ SET_VAR( foReiGn_key_checks = On) */"},
 		},
 		{
 			name:           "Verify casing",
 			comments:       []string{"/*+ SET_VAR(optimizer_switch = 'mrr_cost_b(ased=of\"f') SET_VAR( foReiGn_key_checks = On) */"},
-			key:            sysvars.ForeignKeyChecks.Name,
+			key:            sysvars.ForeignKeyChecks,
 			value:          "off",
 			commentsWanted: []string{"/*+ SET_VAR(optimizer_switch = 'mrr_cost_b(ased=of\"f') SET_VAR(foReiGn_key_checks=off) */"},
 		},
 		{
 			name:           "Leading comment is a normal comment",
 			comments:       []string{"/* This is a normal comment */", "/*+ MAX_EXECUTION_TIME(1000) SET_VAR( foreign_key_checks = 1) */"},
-			key:            sysvars.ForeignKeyChecks.Name,
+			key:            sysvars.ForeignKeyChecks,
 			value:          "Off",
 			commentsWanted: []string{"/* This is a normal comment */", "/*+ MAX_EXECUTION_TIME(1000) SET_VAR(foreign_key_checks=Off) */"},
 		},

go/vt/sysvars/sysvars.go

@@ -57,6 +57,8 @@ var (
 	off     = "0"
 	utf8mb4 = "'utf8mb4'"
 
+	ForeignKeyChecks = "foreign_key_checks"
+
 	Autocommit                  = SystemVariable{Name: "autocommit", IsBoolean: true, Default: on}
 	Charset                     = SystemVariable{Name: "charset", Default: utf8mb4, IdentifierAsString: true}
 	ClientFoundRows             = SystemVariable{Name: "client_found_rows", IsBoolean: true, Default: off}
@@ -71,7 +73,6 @@ var (
 	TxReadOnly                  = SystemVariable{Name: "tx_read_only", IsBoolean: true, Default: off}
 	Workload                    = SystemVariable{Name: "workload", IdentifierAsString: true}
 	QueryTimeout                = SystemVariable{Name: "query_timeout"}
-	ForeignKeyChecks            = SystemVariable{Name: "foreign_key_checks", IsBoolean: true, SupportSetVar: true}
 
 	// Online DDL
 	DDLStrategy      = SystemVariable{Name: "ddl_strategy", IdentifierAsString: true}
@@ -105,7 +106,6 @@ var (
 		ReadAfterWriteTimeOut,
 		SessionTrackGTIDs,
 		QueryTimeout,
-		ForeignKeyChecks,
 	}
 
 	ReadOnly = []SystemVariable{
@@ -188,6 +188,7 @@ var (
 		{Name: "end_markers_in_json", IsBoolean: true, SupportSetVar: true},
 		{Name: "eq_range_index_dive_limit", SupportSetVar: true},
 		{Name: "explicit_defaults_for_timestamp"},
+		{Name: ForeignKeyChecks, IsBoolean: true, SupportSetVar: true},
 		{Name: "group_concat_max_len", SupportSetVar: true},
 		{Name: "information_schema_stats_expiry"},
 		{Name: "max_heap_table_size", SupportSetVar: true},

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -275,10 +275,6 @@ func (t *noopVCursor) SetAutocommit(context.Context, bool) error {
 	panic("implement me")
 }
 
-func (t *noopVCursor) SetSessionForeignKeyChecks(ctx context.Context, foreignKeyChecks bool) error {
-	panic("implement me")
-}
-
 func (t *noopVCursor) SetClientFoundRows(context.Context, bool) error {
 	panic("implement me")
 }
@@ -747,10 +743,6 @@ func (f *loggingVCursor) SetAutocommit(context.Context, bool) error {
 	panic("implement me")
 }
 
-func (f *loggingVCursor) SetSessionForeignKeyChecks(ctx context.Context, foreignKeyChecks bool) error {
-	panic("implement me")
-}
-
 func (f *loggingVCursor) SetClientFoundRows(context.Context, bool) error {
 	panic("implement me")
 }

go/vt/vtgate/engine/primitive.go

@@ -154,7 +154,6 @@ type (
 
 		SetAutocommit(ctx context.Context, autocommit bool) error
 		SetClientFoundRows(context.Context, bool) error
-		SetSessionForeignKeyChecks(ctx context.Context, autocommit bool) error
 		SetSkipQueryPlanCache(context.Context, bool) error
 		SetSQLSelectLimit(int64) error
 		SetTransactionMode(vtgatepb.TransactionMode)

go/vt/vtgate/engine/set.go

@@ -445,8 +445,6 @@ func (svss *SysVarSetAware) Execute(ctx context.Context, vcursor VCursor, env *e
 	switch svss.Name {
 	case sysvars.Autocommit.Name:
 		err = svss.setBoolSysVar(ctx, env, vcursor.Session().SetAutocommit)
-	case sysvars.ForeignKeyChecks.Name:
-		err = svss.setBoolSysVar(ctx, env, vcursor.Session().SetSessionForeignKeyChecks)
 	case sysvars.ClientFoundRows.Name:
 		err = svss.setBoolSysVar(ctx, env, vcursor.Session().SetClientFoundRows)
 	case sysvars.SkipQueryPlanCache.Name:

go/vt/vtgate/executor_set_test.go

@@ -319,9 +319,12 @@ func TestExecutorSetOp(t *testing.T) {
 		sysVars: map[string]string{"sql_quote_show_create": "0"},
 		result:  returnResult("sql_quote_show_create", "int64", "0"),
 	}, {
-		in:      "set foreign_key_checks = 1",
-		sysVars: map[string]string{"foreign_key_checks": "1"},
-		result:  returnNoResult("foreign_key_checks", "int64"),
+		in:     "set foreign_key_checks = 1",
+		result: returnNoResult("foreign_key_checks", "int64"),
+	}, {
+		in:      "set foreign_key_checks = 0",
+		sysVars: map[string]string{"foreign_key_checks": "0"},
+		result:  returnResult("foreign_key_checks", "int64", "0"),
 	}, {
 		in:      "set unique_checks = 0",
 		sysVars: map[string]string{"unique_checks": "0"},

go/vt/vtgate/planbuilder/operators/insert.go

@@ -134,6 +134,7 @@ func createOperatorFromInsert(ctx *plancontext.PlanningContext, ins *sqlparser.I
 	whereExpr := getWhereCondExpr(append(uniqKeyCompExprs, pkCompExpr))
 
 	delStmt := &sqlparser.Delete{
+		Comments:   ins.Comments,
 		TableExprs: sqlparser.TableExprs{sqlparser.CloneRefOfAliasedTableExpr(ins.Table)},
 		Where:      sqlparser.NewWhere(sqlparser.WhereClause, whereExpr),
 	}

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

@@ -442,7 +442,7 @@ func buildChildUpdOpForCascade(ctx *plancontext.PlanningContext, fk vindexes.Chi
 	// Because we could be updating the child to a non-null value,
 	// We have to run with foreign key checks OFF because the parent isn't guaranteed to have
 	// the data being updated to.
-	parsedComments := (&sqlparser.ParsedComments{}).SetMySQLSetVarValue(sysvars.ForeignKeyChecks.Name, "OFF").Parsed()
+	parsedComments := (&sqlparser.ParsedComments{}).SetMySQLSetVarValue(sysvars.ForeignKeyChecks, "OFF").Parsed()
 	childUpdStmt := &sqlparser.Update{
 		Comments:   parsedComments,
 		Exprs:      childUpdateExprs,
@@ -515,7 +515,7 @@ func buildChildUpdOpForSetNull(
 func getParsedCommentsForFkChecks(ctx *plancontext.PlanningContext) (parsedComments *sqlparser.ParsedComments) {
 	fkState := ctx.VSchema.GetForeignKeyChecksState()
 	if fkState != nil && *fkState {
-		parsedComments = parsedComments.SetMySQLSetVarValue(sysvars.ForeignKeyChecks.Name, "ON").Parsed()
+		parsedComments = parsedComments.SetMySQLSetVarValue(sysvars.ForeignKeyChecks, "ON").Parsed()
 	}
 	return parsedComments
 }

go/vt/vtgate/planbuilder/testdata/foreignkey_checks_on_cases.json

@@ -1772,7 +1772,7 @@
                   "Sharded": false
                 },
                 "TargetTabletType": "PRIMARY",
-                "Query": "delete from u_tbl1 where (id) in ((1))",
+                "Query": "delete /*+ SET_VAR(foreign_key_checks=On) */ from u_tbl1 where (id) in ((1))",
                 "Table": "u_tbl1"
               }
             ]

go/vt/vtgate/planbuilder/update.go

@@ -51,7 +51,7 @@ func gen4UpdateStmtPlanner(
 	if ctx.SemTable.HasNonLiteralForeignKeyUpdate(updStmt.Exprs) {
 		// Since we are running the query with foreign key checks off, we have to verify all the foreign keys validity on vtgate.
 		ctx.VerifyAllFKs = true
-		updStmt.SetComments(updStmt.GetParsedComments().SetMySQLSetVarValue(sysvars.ForeignKeyChecks.Name, "OFF"))
+		updStmt.SetComments(updStmt.GetParsedComments().SetMySQLSetVarValue(sysvars.ForeignKeyChecks, "OFF"))
 	}
 
 	// Remove all the foreign keys that don't require any handling.

go/vt/vtgate/safe_session.go

@@ -572,6 +572,26 @@ func (session *SafeSession) TimeZone() *time.Location {
 	return loc
 }
 
+// ForeignKeyChecks returns the foreign_key_checks stored in system_variables map in the session.
+func (session *SafeSession) ForeignKeyChecks() *bool {
+	session.mu.Lock()
+	fkVal, ok := session.SystemVariables[sysvars.ForeignKeyChecks]
+	session.mu.Unlock()
+
+	if !ok {
+		return nil
+	}
+	switch strings.ToLower(fkVal) {
+	case "off", "0":
+		fkCheckBool := false
+		return &fkCheckBool
+	case "on", "1":
+		fkCheckBool := true
+		return &fkCheckBool
+	}
+	return nil
+}
+
 // SetOptions sets the options
 func (session *SafeSession) SetOptions(options *querypb.ExecuteOptions) {
 	session.mu.Lock()

go/vt/vtgate/vcursor_impl.go

@@ -43,7 +43,6 @@ import (
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/srvtopo"
-	"vitess.io/vitess/go/vt/sysvars"
 	"vitess.io/vitess/go/vt/topo"
 	topoprotopb "vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/topotools"
@@ -832,16 +831,6 @@ func (vc *vcursorImpl) SetAutocommit(ctx context.Context, autocommit bool) error
 	return nil
 }
 
-// SetSessionForeignKeyChecks implements the SessionActions interface
-func (vc *vcursorImpl) SetSessionForeignKeyChecks(ctx context.Context, foreignKeyChecks bool) error {
-	if foreignKeyChecks {
-		vc.safeSession.SetSystemVariable(sysvars.ForeignKeyChecks.Name, "1")
-	} else {
-		vc.safeSession.SetSystemVariable(sysvars.ForeignKeyChecks.Name, "0")
-	}
-	return nil
-}
-
 // SetQueryTimeout implements the SessionActions interface
 func (vc *vcursorImpl) SetQueryTimeout(maxExecutionTime int64) {
 	vc.safeSession.QueryTimeout = maxExecutionTime
@@ -1365,17 +1354,7 @@ func (vc *vcursorImpl) UpdateForeignKeyChecksState(fkStateFromQuery *bool) {
 		return
 	}
 	// If the query doesn't have anything, then we consult the session state.
-	fkVal, isPresent := vc.safeSession.SystemVariables[sysvars.ForeignKeyChecks.Name]
-	if isPresent {
-		switch strings.ToLower(fkVal) {
-		case "on", "1":
-			val := true
-			vc.fkChecksState = &val
-		case "off", "0":
-			val := false
-			vc.fkChecksState = &val
-		}
-	}
+	vc.fkChecksState = vc.safeSession.ForeignKeyChecks()
 }
 
 // GetForeignKeyChecksState gets the stored foreign key checks state in the vcursor.