Foreign Keys: INSERT planning (#13676)

Signed-off-by: Harshit Gangal <[email protected]>
Signed-off-by: Manan Gupta <[email protected]>
Co-authored-by: Harshit Gangal <[email protected]>

go/test/endtoend/vtgate/foreignkey/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 foreignkey
+
+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) {
+	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)`)
+
+	// Verify that inserting data into a table that has shard scoped foreign keys works.
+	utils.Exec(t, conn, `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, 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")
+	// 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")
+
+	// 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')`)
+	// Verify that inserting data into a table that has shard scoped multi-column foreign keys works.
+	utils.Exec(t, conn, `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, 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")
+}

go/test/endtoend/vtgate/foreignkey/main_test.go

@@ -0,0 +1,106 @@
+/*
+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 foreignkey
+
+import (
+	"context"
+	_ "embed"
+	"flag"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/test/endtoend/utils"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	shardedKs       = "ks"
+	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,
+		}
+
+		err = clusterInstance.StartKeyspace(*sKs, []string{"-80", "80-"}, 0, false)
+		if err != nil {
+			return 1
+		}
+
+		// Start vtgate
+		err = clusterInstance.StartVtgate()
+		if err != nil {
+			return 1
+		}
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+
+		return m.Run()
+	}()
+	os.Exit(exitCode)
+}
+
+func start(t *testing.T) (*mysql.Conn, func()) {
+	conn, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+
+	deleteAll := func() {
+		tables := []string{"t3", "t2", "t1", "multicol_tbl2", "multicol_tbl1"}
+		for _, table := range tables {
+			_ = utils.Exec(t, conn, "delete from "+table)
+		}
+	}
+
+	deleteAll()
+
+	return conn, func() {
+		deleteAll()
+		conn.Close()
+		cluster.PanicHandler(t)
+	}
+}

go/test/endtoend/vtgate/foreignkey/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/foreignkey/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/sqlparser/ast_funcs.go

@@ -2472,3 +2472,33 @@ func (ty KillType) ToString() string {
 		return ConnectionStr
 	}
 }
+
+// Indexes returns true, if the list of columns contains all the elements in the other list.
+// It also returns the indexes of the columns in the list.
+func (cols Columns) Indexes(subSetCols Columns) (bool, []int) {
+	var indexes []int
+	for _, subSetCol := range subSetCols {
+		colFound := false
+		for idx, col := range cols {
+			if col.Equal(subSetCol) {
+				colFound = true
+				indexes = append(indexes, idx)
+				break
+			}
+		}
+		if !colFound {
+			return false, nil
+		}
+	}
+	return true, indexes
+}
+
+// MakeColumns is used to make a list of columns from a list of strings.
+// This function is meant to be used in testing code.
+func MakeColumns(colNames ...string) Columns {
+	var cols Columns
+	for _, name := range colNames {
+		cols = append(cols, NewIdentifierCI(name))
+	}
+	return cols
+}

go/vt/sqlparser/ast_funcs_test.go

@@ -134,3 +134,41 @@ func TestSQLTypeToQueryType(t *testing.T) {
 		})
 	}
 }
+
+// TestColumns_Indexes verifies the functionality of Indexes method on Columns.
+func TestColumns_Indexes(t *testing.T) {
+	tests := []struct {
+		name          string
+		cols          Columns
+		subSetCols    Columns
+		indexesWanted []int
+	}{
+		{
+			name:       "Not a subset",
+			cols:       MakeColumns("col1", "col2", "col3"),
+			subSetCols: MakeColumns("col2", "col4"),
+		}, {
+			name:          "Subset with 1 value",
+			cols:          MakeColumns("col1", "col2", "col3"),
+			subSetCols:    MakeColumns("col2"),
+			indexesWanted: []int{1},
+		}, {
+			name:          "Subset with multiple values",
+			cols:          MakeColumns("col1", "col2", "col3", "col4", "col5"),
+			subSetCols:    MakeColumns("col3", "col5", "col1"),
+			indexesWanted: []int{2, 4, 0},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			isSubset, indexes := tt.cols.Indexes(tt.subSetCols)
+			if tt.indexesWanted == nil {
+				require.False(t, isSubset)
+				require.Nil(t, indexes)
+				return
+			}
+			require.True(t, isSubset)
+			require.EqualValues(t, tt.indexesWanted, indexes)
+		})
+	}
+}

go/vt/sqlparser/normalizer_test.go

@@ -436,6 +436,34 @@ func TestNormalizeValidSQL(t *testing.T) {
 	}
 }
 
+func TestNormalizeOneCasae(t *testing.T) {
+	testOne := struct {
+		input, output string
+	}{
+		input:  "",
+		output: "",
+	}
+	if testOne.input == "" {
+		t.Skip("empty test case")
+	}
+	tree, err := Parse(testOne.input)
+	require.NoError(t, err, testOne.input)
+	// Skip the test for the queries that do not run the normalizer
+	if !CanNormalize(tree) {
+		return
+	}
+	bv := make(map[string]*querypb.BindVariable)
+	known := make(BindVars)
+	err = Normalize(tree, NewReservedVars("vtg", known), bv)
+	require.NoError(t, err)
+	normalizerOutput := String(tree)
+	if normalizerOutput == "otheradmin" || normalizerOutput == "otherread" {
+		return
+	}
+	_, err = Parse(normalizerOutput)
+	require.NoError(t, err, normalizerOutput)
+}
+
 func TestGetBindVars(t *testing.T) {
 	stmt, err := Parse("select * from t where :v1 = :v2 and :v2 = :v3 and :v4 in ::v5")
 	if err != nil {

go/vt/vterrors/code.go

@@ -80,6 +80,7 @@ var (
 	VT10001 = errorWithoutState("VT10001", vtrpcpb.Code_ABORTED, "foreign key constraints are not allowed", "Foreign key constraints are not allowed, see https://vitess.io/blog/2021-06-15-online-ddl-why-no-fk/.")
 
 	VT12001 = errorWithoutState("VT12001", vtrpcpb.Code_UNIMPLEMENTED, "unsupported: %s", "This statement is unsupported by Vitess. Please rewrite your query to use supported syntax.")
+	VT12002 = errorWithoutState("VT12002", vtrpcpb.Code_UNIMPLEMENTED, "unsupported: cross-shard foreign keys", "Vitess does not support cross shard foreign keys.")
 
 	// VT13001 General Error
 	VT13001 = errorWithoutState("VT13001", vtrpcpb.Code_INTERNAL, "[BUG] %s", "This error should not happen and is a bug. Please file an issue on GitHub: https://github.com/vitessio/vitess/issues/new/choose.")
@@ -143,6 +144,7 @@ var (
 		VT09015,
 		VT10001,
 		VT12001,
+		VT12002,
 		VT13001,
 		VT13002,
 		VT14001,