schemadiff: RelatedForeignKeyTables()

go/vt/schemadiff/schema.go

@@ -35,8 +35,8 @@ type Schema struct {
 	named  map[string]Entity
 	sorted []Entity
 
-	foreignKeyParents  []*CreateTableEntity // subset of tables
-	foreignKeyChildren []*CreateTableEntity // subset of tables
+	fkChildToParents   map[string][]*CreateTableEntity
+	fkParentToChildren map[string][]*CreateTableEntity
 
 	env *Environment
 }
@@ -49,8 +49,8 @@ func newEmptySchema(env *Environment) *Schema {
 		named:  map[string]Entity{},
 		sorted: []Entity{},
 
-		foreignKeyParents:  []*CreateTableEntity{},
-		foreignKeyChildren: []*CreateTableEntity{},
+		fkChildToParents:   map[string][]*CreateTableEntity{},
+		fkParentToChildren: map[string][]*CreateTableEntity{},
 
 		env: env,
 	}
@@ -236,8 +236,13 @@ func (s *Schema) normalize(hints *DiffHints) error {
 			}
 			// Not handled. Does this table reference an already handled table?
 			referencedTableNames := getForeignKeyParentTableNames(t.CreateTable)
-			if len(referencedTableNames) > 0 {
-				s.foreignKeyChildren = append(s.foreignKeyChildren, t)
+			for _, referencedTableName := range referencedTableNames {
+				referencedTableEntity := s.Table(referencedTableName)
+				if referencedTableEntity == nil {
+					continue
+				}
+				s.fkParentToChildren[referencedTableName] = append(s.fkParentToChildren[referencedTableName], t)
+				s.fkChildToParents[name] = append(s.fkChildToParents[name], referencedTableEntity)
 			}
 			nonSelfReferenceNames := []string{}
 			for _, referencedTableName := range referencedTableNames {
@@ -269,11 +274,6 @@ func (s *Schema) normalize(hints *DiffHints) error {
 		}
 		iterationLevel++
 	}
-	for _, t := range s.tables {
-		if fkParents[t.Name()] {
-			s.foreignKeyParents = append(s.foreignKeyParents, t)
-		}
-	}
 	if len(dependencyLevels) != len(s.tables) {
 		// We have leftover tables. This can happen if there's foreign key loops
 		for _, t := range s.tables {
@@ -890,7 +890,7 @@ func (s *Schema) SchemaDiff(other *Schema, hints *DiffHints) (*SchemaDiff, error
 	if err != nil {
 		return nil, err
 	}
-	schemaDiff := NewSchemaDiff(s, hints)
+	schemaDiff := NewSchemaDiff(s, other, hints)
 	schemaDiff.loadDiffs(diffs)
 
 	// Utility function to see whether the given diff has dependencies on diffs that operate on any of the given named entities,

go/vt/schemadiff/schema_diff.go

@@ -193,6 +193,7 @@ func permDiff(
 // Operations on SchemaDiff are not concurrency-safe.
 type SchemaDiff struct {
 	schema *Schema
+	from   *Schema
 	hints  *DiffHints
 	diffs  []EntityDiff
 
@@ -202,9 +203,10 @@ type SchemaDiff struct {
 	r *mathutil.EquivalenceRelation // internal structure to help determine diffs
 }
 
-func NewSchemaDiff(schema *Schema, hints *DiffHints) *SchemaDiff {
+func NewSchemaDiff(schema *Schema, to *Schema, hints *DiffHints) *SchemaDiff {
 	return &SchemaDiff{
 		schema:       schema,
+		from:         to,
 		hints:        hints,
 		dependencies: make(map[string]*DiffDependency),
 		diffMap:      make(map[string]EntityDiff),
@@ -451,3 +453,30 @@ func (d *SchemaDiff) InstantDDLCapability() InstantDDLCapability {
 	}
 	return capability
 }
+
+// RelatedForeignKeyTables returns the foreign key tables, either parent or child, that are affected by the
+// diffs in this schema diff, either explicitly (table is modified) or implicitly (table's parent or
+// child is modified).
+func (d *SchemaDiff) RelatedForeignKeyTables() (fkTables map[string]*CreateTableEntity) {
+	fkTables = make(map[string]*CreateTableEntity)
+
+	scanForAffectedFKs := func(entityName string) {
+		for _, schema := range []*Schema{d.from, d.schema} {
+			for _, child := range schema.fkParentToChildren[entityName] {
+				fkTables[entityName] = schema.Table(entityName) // this is a parent
+				fkTables[child.Name()] = child
+			}
+			for _, parent := range schema.fkChildToParents[entityName] {
+				fkTables[entityName] = schema.Table(entityName) // this is a child
+				fkTables[parent.Name()] = parent
+			}
+		}
+	}
+	for _, diff := range d.UnorderedDiffs() {
+		switch diff := diff.(type) {
+		case *CreateTableEntityDiff, *AlterTableEntityDiff, *DropTableEntityDiff:
+			scanForAffectedFKs(diff.EntityName())
+		}
+	}
+	return fkTables
+}