Disable FK checks on source/target table drops

Signed-off-by: Matt Lord <[email protected]>

go/vt/vtctl/workflow/traffic_switcher.go

@@ -510,9 +510,10 @@ func (ts *trafficSwitcher) removeSourceTables(ctx context.Context, removalType T
 				query = fmt.Sprintf("rename table %s.%s TO %s.%s", primaryDbName, tableNameEscaped, primaryDbName, renameName)
 			}
 			_, err = ts.ws.tmc.ExecuteFetchAsDba(ctx, source.GetPrimary().Tablet, false, &tabletmanagerdatapb.ExecuteFetchAsDbaRequest{
-				Query:        []byte(query),
-				MaxRows:      1,
-				ReloadSchema: true,
+				Query:                   []byte(query),
+				MaxRows:                 1,
+				ReloadSchema:            true,
+				DisableForeignKeyChecks: true,
 			})
 			if err != nil {
 				ts.Logger().Errorf("%s: Error removing table %s: %v", source.GetPrimary().String(), tableName, err)
@@ -1067,7 +1068,6 @@ func (ts *trafficSwitcher) dropSourceReverseVReplicationStreams(ctx context.Cont
 }
 
 func (ts *trafficSwitcher) removeTargetTables(ctx context.Context) error {
-	log.Flush()
 	err := ts.ForAllTargets(func(target *MigrationTarget) error {
 		log.Infof("ForAllTargets: %+v", target)
 		for _, tableName := range ts.Tables() {
@@ -1083,12 +1083,12 @@ func (ts *trafficSwitcher) removeTargetTables(ctx context.Context) error {
 			ts.Logger().Infof("%s: Dropping table %s.%s\n",
 				target.GetPrimary().String(), target.GetPrimary().DbName(), tableName)
 			res, err := ts.ws.tmc.ExecuteFetchAsDba(ctx, target.GetPrimary().Tablet, false, &tabletmanagerdatapb.ExecuteFetchAsDbaRequest{
-				Query:        []byte(query),
-				MaxRows:      1,
-				ReloadSchema: true,
+				Query:                   []byte(query),
+				MaxRows:                 1,
+				ReloadSchema:            true,
+				DisableForeignKeyChecks: true,
 			})
 			log.Infof("Removed target table with result: %+v", res)
-			log.Flush()
 			if err != nil {
 				ts.Logger().Errorf("%s: Error removing table %s: %v",
 					target.GetPrimary().String(), tableName, err)

go/vt/vttablet/grpctmclient/client.go

@@ -488,11 +488,12 @@ func (client *Client) ExecuteFetchAsDba(ctx context.Context, tablet *topodatapb.
 	}
 
 	response, err := c.ExecuteFetchAsDba(ctx, &tabletmanagerdatapb.ExecuteFetchAsDbaRequest{
-		Query:          req.Query,
-		DbName:         topoproto.TabletDbName(tablet),
-		MaxRows:        req.MaxRows,
-		DisableBinlogs: req.DisableBinlogs,
-		ReloadSchema:   req.DisableBinlogs,
+		Query:                   req.Query,
+		DbName:                  topoproto.TabletDbName(tablet),
+		MaxRows:                 req.MaxRows,
+		DisableBinlogs:          req.DisableBinlogs,
+		ReloadSchema:            req.DisableBinlogs,
+		DisableForeignKeyChecks: req.DisableForeignKeyChecks,
 	})
 	if err != nil {
 		return nil, err

go/vt/vttablet/tabletmanager/rpc_query.go

@@ -73,21 +73,29 @@ func (tm *TabletManager) ExecuteFetchAsDba(ctx context.Context, req *tabletmanag
 	if err := tm.waitForGrantsToHaveApplied(ctx); err != nil {
 		return nil, err
 	}
-	// get a connection
+	// Get a connection.
 	conn, err := tm.MysqlDaemon.GetDbaConnection(ctx)
 	if err != nil {
 		return nil, err
 	}
 	defer conn.Close()
 
-	// disable binlogs if necessary
+	// Disable binlogs if necessary.
 	if req.DisableBinlogs {
 		_, err := conn.ExecuteFetch("SET sql_log_bin = OFF", 0, false)
 		if err != nil {
 			return nil, err
 		}
 	}
 
+	// Disable FK checks if requested.
+	if req.DisableForeignKeyChecks {
+		_, err := conn.ExecuteFetch("SET SESSION foreign_key_checks = OFF", 0, false)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if req.DbName != "" {
 		// This execute might fail if db does not exist.
 		// Error is ignored because given query might create this database.
@@ -130,7 +138,17 @@ func (tm *TabletManager) ExecuteFetchAsDba(ctx context.Context, req *tabletmanag
 		}
 	}
 
-	// re-enable binlogs if necessary
+	// Re-enable FK checks if necessary.
+	if req.DisableForeignKeyChecks && !conn.IsClosed() {
+		_, err := conn.ExecuteFetch("SET SESSION foreign_key_checks = ON", 0, false)
+		if err != nil {
+			// If we can't reset the FK checks flag,
+			// let's just close the connection.
+			conn.Close()
+		}
+	}
+
+	// Re-enable binlogs if necessary.
 	if req.DisableBinlogs && !conn.IsClosed() {
 		_, err := conn.ExecuteFetch("SET sql_log_bin = ON", 0, false)
 		if err != nil {

go/vt/wrangler/traffic_switcher.go

@@ -49,6 +49,7 @@ import (
 
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
 	querypb "vitess.io/vitess/go/vt/proto/query"
+	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
@@ -1790,7 +1791,12 @@ func (ts *trafficSwitcher) removeSourceTables(ctx context.Context, removalType w
 					source.GetPrimary().String(), source.GetPrimary().DbName(), tableName, source.GetPrimary().DbName(), renameName)
 				query = fmt.Sprintf("rename table %s.%s TO %s.%s", primaryDbName, tableNameEscaped, primaryDbName, renameName)
 			}
-			_, err = ts.wr.ExecuteFetchAsDba(ctx, source.GetPrimary().Alias, query, 1, false, true)
+			_, err = ts.wr.tmc.ExecuteFetchAsDba(ctx, source.GetPrimary().Tablet, false, &tabletmanagerdatapb.ExecuteFetchAsDbaRequest{
+				Query:                   []byte(query),
+				MaxRows:                 1,
+				ReloadSchema:            true,
+				DisableForeignKeyChecks: true,
+			})
 			if err != nil {
 				ts.Logger().Errorf("%s: Error removing table %s: %v", source.GetPrimary().String(), tableName, err)
 				return err
@@ -1896,7 +1902,12 @@ func (ts *trafficSwitcher) removeTargetTables(ctx context.Context) error {
 			query := fmt.Sprintf("drop table %s.%s", primaryDbName, tableName)
 			ts.Logger().Infof("%s: Dropping table %s.%s\n",
 				target.GetPrimary().String(), target.GetPrimary().DbName(), tableName)
-			_, err = ts.wr.ExecuteFetchAsDba(ctx, target.GetPrimary().Alias, query, 1, false, true)
+			_, err = ts.wr.tmc.ExecuteFetchAsDba(ctx, target.GetPrimary().Tablet, false, &tabletmanagerdatapb.ExecuteFetchAsDbaRequest{
+				Query:                   []byte(query),
+				MaxRows:                 1,
+				ReloadSchema:            true,
+				DisableForeignKeyChecks: true,
+			})
 			if err != nil {
 				ts.Logger().Errorf("%s: Error removing table %s: %v",
 					target.GetPrimary().String(), tableName, err)

proto/tabletmanagerdata.proto

@@ -259,6 +259,7 @@ message ExecuteFetchAsDbaRequest {
   uint64 max_rows = 3;
   bool disable_binlogs = 4;
   bool reload_schema = 5;
+  bool disable_foreign_key_checks = 6;
 }
 
 message ExecuteFetchAsDbaResponse {