sql: fix including all related FK tables into the stmt bundle

Earlier this year in 084a741 we made
a change to include all tables referencing the table from the stmt
bundle (i.e. those that we have an inbound FK relationship with).
However, we forgot to include tables that the referencing table is
dependent on, so this commit fixes that oversight. Namely, we now
include all tables that we reference or that reference us via the FK
relationship, and this rule is applied recursively to every table in
consideration. (Previously, we didn't include inbound FKs when
handling an outbound FK nor outbound FKs when handling an inbound FK.)

Note that I believe we could have avoided including referencing tables
(i.e. with an inbound FK relationship) that only have NO ACTION or
RESTRICT actions in the ON DELETE and ON UPDATE, but I think those cases
aren't very common, and it's unlikely to hurt including all "related"
tables. The case of FK cycles is still not handled correctly on the stmt
bundle recreation, but I think we have yet to run into one in a support
ticket.

Release note: None

pkg/sql/explain_bundle_test.go

@@ -276,11 +276,18 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 	})
 
 	t.Run("foreign keys", func(t *testing.T) {
+		// All tables should be included in the stmt bundle, regardless of which
+		// one we query because all of them are considered "related" (even
+		// though we don't specify ON DELETE and ON UPDATE actions).
+		tableNames := []string{"parent", "child1", "child2", "grandchild1", "grandchild2"}
 		r.Exec(t, "CREATE TABLE parent (pk INT PRIMARY KEY, v INT);")
-		r.Exec(t, "CREATE TABLE child (pk INT PRIMARY KEY, fk INT REFERENCES parent(pk));")
+		r.Exec(t, "CREATE TABLE child1 (pk INT PRIMARY KEY, fk INT REFERENCES parent(pk));")
+		r.Exec(t, "CREATE TABLE child2 (pk INT PRIMARY KEY, fk INT REFERENCES parent(pk));")
+		r.Exec(t, "CREATE TABLE grandchild1 (pk INT PRIMARY KEY, fk INT REFERENCES child1(pk));")
+		r.Exec(t, "CREATE TABLE grandchild2 (pk INT PRIMARY KEY, fk INT REFERENCES child2(pk));")
 		contentCheck := func(name, contents string) error {
 			if name == "schema.sql" {
-				for _, tableName := range []string{"parent", "child"} {
+				for _, tableName := range tableNames {
 					if regexp.MustCompile("CREATE TABLE defaultdb.public."+tableName).FindString(contents) == "" {
 						return errors.Newf(
 							"could not find 'CREATE TABLE defaultdb.public.%s' in schema.sql:\n%s", tableName, contents)
@@ -289,12 +296,12 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 			}
 			return nil
 		}
-		for _, tableName := range []string{"parent", "child"} {
+		for _, tableName := range tableNames {
 			rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM "+tableName)
 			checkBundle(
 				t, fmt.Sprint(rows), "child", contentCheck, false, /* expectErrors */
-				base, plans, "stats-defaultdb.public.parent.sql", "stats-defaultdb.public.child.sql",
-				"distsql.html vec.txt vec-v.txt",
+				base, plans, "stats-defaultdb.public.parent.sql", "stats-defaultdb.public.child1.sql", "stats-defaultdb.public.child2.sql",
+				"stats-defaultdb.public.grandchild1.sql", "stats-defaultdb.public.grandchild2.sql", "distsql.html vec.txt vec-v.txt",
 			)
 		}
 	})

pkg/sql/opt/metadata.go

@@ -1007,49 +1007,43 @@ func (md *Metadata) getAllReferenceTables(
 	var tableSet intsets.Fast
 	var tableList []cat.DataSource
 	var addForeignKeyReferencedTables func(tab cat.Table)
+	var addForeignKeyReferencingTables func(tab cat.Table)
+	// handleRelatedTables is a helper function that processes the given table
+	// if it hasn't been handled yet by adding all referenced and referencing
+	// table of the given one, including via transient (recursive) FK
+	// relationships.
+	handleRelatedTables := func(tabID cat.StableID) {
+		if !tableSet.Contains(int(tabID)) {
+			tableSet.Add(int(tabID))
+			ds, _, err := catalog.ResolveDataSourceByID(ctx, cat.Flags{}, tabID)
+			if err != nil {
+				// This is a best-effort attempt to get all the tables, so don't
+				// error.
+				return
+			}
+			refTab, ok := ds.(cat.Table)
+			if !ok {
+				// This is a best-effort attempt to get all the tables, so don't
+				// error.
+				return
+			}
+			// We want to include all tables that we reference before adding
+			// ourselves, followed by all tables that reference us.
+			addForeignKeyReferencedTables(refTab)
+			tableList = append(tableList, ds)
+			addForeignKeyReferencingTables(refTab)
+		}
+	}
 	addForeignKeyReferencedTables = func(tab cat.Table) {
 		for i := 0; i < tab.OutboundForeignKeyCount(); i++ {
 			tabID := tab.OutboundForeignKey(i).ReferencedTableID()
-			if !tableSet.Contains(int(tabID)) {
-				tableSet.Add(int(tabID))
-				ds, _, err := catalog.ResolveDataSourceByID(ctx, cat.Flags{}, tabID)
-				if err != nil {
-					// This is a best-effort attempt to get all the tables, so don't error.
-					continue
-				}
-				refTab, ok := ds.(cat.Table)
-				if !ok {
-					// This is a best-effort attempt to get all the tables, so don't error.
-					continue
-				}
-				// We want to include all tables that we reference before adding
-				// ourselves.
-				addForeignKeyReferencedTables(refTab)
-				tableList = append(tableList, ds)
-			}
+			handleRelatedTables(tabID)
 		}
 	}
-	var addForeignKeyReferencingTables func(tab cat.Table)
 	addForeignKeyReferencingTables = func(tab cat.Table) {
 		for i := 0; i < tab.InboundForeignKeyCount(); i++ {
 			tabID := tab.InboundForeignKey(i).OriginTableID()
-			if !tableSet.Contains(int(tabID)) {
-				tableSet.Add(int(tabID))
-				ds, _, err := catalog.ResolveDataSourceByID(ctx, cat.Flags{}, tabID)
-				if err != nil {
-					// This is a best-effort attempt to get all the tables, so don't error.
-					continue
-				}
-				refTab, ok := ds.(cat.Table)
-				if !ok {
-					// This is a best-effort attempt to get all the tables, so don't error.
-					continue
-				}
-				// We want to include ourselves before all tables that reference
-				// us.
-				tableList = append(tableList, ds)
-				addForeignKeyReferencingTables(refTab)
-			}
+			handleRelatedTables(tabID)
 		}
 	}
 	for i := range md.tables {