Merge #155541

155541: stmtbundle: include skipped FKs into `schema.sql` r=yuzefovich a=yuzefovich

Earlier this year we refactored how we handle FKs in stmt bundles. Namely, we now use IGNORE_FOREIGN_KEYS option of SHOW CREATE TABLE (which produces DDLs with FK constraints omitted) and then we include only "relevant" FKs as ALTER TABLE ... ADD CONSTRAINT stmts later. This behavior can be confusing since now CREATE TABLE stmt from the bundle cannot be relied upon as the source of truth for the table's schema. In order to partially remedy the situation, this commit extends the logic to also include ALTER TABLE stmts for "skipped" FKs where the origin (i.e. referencing) table is included into the bundle, but these ALTERs are commented out (to keep the bundle being recreatable).

Informs: https://github.com/cockroachlabs/support/issues/3459
Epic: None
Release note: None

Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: craig[bot]

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

@@ -4353,12 +4353,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.