Merge pull request #155560 from yuzefovich/blathers/backport-release-25.4-155541

release-25.4: stmtbundle: include skipped FKs into `schema.sql`

Author: yuzefovich

pkg/sql/explain_bundle.go

@@ -599,7 +599,7 @@ func (b *stmtBundleBuilder) addEnv(ctx context.Context) {
 	// update this logic to not include virtual tables into schema.sql but still
 	// create stats files for them.
 	var tables, sequences, views []tree.TableName
-	var addFKs []*tree.AlterTable
+	var addFKs, skipFKs []*tree.AlterTable
 	err := b.db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
 		// Catalog objects can show up multiple times in our lists, so
 		// deduplicate them.
@@ -778,9 +778,13 @@ func (b *stmtBundleBuilder) addEnv(ctx context.Context) {
 				include = hasDelete || hasUpdate || hasUpsert
 			},
 		)
-		addFKs = opt.GetAllFKsAmongTables(refTables, func(t cat.Table) (tree.TableName, error) {
-			return b.plan.catalog.fullyQualifiedNameWithTxn(ctx, t, txn)
-		})
+		addFKs, skipFKs = opt.GetAllFKs(
+			ctx,
+			b.plan.catalog,
+			refTables,
+			func(t cat.Table) (tree.TableName, error) {
+				return b.plan.catalog.fullyQualifiedNameWithTxn(ctx, t, txn)
+			})
 		var err error
 		tables, err = getNames(len(refTables), func(i int) cat.DataSource {
 			return refTables[i]
@@ -890,6 +894,10 @@ func (b *stmtBundleBuilder) addEnv(ctx context.Context) {
 		for _, addFK := range addFKs {
 			fmt.Fprintf(&buf, "%s;\n", addFK)
 		}
+		// Include FK constraints that were skipped in commented out form.
+		for _, skipFK := range skipFKs {
+			fmt.Fprintf(&buf, "-- %s;\n", skipFK)
+		}
 	}
 	for i := range views {
 		blankLine()

pkg/sql/explain_bundle_test.go

@@ -536,9 +536,15 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 		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));")
+		getFK := func(table string) string {
+			if table == "parent" {
+				return ""
+			}
+			return fmt.Sprintf("ALTER TABLE defaultdb.public.%s ADD CONSTRAINT", table)
+		}
 		// Only the target tables should be included since we perform a
 		// read-only stmt.
-		getContentCheckFn := func(targetTableNames, targetFKs []string) func(name, contents string) error {
+		getContentCheckFn := func(targetTableNames, addFKs, skipFKs []string) func(name, contents string) error {
 			return func(name, contents string) error {
 				if name == "schema.sql" {
 					for _, targetTableName := range targetTableNames {
@@ -563,14 +569,27 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 								"unexpectedly found non-target table 'USE defaultdb;\nCREATE TABLE public.%s' in schema.sql:\n%s", tableName, contents)
 						}
 					}
-					// Now confirm that only relevant FKs are included.
+					// Sanity-check that all FKs in the output are either in the
+					// "added" or "skipped" set.
 					numFoundFKs := strings.Count(contents, "FOREIGN KEY")
-					if numFoundFKs != len(targetFKs) {
-						return errors.Newf("found %d FKs, expected %d\n%s", numFoundFKs, len(targetFKs), contents)
+					if numFoundFKs != len(addFKs)+len(skipFKs) {
+						return errors.Newf(
+							"found %d FKs total whereas %d added and %d skipped were passed\n%s",
+							numFoundFKs, len(addFKs), len(skipFKs), contents,
+						)
+					}
+					// Now check that all expected added and skipped FKs are
+					// present.
+					for _, addFK := range addFKs {
+						if !strings.Contains(contents, addFK) {
+							return errors.Newf("didn't find added FK: %s\n%s", addFK, contents)
+						} else if strings.Contains(contents, "-- "+addFK) {
+							return errors.Newf("added FK shouldn't be commented out: %s\n%s", addFK, contents)
+						}
 					}
-					for _, fk := range targetFKs {
-						if !strings.Contains(contents, fk) {
-							return errors.Newf("didn't find target FK: %s\n%s", fk, contents)
+					for _, skipFK := range skipFKs {
+						if !strings.Contains(contents, "-- "+skipFK) {
+							return errors.Newf("didn't find skipped FK in commented out form: %s\n%s", skipFK, contents)
 						}
 					}
 				}
@@ -580,8 +599,12 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 		// First read each table separately.
 		for _, tableName := range tableNames {
 			targetTableName := tableName
-			// There should be no FKs included.
-			contentCheck := getContentCheckFn([]string{targetTableName}, nil /* targetFKs */)
+			// No FKs should be added, but 1 FK might be skipped.
+			var skipFKs []string
+			if skipFK := getFK(tableName); skipFK != "" {
+				skipFKs = []string{skipFK}
+			}
+			contentCheck := getContentCheckFn([]string{targetTableName}, nil /* addFKS */, skipFKs)
 			rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM "+targetTableName)
 			checkBundle(
 				t, fmt.Sprint(rows), targetTableName, contentCheck, false, /* expectErrors */
@@ -590,26 +613,27 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 			)
 		}
 		// Now read different combinations of tables which will influence
-		// whether ADD CONSTRAINT ... FOREIGN KEY statements are included.
-		contentCheck := getContentCheckFn([]string{"parent", "child1"}, []string{"ALTER TABLE defaultdb.public.child1 ADD CONSTRAINT"})
+		// whether ADD CONSTRAINT ... FOREIGN KEY statements are added or
+		// skipped.
+		contentCheck := getContentCheckFn([]string{"parent", "child1"}, []string{getFK("child1")}, nil)
 		rows := r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM parent, child1")
 		checkBundle(
 			t, fmt.Sprint(rows), "parent", contentCheck, false, /* expectErrors */
 			base, plans, "stats-defaultdb.public.parent.sql stats-defaultdb.public.child1.sql distsql.html vec.txt vec-v.txt",
 		)
 
-		// There should be no FKs since there isn't a direct link between the
-		// tables.
-		contentCheck = getContentCheckFn([]string{"parent", "grandchild1"}, nil /* targetFKs */)
+		// There should be no added FKs since there isn't a direct link between
+		// the tables.
+		contentCheck = getContentCheckFn([]string{"parent", "grandchild1"}, nil /* addFKS */, []string{getFK("grandchild1")})
 		rows = r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM parent, grandchild1")
 		checkBundle(
 			t, fmt.Sprint(rows), "parent", contentCheck, false, /* expectErrors */
 			base, plans, "stats-defaultdb.public.parent.sql stats-defaultdb.public.grandchild1.sql distsql.html vec.txt vec-v.txt",
 		)
 
-		// Note that we omit the FK from grandchild1 since the FK referenced
+		// Note that we skip the FK from grandchild1 since the FK referenced
 		// table isn't being read.
-		contentCheck = getContentCheckFn([]string{"parent", "child2", "grandchild1"}, []string{"ALTER TABLE defaultdb.public.child2 ADD CONSTRAINT"})
+		contentCheck = getContentCheckFn([]string{"parent", "child2", "grandchild1"}, []string{getFK("child2")}, []string{getFK("grandchild1")})
 		rows = r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM parent, child2, grandchild1")
 		checkBundle(
 			t, fmt.Sprint(rows), "parent", contentCheck, false, /* expectErrors */
@@ -618,10 +642,9 @@ CREATE TABLE users(id UUID DEFAULT gen_random_uuid() PRIMARY KEY, promo_id INT R
 
 		contentCheck = getContentCheckFn(
 			[]string{"parent", "child1", "grandchild1"},
-			[]string{
-				"ALTER TABLE defaultdb.public.child1 ADD CONSTRAINT",
-				"ALTER TABLE defaultdb.public.grandchild1 ADD CONSTRAINT",
-			})
+			[]string{getFK("child1"), getFK("grandchild1")},
+			nil, /* skipFKs */
+		)
 		rows = r.QueryStr(t, "EXPLAIN ANALYZE (DEBUG) SELECT * FROM parent, child1, grandchild1")
 		checkBundle(
 			t, fmt.Sprint(rows), "parent", contentCheck, false, /* expectErrors */

pkg/sql/opt/exec/execbuilder/relational.go

@@ -4356,12 +4356,23 @@ func (b *Builder) getEnvData() (exec.ExplainEnvData, error) {
 			refTableIncluded.Add(int(table.ID()))
 		},
 	)
-	envOpts.AddFKs = opt.GetAllFKsAmongTables(
+	var skipFKs []*tree.AlterTable
+	envOpts.AddFKs, skipFKs = opt.GetAllFKs(
+		b.ctx,
+		b.catalog,
 		refTables,
 		func(t cat.Table) (tree.TableName, error) {
 			return b.catalog.FullyQualifiedName(b.ctx, t)
 		},
 	)
+	// Given that we include all FK-related tables when visiting them, we should
+	// not skip any FKs - we'll return an error if we find any in test-only
+	// builds.
+	if buildutil.CrdbTestBuild {
+		if len(skipFKs) > 0 {
+			return envOpts, errors.AssertionFailedf("unexpectedly skipped adding FKs in EXPLAIN (OPT): %v", skipFKs)
+		}
+	}
 	var err error
 	envOpts.Tables, err = getNames(len(refTables), func(i int) cat.DataSource {
 		return refTables[i]

pkg/sql/opt/util.go

@@ -104,26 +104,53 @@ func VisitFKReferenceTables(
 	}
 }
 
-// GetAllFKsAmongTables returns a list of ALTER statements that corresponds to
-// all FOREIGN KEY constraints where both the origin and the referenced tables
-// are present in the given set of tables. List of the given tables is assumed
-// to be unique.
-func GetAllFKsAmongTables(
-	tables []cat.Table, fullyQualifiedName func(cat.Table) (tree.TableName, error),
-) []*tree.AlterTable {
+// GetAllFKs returns a list of ALTER statements that corresponds to all FOREIGN
+// KEY constraints where both the origin and the referenced tables are present
+// in the given set of tables as well as a list of ALTER statements for all
+// FOREIGN KEY constraints where the origin table is in the given set while the
+// referenced one isn't.
+//
+// List of the given tables is assumed to be unique.
+func GetAllFKs(
+	ctx context.Context,
+	catalog cat.Catalog,
+	tables []cat.Table,
+	fullyQualifiedName func(cat.Table) (tree.TableName, error),
+) (addFKs, skipFKs []*tree.AlterTable) {
 	idToTable := make(map[cat.StableID]cat.Table)
 	for _, table := range tables {
 		idToTable[table.ID()] = table
 	}
-	var addFKs []*tree.AlterTable
+	skippedIDToTable := make(map[cat.StableID]cat.Table)
+	resolveSkippedTable := func(tabID cat.StableID) cat.Table {
+		if table, ok := skippedIDToTable[tabID]; ok {
+			return table
+		}
+		ds, _, err := catalog.ResolveDataSourceByID(ctx, cat.Flags{}, tabID)
+		if err != nil {
+			return nil
+		}
+		t, ok := ds.(cat.Table)
+		if !ok {
+			return nil
+		}
+		skippedIDToTable[tabID] = t
+		return t
+	}
 	for _, origTable := range tables {
 		for i := 0; i < origTable.OutboundForeignKeyCount(); i++ {
+			includeInto := &addFKs
 			fk := origTable.OutboundForeignKey(i)
 			refTable, ok := idToTable[fk.ReferencedTableID()]
 			if !ok {
-				// The referenced table is not in the given list, so we skip
-				// this FK constraint.
-				continue
+				// The referenced table is not in the given list, so we include
+				// this FK into the skipped list.
+				includeInto = &skipFKs
+				refTable = resolveSkippedTable(fk.ReferencedTableID())
+				if refTable == nil {
+					// This is a best-effort attempt to get skipped FKs.
+					continue
+				}
 			}
 			fromCols, toCols := make(tree.NameList, fk.ColumnCount()), make(tree.NameList, fk.ColumnCount())
 			for j := range fromCols {
@@ -138,7 +165,7 @@ func GetAllFKsAmongTables(
 			if err != nil {
 				continue
 			}
-			addFKs = append(addFKs, &tree.AlterTable{
+			*includeInto = append(*includeInto, &tree.AlterTable{
 				Table: origTableName.ToUnresolvedObjectName(),
 				Cmds: []tree.AlterTableCmd{
 					&tree.AlterTableAddConstraint{
@@ -158,7 +185,7 @@ func GetAllFKsAmongTables(
 			})
 		}
 	}
-	return addFKs
+	return addFKs, skipFKs
 }
 
 // FamiliesForCols returns the set of column families for the set of cols.