Foreign Keys: UPDATE planning (#13762)

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
Signed-off-by: Andres Taylor <andres@planetscale.com>
Co-authored-by: Andres Taylor <andres@planetscale.com>

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

@@ -19,13 +19,13 @@ package foreignkey
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/assert"
 
 	"vitess.io/vitess/go/test/endtoend/utils"
 )
 
-// TestInsertions tests that insertions work as expected when foreign key management is enabled in Vitess.
-func TestInsertions(t *testing.T) {
+// TestInsertWithFK tests that insertions work as expected when foreign key management is enabled in Vitess.
+func TestInsertWithFK(t *testing.T) {
 	conn, closer := start(t)
 	defer closer()
 
@@ -37,11 +37,11 @@ func TestInsertions(t *testing.T) {
 
 	// Verify that insertion fails if the data doesn't follow the fk constraint.
 	_, err := utils.ExecAllowError(t, conn, `insert into t2(id, col) values (1310, 125)`)
-	require.ErrorContains(t, err, "Cannot add or update a child row: a foreign key constraint fails")
+	assert.ErrorContains(t, err, "Cannot add or update a child row: a foreign key constraint fails")
 
 	// Verify that insertion fails if the table has cross-shard foreign keys (even if the data follows the constraints).
 	_, err = utils.ExecAllowError(t, conn, `insert into t3(id, col) values (100, 100)`)
-	require.ErrorContains(t, err, "VT12002: unsupported: cross-shard foreign keys")
+	assert.ErrorContains(t, err, "VT12002: unsupported: cross-shard foreign keys")
 
 	// insert some data in a table with multicol vindex.
 	utils.Exec(t, conn, `insert into multicol_tbl1(cola, colb, colc, msg) values (100, 'a', 'b', 'msg'), (101, 'c', 'd', 'msg2')`)
@@ -51,11 +51,11 @@ func TestInsertions(t *testing.T) {
 
 	// Verify that insertion fails if the data doesn't follow the fk constraint.
 	_, err = utils.ExecAllowError(t, conn, `insert into multicol_tbl2(cola, colb, colc, msg) values (103, 'c', 'd', 'msg2')`)
-	require.ErrorContains(t, err, "Cannot add or update a child row: a foreign key constraint fails")
+	assert.ErrorContains(t, err, "Cannot add or update a child row: a foreign key constraint fails")
 }
 
-// TestDeletions tests that deletions work as expected when foreign key management is enabled in Vitess.
-func TestDeletions(t *testing.T) {
+// TestDeleteWithFK tests that deletions work as expected when foreign key management is enabled in Vitess.
+func TestDeleteWithFK(t *testing.T) {
 	conn, closer := start(t)
 	defer closer()
 
@@ -68,17 +68,53 @@ func TestDeletions(t *testing.T) {
 
 	// child foreign key is shard scoped. Query will fail at mysql due to On Delete Restrict.
 	_, err := utils.ExecAllowError(t, conn, `delete from t2 where col = 132`)
-	require.ErrorContains(t, err, "Cannot delete or update a parent row: a foreign key constraint fails")
+	assert.ErrorContains(t, err, "Cannot delete or update a parent row: a foreign key constraint fails")
 
 	// child row does not exist so query will succeed.
 	qr := utils.Exec(t, conn, `delete from t2 where col = 125`)
-	require.EqualValues(t, 1, qr.RowsAffected)
+	assert.EqualValues(t, 1, qr.RowsAffected)
 
 	// table's child foreign key has cross shard fk, so query will fail at vtgate.
 	_, err = utils.ExecAllowError(t, conn, `delete from t1 where id = 42`)
-	require.ErrorContains(t, err, "VT12002: unsupported: foreign keys management at vitess (errno 1235) (sqlstate 42000)")
+	assert.ErrorContains(t, err, "VT12002: unsupported: foreign keys management at vitess (errno 1235) (sqlstate 42000)")
 
 	// child foreign key is cascade, so query will fail at vtgate.
 	_, err = utils.ExecAllowError(t, conn, `delete from multicol_tbl1 where cola = 100`)
-	require.ErrorContains(t, err, "VT12002: unsupported: foreign keys management at vitess (errno 1235) (sqlstate 42000)")
+	assert.ErrorContains(t, err, "VT12002: unsupported: foreign keys management at vitess (errno 1235) (sqlstate 42000)")
+}
+
+// TestUpdations tests that update work as expected when foreign key management is enabled in Vitess.
+func TestUpdateWithFK(t *testing.T) {
+	conn, closer := start(t)
+	defer closer()
+
+	// insert some data.
+	utils.Exec(t, conn, `insert into t1(id, col) values (100, 123),(10, 12),(1, 13),(1000, 1234)`)
+	utils.Exec(t, conn, `insert into t2(id, col, mycol) values (100, 125, 'foo'), (1, 132, 'bar')`)
+	utils.Exec(t, conn, `insert into t4(id, col, t2_mycol) values (1, 321, 'bar')`)
+	utils.Exec(t, conn, `insert into t5(pk, sk, col1) values (1, 1, 1),(2, 1, 1),(3, 1, 10),(4, 1, 20),(5, 1, 30),(6, 1, 40)`)
+	utils.Exec(t, conn, `insert into t6(pk, sk, col1) values (10, 1, 1), (20, 1, 20)`)
+
+	// parent foreign key is shard scoped and value is not updated. Query will succeed.
+	_ = utils.Exec(t, conn, `update t4 set t2_mycol = 'bar' where id = 1`)
+
+	// parent foreign key is shard scoped and value does not exists in parent table. Query will fail at mysql due to On Update Restrict.
+	_, err := utils.ExecAllowError(t, conn, `update t4 set t2_mycol = 'foo' where id = 1`)
+	assert.ErrorContains(t, err, "Cannot add or update a child row: a foreign key constraint fails")
+
+	// updating column which does not have foreign key constraint, so query will succeed.
+	qr := utils.Exec(t, conn, `update t4 set col = 20 where id = 1`)
+	assert.EqualValues(t, 1, qr.RowsAffected)
+
+	// child table have cascade which is cross shard. Query will fail at vtgate.
+	_, err = utils.ExecAllowError(t, conn, `update t2 set col = 125 where id = 100`)
+	assert.ErrorContains(t, err, "VT12002: unsupported: foreign keys management at vitess (errno 1235) (sqlstate 42000)")
+
+	// updating column which does not have foreign key constraint, so query will succeed.
+	_ = utils.Exec(t, conn, `update t2 set mycol = 'baz' where id = 100`)
+	assert.EqualValues(t, 1, qr.RowsAffected)
+
+	// child table have restrict in shard scoped and value exists in parent table.
+	_ = utils.Exec(t, conn, `update t6 set col1 = 40 where pk = 20`)
+	assert.EqualValues(t, 1, qr.RowsAffected)
 }

go/test/endtoend/vtgate/foreignkey/sharded_schema.sql

@@ -9,14 +9,17 @@ create table t2
 (
     id    bigint,
     col   bigint,
+    mycol varchar(50),
     primary key (id),
+    index(id, mycol),
+    index(id, col),
     foreign key (id) references t1 (id) on delete restrict
 ) Engine = InnoDB;
 
 create table t3
 (
-    id    bigint,
-    col   bigint,
+    id  bigint,
+    col bigint,
     primary key (id),
     foreign key (col) references t1 (id) on delete restrict
 ) Engine = InnoDB;
@@ -42,8 +45,30 @@ create table multicol_tbl2
 
 create table t4
 (
-    id    bigint,
-    col   bigint,
+    id       bigint,
+    col      bigint,
+    t2_mycol varchar(50),
+    t2_col   bigint,
     primary key (id),
-    foreign key (id) references t2 (id) on delete restrict
+    foreign key (id) references t2 (id) on delete restrict,
+    foreign key (id, t2_mycol) references t2 (id, mycol) on update restrict,
+    foreign key (id, t2_col) references t2 (id, col) on update cascade
+) Engine = InnoDB;
+
+create table t5
+(
+    pk   bigint,
+    sk   bigint,
+    col1 varchar(50),
+    primary key (pk),
+    index(sk, col1)
+) Engine = InnoDB;
+
+create table t6
+(
+    pk   bigint,
+    sk   bigint,
+    col1 varchar(50),
+    primary key (pk),
+    foreign key (sk, col1) references t5 (sk, col1) on delete restrict on update restrict
 ) Engine = InnoDB;

go/test/endtoend/vtgate/foreignkey/sharded_vschema.json

@@ -47,6 +47,22 @@
         }
       ]
     },
+    "t5": {
+      "column_vindexes": [
+        {
+          "column": "sk",
+          "name": "xxhash"
+        }
+      ]
+    },
+    "t6": {
+      "column_vindexes": [
+        {
+          "column": "sk",
+          "name": "xxhash"
+        }
+      ]
+    },
     "multicol_tbl1": {
       "column_vindexes": [
         {

go/vt/vtgate/planbuilder/delete.go

@@ -111,7 +111,7 @@ func fkManagementNotRequired(vschema plancontext.VSchema, vTables []*vindexes.Ta
 		if ksMode != vschemapb.Keyspace_FK_MANAGED {
 			continue
 		}
-		childFks := vTable.ChildFKsNeedsHandling()
+		childFks := vTable.ChildFKsNeedsHandling(vindexes.DeleteAction)
 		if len(childFks) > 0 {
 			return false
 		}

go/vt/vtgate/planbuilder/operators/ast2op.go

@@ -160,6 +160,21 @@ func createOperatorFromUpdate(ctx *plancontext.PlanningContext, updStmt *sqlpars
 		Routing: routing,
 	}
 
+	ksMode, err := ctx.VSchema.ForeignKeyMode(vindexTable.Keyspace.Name)
+	if err != nil {
+		return nil, err
+	}
+	if ksMode == vschemapb.Keyspace_FK_MANAGED {
+		parentFKs := vindexTable.ParentFKsNeedsHandling()
+		childFks := vindexTable.ChildFKsNeedsHandling(vindexes.UpdateAction)
+		if (len(childFks) > 0 || len(parentFKs) > 0) &&
+			ColumnModified(updStmt.Exprs, func(expr *sqlparser.UpdateExpr) ([]vindexes.ParentFKInfo, []vindexes.ChildFKInfo) {
+				return parentFKs, childFks
+			}) {
+			return nil, vterrors.VT12003()
+		}
+	}
+
 	subq, err := createSubqueryFromStatement(ctx, updStmt)
 	if err != nil {
 		return nil, err
@@ -171,6 +186,24 @@ func createOperatorFromUpdate(ctx *plancontext.PlanningContext, updStmt *sqlpars
 	return subq, nil
 }
 
+// ColumnModified checks if any column in the parent table is being updated which has a child foreign key.
+func ColumnModified(exprs sqlparser.UpdateExprs, getFks func(expr *sqlparser.UpdateExpr) ([]vindexes.ParentFKInfo, []vindexes.ChildFKInfo)) bool {
+	for _, updateExpr := range exprs {
+		parentFKs, childFks := getFks(updateExpr)
+		for _, childFk := range childFks {
+			if childFk.ParentColumns.FindColumn(updateExpr.Name.Name) >= 0 {
+				return true
+			}
+		}
+		for _, parentFk := range parentFKs {
+			if parentFk.ChildColumns.FindColumn(updateExpr.Name.Name) >= 0 {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func createOperatorFromDelete(ctx *plancontext.PlanningContext, deleteStmt *sqlparser.Delete) (ops.Operator, error) {
 	tableInfo, qt, err := createQueryTableForDML(ctx, deleteStmt.TableExprs[0], deleteStmt.Where)
 	if err != nil {
@@ -197,7 +230,7 @@ func createOperatorFromDelete(ctx *plancontext.PlanningContext, deleteStmt *sqlp
 		return nil, err
 	}
 	if ksMode == vschemapb.Keyspace_FK_MANAGED {
-		childFks := vindexTable.ChildFKsNeedsHandling()
+		childFks := vindexTable.ChildFKsNeedsHandling(vindexes.DeleteAction)
 		if len(childFks) > 0 {
 			return nil, vterrors.VT12003()
 		}
@@ -279,7 +312,7 @@ func createOperatorFromInsert(ctx *plancontext.PlanningContext, ins *sqlparser.I
 		return nil, err
 	}
 	if ksMode == vschemapb.Keyspace_FK_MANAGED {
-		parentFKs := vindexTable.CrossShardParentFKs()
+		parentFKs := vindexTable.ParentFKsNeedsHandling()
 		if len(parentFKs) > 0 {
 			return nil, vterrors.VT12002()
 		}

go/vt/vtgate/planbuilder/plan_test.go

@@ -99,19 +99,56 @@ func TestPlan(t *testing.T) {
 
 // TestForeignKeyPlanning tests the planning of foreign keys in a managed mode by Vitess.
 func TestForeignKeyPlanning(t *testing.T) {
+	vschema := loadSchema(t, "vschemas/schema.json", true)
+	setFks(t, vschema)
 	vschemaWrapper := &vschemawrapper.VSchemaWrapper{
-		V: loadSchema(t, "vschemas/schema.json", true),
-		// Set the keyspace with foreign keys enabled as the default.
-		Keyspace: &vindexes.Keyspace{
-			Name:    "user_fk_allow",
-			Sharded: true,
-		},
+		V: vschema,
 	}
+
 	testOutputTempDir := makeTestOutput(t)
 
 	testFile(t, "foreignkey_cases.json", testOutputTempDir, vschemaWrapper, false)
 }
 
+func setFks(t *testing.T, vschema *vindexes.VSchema) {
+	if vschema.Keyspaces["sharded_fk_allow"] != nil {
+		// FK from multicol_tbl2 referencing multicol_tbl1 that is shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "multicol_tbl2", createFkDefinition([]string{"colb", "cola", "x", "colc", "y"}, "multicol_tbl1", []string{"colb", "cola", "y", "colc", "x"}, sqlparser.Cascade, sqlparser.Cascade))
+
+		// FK from tbl2 referencing tbl1 that is shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl2", createFkDefinition([]string{"col2"}, "tbl1", []string{"col1"}, sqlparser.Restrict, sqlparser.Restrict))
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl2", createFkDefinition([]string{"col2", "col"}, "tbl1", []string{"col1", "col"}, sqlparser.Restrict, sqlparser.Restrict))
+		// FK from tbl3 referencing tbl1 that is not shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl3", createFkDefinition([]string{"coly"}, "tbl1", []string{"t1col1"}, sqlparser.DefaultAction, sqlparser.DefaultAction))
+		// FK from tbl10 referencing tbl2 that is shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl10", createFkDefinition([]string{"sk", "col"}, "tbl2", []string{"col2", "col"}, sqlparser.Restrict, sqlparser.Restrict))
+		// FK from tbl10 referencing tbl3 that is not shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl10", createFkDefinition([]string{"col"}, "tbl3", []string{"col"}, sqlparser.Restrict, sqlparser.Restrict))
+
+		// FK from tbl4 referencing tbl5 that is shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl4", createFkDefinition([]string{"col4"}, "tbl5", []string{"col5"}, sqlparser.SetNull, sqlparser.SetNull))
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl4", createFkDefinition([]string{"t4col4"}, "tbl5", []string{"t5col5"}, sqlparser.SetNull, sqlparser.SetNull))
+
+		// FK from tbl6 referencing tbl7 that is shard scoped.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl6", createFkDefinition([]string{"col6"}, "tbl7", []string{"col7"}, sqlparser.NoAction, sqlparser.NoAction))
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl6", createFkDefinition([]string{"t6col6"}, "tbl7", []string{"t7col7"}, sqlparser.NoAction, sqlparser.NoAction))
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl6", createFkDefinition([]string{"t6col62"}, "tbl7", []string{"t7col72"}, sqlparser.NoAction, sqlparser.NoAction))
+	}
+	if vschema.Keyspaces["unsharded_fk_allow"] != nil {
+		// u_tbl2(col2)  -> u_tbl1(col1)  Cascade.
+		// u_tbl3(col2)  -> u_tbl2(col2)  Cascade Null.
+		// u_tbl4(col41) -> u_tbl1(col14) Restrict.
+		// u_tbl4(col4)  -> u_tbl3(col3)  Restrict.
+		// u_tbl6(col6)  -> u_tbl5(col5)  Restrict.
+
+		_ = vschema.AddForeignKey("unsharded_fk_allow", "u_tbl2", createFkDefinition([]string{"col2"}, "u_tbl1", []string{"col1"}, sqlparser.Cascade, sqlparser.Cascade))
+		_ = vschema.AddForeignKey("unsharded_fk_allow", "u_tbl3", createFkDefinition([]string{"col3"}, "u_tbl2", []string{"col2"}, sqlparser.SetNull, sqlparser.SetNull))
+		_ = vschema.AddForeignKey("unsharded_fk_allow", "u_tbl4", createFkDefinition([]string{"col41"}, "u_tbl1", []string{"col14"}, sqlparser.NoAction, sqlparser.NoAction))
+		_ = vschema.AddForeignKey("unsharded_fk_allow", "u_tbl4", createFkDefinition([]string{"col4"}, "u_tbl3", []string{"col3"}, sqlparser.Restrict, sqlparser.Restrict))
+		_ = vschema.AddForeignKey("unsharded_fk_allow", "u_tbl6", createFkDefinition([]string{"col6"}, "u_tbl5", []string{"col5"}, sqlparser.DefaultAction, sqlparser.DefaultAction))
+	}
+}
+
 func TestSystemTables57(t *testing.T) {
 	// first we move everything to use 5.7 logic
 	oldVer := servenv.MySQLServerVersion()
@@ -436,23 +473,6 @@ func loadSchema(t testing.TB, filename string, setCollation bool) *vindexes.VSch
 			}
 		}
 	}
-	if vschema.Keyspaces["user_fk_allow"] != nil {
-		// FK from multicol_tbl2 referencing multicol_tbl1 that is shard scoped.
-		err = vschema.AddForeignKey("user_fk_allow", "multicol_tbl2", createFkDefinition([]string{"colb", "cola", "x", "colc", "y"}, "multicol_tbl1", []string{"colb", "cola", "y", "colc", "x"}, sqlparser.Cascade, sqlparser.Cascade))
-		require.NoError(t, err)
-		// FK from tbl2 referencing tbl1 that is shard scoped.
-		err = vschema.AddForeignKey("user_fk_allow", "tbl2", createFkDefinition([]string{"col2"}, "tbl1", []string{"col1"}, sqlparser.Restrict, sqlparser.Restrict))
-		require.NoError(t, err)
-		// FK from tbl3 referencing tbl1 that is not shard scoped.
-		err = vschema.AddForeignKey("user_fk_allow", "tbl3", createFkDefinition([]string{"coly"}, "tbl1", []string{"col1"}, sqlparser.DefaultAction, sqlparser.DefaultAction))
-		require.NoError(t, err)
-		// FK from tbl4 referencing tbl5 that is shard scoped.
-		err = vschema.AddForeignKey("user_fk_allow", "tbl4", createFkDefinition([]string{"col4"}, "tbl5", []string{"col5"}, sqlparser.SetNull, sqlparser.SetNull))
-		require.NoError(t, err)
-		// FK from tbl6 referencing tbl7 that is shard scoped.
-		err = vschema.AddForeignKey("user_fk_allow", "tbl6", createFkDefinition([]string{"col6"}, "tbl7", []string{"col7"}, sqlparser.NoAction, sqlparser.NoAction))
-		require.NoError(t, err)
-	}
 	return vschema
 }