feat: fix vschema updation with foreign keys and add tests

Signed-off-by: Manan Gupta <[email protected]>
vitessio · Aug 25, 2023 · 21f7aa5 · 21f7aa5
1 parent 7e589ed
commit 21f7aa5
Show file tree

Hide file tree

Showing 2 changed files with 163 additions and 2 deletions.
diff --git a/go/vt/vtgate/vschema_manager.go b/go/vt/vtgate/vschema_manager.go
@@ -198,9 +198,16 @@ func (vm *VSchemaManager) buildAndEnhanceVSchema(v *vschemapb.SrvVSchema) *vinde
 func (vm *VSchemaManager) updateFromSchema(vschema *vindexes.VSchema) {
 	for ksName, ks := range vschema.Keyspaces {
 		m := vm.schema.Tables(ksName)
+		// Before we add the foreign key definitions in the tables, we need to make sure that all the tables
+		// are created in the Vschema, so that later when we try to find the routed tables, we don't end up
+		// getting dummy tables.
+		for tblName, tblInfo := range m {
+			setColumns(ks, tblName, tblInfo.Columns)
+		}
 
+		// Now that we have ensured that all the tables are created, we can start populating the foreign keys
+		// in the tables.
 		for tblName, tblInfo := range m {
-			vTbl := setColumns(ks, tblName, tblInfo.Columns)
 			for _, fkDef := range tblInfo.ForeignKeys {
 				parentTbl, err := vschema.FindRoutedTable(ksName, fkDef.ReferenceDefinition.ReferencedTable.Name.String(), topodatapb.TabletType_PRIMARY)
 				if err != nil {
@@ -212,7 +219,7 @@ func (vm *VSchemaManager) updateFromSchema(vschema *vindexes.VSchema) {
 					log.Errorf("error finding child table %s: %v", tblName, err)
 					continue
 				}
-				vTbl.ParentForeignKeys = append(vTbl.ParentForeignKeys, vindexes.NewParentFkInfo(parentTbl, fkDef))
+				childTbl.ParentForeignKeys = append(childTbl.ParentForeignKeys, vindexes.NewParentFkInfo(parentTbl, fkDef))
 				parentTbl.ChildForeignKeys = append(parentTbl.ChildForeignKeys, vindexes.NewChildFkInfo(childTbl, fkDef))
 			}
 		}

diff --git a/go/vt/vtgate/vschema_manager_test.go b/go/vt/vtgate/vschema_manager_test.go
@@ -29,6 +29,58 @@ func TestVSchemaUpdate(t *testing.T) {
 	tblCol2 := &vindexes.Table{Name: sqlparser.NewIdentifierCS("tbl"), Keyspace: ks, Columns: cols2, ColumnListAuthoritative: true}
 	tblCol2NA := &vindexes.Table{Name: sqlparser.NewIdentifierCS("tbl"), Keyspace: ks, Columns: cols2}
 
+	vindexTable_multicol_t1 := &vindexes.Table{
+		Name:                    sqlparser.NewIdentifierCS("multicol_t1"),
+		Keyspace:                ks,
+		Columns:                 cols2,
+		ColumnListAuthoritative: true,
+	}
+	vindexTable_multicol_t2 := &vindexes.Table{
+		Name:                    sqlparser.NewIdentifierCS("multicol_t2"),
+		Keyspace:                ks,
+		Columns:                 cols2,
+		ColumnListAuthoritative: true,
+	}
+	vindexTable_t1 := &vindexes.Table{
+		Name:                    sqlparser.NewIdentifierCS("t1"),
+		Keyspace:                ks,
+		Columns:                 cols1,
+		ColumnListAuthoritative: true,
+	}
+	vindexTable_t2 := &vindexes.Table{
+		Name:                    sqlparser.NewIdentifierCS("t2"),
+		Keyspace:                ks,
+		Columns:                 cols1,
+		ColumnListAuthoritative: true,
+	}
+	sqlparserCols1 := sqlparser.MakeColumns("id")
+	sqlparserCols2 := sqlparser.MakeColumns("uid", "name")
+
+	vindexTable_multicol_t1.ChildForeignKeys = append(vindexTable_multicol_t1.ChildForeignKeys, vindexes.ChildFKInfo{
+		Table:         vindexTable_multicol_t2,
+		ChildColumns:  sqlparserCols2,
+		ParentColumns: sqlparserCols2,
+		OnDelete:      sqlparser.NoAction,
+		OnUpdate:      sqlparser.Restrict,
+	})
+	vindexTable_multicol_t2.ParentForeignKeys = append(vindexTable_multicol_t2.ParentForeignKeys, vindexes.ParentFKInfo{
+		Table:         vindexTable_multicol_t1,
+		ChildColumns:  sqlparserCols2,
+		ParentColumns: sqlparserCols2,
+	})
+	vindexTable_t1.ChildForeignKeys = append(vindexTable_t1.ChildForeignKeys, vindexes.ChildFKInfo{
+		Table:         vindexTable_t2,
+		ChildColumns:  sqlparserCols1,
+		ParentColumns: sqlparserCols1,
+		OnDelete:      sqlparser.SetNull,
+		OnUpdate:      sqlparser.Cascade,
+	})
+	vindexTable_t2.ParentForeignKeys = append(vindexTable_t2.ParentForeignKeys, vindexes.ParentFKInfo{
+		Table:         vindexTable_t1,
+		ChildColumns:  sqlparserCols1,
+		ParentColumns: sqlparserCols1,
+	})
+
 	tcases := []struct {
 		name           string
 		srvVschema     *vschemapb.SrvVSchema
@@ -94,6 +146,108 @@ func TestVSchemaUpdate(t *testing.T) {
 		currentVSchema: &vindexes.VSchema{},
 		schema:         map[string]*vindexes.TableInfo{"tbl": {Columns: cols1}},
 		expected:       &vindexes.VSchema{},
+	}, {
+		name:           "foreign keys in schema",
+		currentVSchema: &vindexes.VSchema{},
+		schema: map[string]*vindexes.TableInfo{
+			"t1": {
+				Columns: cols1,
+			},
+			"t2": {
+				Columns: cols1,
+				ForeignKeys: []*sqlparser.ForeignKeyDefinition{
+					{
+						Source: sqlparser.MakeColumns("id"),
+						ReferenceDefinition: &sqlparser.ReferenceDefinition{
+							ReferencedTable:   sqlparser.NewTableName("t1"),
+							ReferencedColumns: sqlparserCols1,
+							OnUpdate:          sqlparser.Cascade,
+							OnDelete:          sqlparser.SetNull,
+						},
+					},
+				},
+			},
+			"multicol_t1": {
+				Columns: cols2,
+			},
+			"multicol_t2": {
+				Columns: cols2,
+				ForeignKeys: []*sqlparser.ForeignKeyDefinition{
+					{
+						Source: sqlparser.MakeColumns("uid", "name"),
+						ReferenceDefinition: &sqlparser.ReferenceDefinition{
+							ReferencedTable:   sqlparser.NewTableName("multicol_t1"),
+							ReferencedColumns: sqlparserCols2,
+							OnUpdate:          sqlparser.Restrict,
+							OnDelete:          sqlparser.NoAction,
+						},
+					},
+				},
+			},
+		},
+		srvVschema: &vschemapb.SrvVSchema{
+			Keyspaces: map[string]*vschemapb.Keyspace{
+				"ks": {
+					Sharded:        false,
+					ForeignKeyMode: vschemapb.Keyspace_FK_MANAGED,
+					Tables: map[string]*vschemapb.Table{
+						"t1": {
+							Columns: []*vschemapb.Column{
+								{
+									Name: "id",
+									Type: querypb.Type_INT64,
+								},
+							},
+						},
+						"t2": {
+							Columns: []*vschemapb.Column{
+								{
+									Name: "id",
+									Type: querypb.Type_INT64,
+								},
+							},
+						},
+						"multicol_t1": {
+							Columns: []*vschemapb.Column{
+								{
+									Name: "uid",
+									Type: querypb.Type_INT64,
+								}, {
+									Name: "name",
+									Type: querypb.Type_VARCHAR,
+								},
+							},
+						}, "multicol_t2": {
+							Columns: []*vschemapb.Column{
+								{
+									Name: "uid",
+									Type: querypb.Type_INT64,
+								}, {
+									Name: "name",
+									Type: querypb.Type_VARCHAR,
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		expected: &vindexes.VSchema{
+			RoutingRules: map[string]*vindexes.RoutingRule{},
+			Keyspaces: map[string]*vindexes.KeyspaceSchema{
+				"ks": {
+					Keyspace:       ks,
+					ForeignKeyMode: vschemapb.Keyspace_FK_MANAGED,
+					Vindexes:       map[string]vindexes.Vindex{},
+					Tables: map[string]*vindexes.Table{
+						"t1":          vindexTable_t1,
+						"t2":          vindexTable_t2,
+						"multicol_t1": vindexTable_multicol_t1,
+						"multicol_t2": vindexTable_multicol_t2,
+					},
+				},
+			},
+		},
 	}}
 
 	vm := &VSchemaManager{}