release-22.1: sql: fix missing backrefs in mr type in DROP REGION

This commit is featured only in the release-22.1 branch and
repairs missing table back-references in the region enum type descriptor
before performing a DROP REGION or DROP SUPER REGION.

Partially addresses cockroachdb#84322.

See also cockroachdb#84144.

Release note: None

pkg/ccl/logictestccl/testdata/logic_test/multi_region_drop_region

@@ -9,11 +9,38 @@ ca-central-1    {ca-az1,ca-az2,ca-az3}  {}  {}
 us-east-1       {us-az1,us-az2,us-az3}  {}  {}
 
 statement ok
-CREATE DATABASE mr primary region "ca-central-1" regions "ap-southeast-2", "us-east-1"
+CREATE DATABASE mr PRIMARY REGION "ca-central-1" REGIONS "ap-southeast-2", "us-east-1";
+USE mr
+
+statement ok
+CREATE TABLE kv(k INT PRIMARY KEY, v INT) LOCALITY REGIONAL BY TABLE IN PRIMARY REGION
+
+statement ok
+ALTER TABLE kv SET LOCALITY REGIONAL BY ROW
+
+statement ok
+INSERT INTO kv (crdb_region, k, v) VALUES ('us-east-1', 1, 1);
+INSERT INTO kv (crdb_region, k, v) VALUES ('ca-central-1', 2, 2)
+
+query I retry
+SELECT count(*) FROM [SHOW JOBS] WHERE status = 'running' AND job_type LIKE '%SCHEMA CHANGE'
+----
+0
+
+statement error pgcode 2BP01 could not remove enum value \"us-east-1\" as it is being used by \"kv\" in row: k=1, v=1, crdb_region=\'us-east-1\'
+ALTER DATABASE mr DROP REGION "us-east-1"
+
+statement error pgcode 42P12 cannot drop region \"ca-central-1\"
+ALTER DATABASE mr DROP REGION "ca-central-1"
+
+statement ok
+DROP TABLE kv
 
 statement ok
-USE mr;
 CREATE TABLE kv (k INT PRIMARY KEY, v INT) LOCALITY REGIONAL BY ROW
 
 statement ok
 ALTER DATABASE mr DROP REGION "us-east-1"
+
+statement ok
+DROP DATABASE mr

pkg/sql/alter_database.go

@@ -351,6 +351,10 @@ func (p *planner) AlterDatabaseDropRegion(
 		return nil, err
 	}
 
+	if err := p.addMissingTableBackReferencesInAllTypeDescriptors(ctx, dbDesc); err != nil {
+		return nil, err
+	}
+
 	removingPrimaryRegion := false
 	var toDrop []*typedesc.Mutable
 
@@ -1409,6 +1413,10 @@ func (p *planner) AlterDatabaseDropSuperRegion(
 		return nil, err
 	}
 
+	if err := p.addMissingTableBackReferencesInAllTypeDescriptors(ctx, dbDesc); err != nil {
+		return nil, err
+	}
+
 	return &alterDatabaseDropSuperRegion{n: n, desc: dbDesc}, nil
 }
 
@@ -1650,3 +1658,71 @@ func (p *planner) addSuperRegion(
 		desc,
 	)
 }
+
+// addMissingTableBackReferencesInAllTypeDescriptors was introduced to fix
+// instances where, due to bugs, back-references in types would not be properly
+// updated. See github issues #84322 and #84144.
+func (p *planner) addMissingTableBackReferencesInAllTypeDescriptors(
+	ctx context.Context, db catalog.DatabaseDescriptor,
+) error {
+	// Collect all table and type descriptors in database.
+	dbTypes := make(map[descpb.ID]catalog.TypeDescriptor)
+	dbTables := make(map[descpb.ID]catalog.TableDescriptor)
+	{
+		all, err := p.Descriptors().GetAllDescriptors(ctx, p.Txn())
+		if err != nil {
+			return err
+		}
+		_ = all.ForEachDescriptorEntry(func(desc catalog.Descriptor) error {
+			if desc.GetParentID() != db.GetID() {
+				return nil
+			}
+			switch t := desc.(type) {
+			case catalog.TypeDescriptor:
+				dbTypes[t.GetID()] = t
+			case catalog.TableDescriptor:
+				dbTables[t.GetID()] = t
+			}
+			return nil
+		})
+	}
+
+	// Check that each forward reference in each table descriptor in the database
+	// has a back-reference in the type descriptor.
+	for _, tbl := range dbTables {
+		if tbl.Dropped() {
+			continue
+		}
+		getType := func(id descpb.ID) (catalog.TypeDescriptor, error) {
+			if typ, ok := dbTypes[id]; ok {
+				return typ, nil
+			}
+			return nil, errors.Wrapf(catalog.WrapTypeDescRefErr(id, catalog.ErrDescriptorNotFound),
+				"relation %q (%d)", tbl.GetName(), tbl.GetID())
+		}
+		typIDs, _, err := tbl.GetAllReferencedTypeIDs(db, getType)
+		if err != nil {
+			return err
+		}
+		for _, typID := range typIDs {
+			typ, err := getType(typID)
+			if err != nil {
+				return err
+			}
+			if catalog.MakeDescriptorIDSet(typ.TypeDesc().ReferencingDescriptorIDs...).Contains(tbl.GetID()) {
+				continue
+			}
+			// Fix missing back-reference in type descriptor to table descriptor.
+			mut, err := p.Descriptors().GetMutableTypeByID(ctx, p.Txn(), typ.GetID(), tree.ObjectLookupFlagsWithRequired())
+			if err != nil {
+				return err
+			}
+			mut.AddReferencingDescriptorID(tbl.GetID())
+			if err := p.writeTypeDesc(ctx, mut); err != nil {
+				return err
+			}
+			dbTypes[typID] = mut
+		}
+	}
+	return nil
+}