feat: add e2e tests and fix infinite loop in vtgate :3

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

go/test/endtoend/vtgate/foreign_keys/fk_test.go

@@ -0,0 +1,51 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package foreign_keys
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"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) {
+	mcmp, closer := start(t)
+	defer closer()
+
+	// insert some data.
+	utils.Exec(t, mcmp.VtConn, `insert into t1(id, col) values (100, 123),(10, 12),(1, 13),(1000, 1234)`)
+
+	// Verify that inserting data into a table that has shard scoped foreign keys works.
+	utils.Exec(t, mcmp.VtConn, `insert into t2(id, col) values (100, 125), (1, 132)`)
+	// Verify that insertion fails if the data doesn't follow the fk constraint.
+	_, err := utils.ExecAllowError(t, mcmp.VtConn, `insert into t2(id, col) values (1310, 125)`)
+	require.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, mcmp.VtConn, `insert into t3(id, col) values (100, 100)`)
+	require.ErrorContains(t, err, "VT12002: unsupported cross-shard foreign keys")
+
+	// insert some data in a table with multicol vindex.
+	utils.Exec(t, mcmp.VtConn, `insert into multicol_tbl1(cola, colb, colc, msg) values (100, 'a', 'b', 'msg'), (101, 'c', 'd', 'msg2')`)
+	// Verify that inserting data into a table that has shard scoped multi-column foreign keys works.
+	utils.Exec(t, mcmp.VtConn, `insert into multicol_tbl2(cola, colb, colc, msg) values (100, 'a', 'b', 'msg3')`)
+	// Verify that insertion fails if the data doesn't follow the fk constraint.
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, `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")
+}

go/test/endtoend/vtgate/foreign_keys/main_test.go

@@ -0,0 +1,125 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package foreign_keys
+
+import (
+	_ "embed"
+	"flag"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/test/endtoend/utils"
+
+	"vitess.io/vitess/go/vt/vtgate/planbuilder"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	mysqlParams     mysql.ConnParams
+	shardedKs       = "ks"
+	shardedKsShards = []string{"-19a0", "19a0-20", "20-20c0", "20c0-"}
+	Cell            = "test"
+	//go:embed sharded_schema.sql
+	shardedSchemaSQL string
+
+	//go:embed sharded_vschema.json
+	shardedVSchema string
+)
+
+func TestMain(m *testing.M) {
+	defer cluster.PanicHandler(nil)
+	flag.Parse()
+
+	exitCode := func() int {
+		clusterInstance = cluster.NewCluster(Cell, "localhost")
+		defer clusterInstance.Teardown()
+
+		// Start topo server
+		err := clusterInstance.StartTopo()
+		if err != nil {
+			return 1
+		}
+
+		// Start keyspace
+		sKs := &cluster.Keyspace{
+			Name:      shardedKs,
+			SchemaSQL: shardedSchemaSQL,
+			VSchema:   shardedVSchema,
+		}
+
+		clusterInstance.VtGateExtraArgs = []string{"--schema_change_signal"}
+		clusterInstance.VtTabletExtraArgs = []string{"--queryserver-config-schema-change-signal", "--queryserver-config-schema-change-signal-interval", "0.1"}
+		err = clusterInstance.StartKeyspace(*sKs, shardedKsShards, 0, false)
+		if err != nil {
+			return 1
+		}
+
+		err = clusterInstance.VtctlclientProcess.ExecuteCommand("RebuildVSchemaGraph")
+		if err != nil {
+			return 1
+		}
+
+		// Start vtgate
+		clusterInstance.VtGatePlannerVersion = planbuilder.Gen4 // enable Gen4 planner.
+		err = clusterInstance.StartVtgate()
+		if err != nil {
+			return 1
+		}
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+
+		conn, closer, err := utils.NewMySQL(clusterInstance, shardedKs, shardedSchemaSQL)
+		if err != nil {
+			fmt.Println(err)
+			return 1
+		}
+		defer closer()
+		mysqlParams = conn
+		return m.Run()
+	}()
+	os.Exit(exitCode)
+}
+
+func start(t *testing.T) (utils.MySQLCompare, func()) {
+	mcmp, err := utils.NewMySQLCompare(t, vtParams, mysqlParams)
+	require.NoError(t, err)
+	deleteAll := func() {
+		_, _ = utils.ExecAllowError(t, mcmp.VtConn, "set workload = oltp")
+
+		tables := []string{"t3", "t2", "t1", "multicol_tbl2", "multicol_tbl1"}
+		for _, table := range tables {
+			_, _ = mcmp.ExecAndIgnore("delete from " + table)
+		}
+	}
+
+	deleteAll()
+
+	return mcmp, func() {
+		deleteAll()
+		mcmp.Close()
+		cluster.PanicHandler(t)
+	}
+}

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

@@ -0,0 +1,41 @@
+create table t1
+(
+    id  bigint,
+    col bigint,
+    primary key (id)
+) Engine = InnoDB;
+
+create table t2
+(
+    id    bigint,
+    col   bigint,
+    primary key (id),
+    foreign key (id) references t1 (id)
+) Engine = InnoDB;
+
+create table t3
+(
+    id    bigint,
+    col   bigint,
+    primary key (id),
+    foreign key (col) references t1 (id)
+) Engine = InnoDB;
+
+create table multicol_tbl1
+(
+    cola bigint,
+    colb varbinary(50),
+    colc varchar(50),
+    msg  varchar(50),
+    primary key (cola, colb, colc)
+) Engine = InnoDB;
+
+create table multicol_tbl2
+(
+    cola bigint,
+    colb varbinary(50),
+    colc varchar(50),
+    msg  varchar(50),
+    primary key (cola, colb, colc),
+    foreign key (cola, colb, colc) references multicol_tbl1 (cola, colb, colc)
+) Engine = InnoDB;

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

@@ -0,0 +1,67 @@
+{
+  "sharded": true,
+  "foreignKeyMode": "FK_MANAGED",
+  "vindexes": {
+    "xxhash": {
+      "type": "xxhash"
+    },
+    "multicol_vdx": {
+      "type": "multicol",
+      "params": {
+        "column_count": "3",
+        "column_bytes": "1,3,4",
+        "column_vindex": "hash,binary,unicode_loose_xxhash"
+      }
+    }
+  },
+  "tables": {
+    "t1": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ]
+    },
+    "t2": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ]
+    },
+    "t3": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ]
+    },
+    "multicol_tbl1": {
+      "column_vindexes": [
+        {
+          "columns": [
+            "cola",
+            "colb",
+            "colc"
+          ],
+          "name": "multicol_vdx"
+        }
+      ]
+    },
+    "multicol_tbl2": {
+      "column_vindexes": [
+        {
+          "columns": [
+            "cola",
+            "colb",
+            "colc"
+          ],
+          "name": "multicol_vdx"
+        }
+      ]
+    }
+  }
+}

go/vt/vtgate/planbuilder/operators/fk_verify.go

@@ -87,11 +87,11 @@ func NewFkVerify(ctx *plancontext.PlanningContext, parentFKs []vindexes.ParentFK
 		// Build the offsets
 		var offsets []int
 	outer:
-		for _, column := range fk.ParentColumns {
+		for _, column := range fk.ChildColumns {
 			for idx, insertColumn := range insertColumns {
 				if column.Equal(insertColumn) {
 					offsets = append(offsets, idx)
-					goto outer
+					continue outer
 				}
 			}
 			offsets = append(offsets, -1)

test/config.json

@@ -842,6 +842,15 @@
 			"RetryMax": 1,
 			"Tags": []
 		},
+		"vtgate_foreign_keys": {
+			"File": "unused.go",
+			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/foreign_keys"],
+			"Command": [],
+			"Manual": false,
+			"Shard": "vtgate_gen4",
+			"RetryMax": 2,
+			"Tags": []
+		},
 		"vtgate_gen4": {
 			"File": "unused.go",
 			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/gen4"],