scbuild,scdeps: replace deprecated descs.Collection method calls

Informs cockroachdb#64089.

Release note: None

pkg/sql/schemachanger/scbuild/BUILD.bazel

@@ -23,8 +23,8 @@ go_library(
         "//pkg/sql/catalog/colinfo",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/catalog/descs",
+        "//pkg/sql/catalog/nstree",
         "//pkg/sql/catalog/resolver",
-        "//pkg/sql/catalog/schemadesc",
         "//pkg/sql/catalog/schemaexpr",
         "//pkg/sql/catalog/seqexpr",
         "//pkg/sql/catalog/tabledesc",

pkg/sql/schemachanger/scbuild/builder_state.go

@@ -12,15 +12,13 @@ package scbuild
 
 import (
 	"sort"
-	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemadesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemaexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/seqexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
@@ -988,13 +986,21 @@ func (b *builderState) ensureDescriptor(id catid.DescID) {
 		// Handle special case of database children, which may include temporary
 		// schemas, which aren't explicitly referenced in the database's schemas
 		// map.
-		childSchemas := b.cr.MustGetSchemasForDatabase(b.ctx, d)
-		for schemaID, schemaName := range childSchemas {
-			c.backrefs.Add(schemaID)
-			if strings.HasPrefix(schemaName, catconstants.PgTempSchemaName) {
-				b.tempSchemas[schemaID] = schemadesc.NewTemporarySchema(schemaName, schemaID, d.GetID())
+		childSchemas := b.cr.GetAllSchemasInDatabase(b.ctx, d)
+		_ = childSchemas.ForEachDescriptor(func(desc catalog.Descriptor) error {
+			sc, err := catalog.AsSchemaDescriptor(desc)
+			if err != nil {
+				panic(err)
 			}
-		}
+			switch sc.SchemaKind() {
+			case catalog.SchemaVirtual:
+				return nil
+			case catalog.SchemaTemporary:
+				b.tempSchemas[sc.GetID()] = sc
+			}
+			c.backrefs.Add(sc.GetID())
+			return nil
+		})
 	case catalog.SchemaDescriptor:
 		b.ensureDescriptor(c.desc.GetParentID())
 		db := b.descCache[c.desc.GetParentID()].desc
@@ -1003,13 +1009,11 @@ func (b *builderState) ensureDescriptor(id catid.DescID) {
 		// Handle special case of schema children, which have to be added to
 		// the back-referenced ID set but which aren't explicitly referenced in
 		// the schema descriptor itself.
-		b.cr.ReadObjectIDs(b.ctx, db.(catalog.DatabaseDescriptor), d).ForEach(c.backrefs.Add)
-		if err := d.ForEachFunctionOverload(func(overload descpb.SchemaDescriptor_FunctionOverload) error {
-			c.backrefs.Add(overload.ID)
+		objects := b.cr.GetAllObjectsInSchema(b.ctx, db.(catalog.DatabaseDescriptor), d)
+		_ = objects.ForEachDescriptor(func(desc catalog.Descriptor) error {
+			c.backrefs.Add(desc.GetID())
 			return nil
-		}); err != nil {
-			panic(err)
-		}
+		})
 	default:
 		b.ensureDescriptor(c.desc.GetParentID())
 		db := b.descCache[c.desc.GetParentID()].desc

pkg/sql/schemachanger/scbuild/dependencies.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/nstree"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/resolver"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scbuild/internal/scbuildstmt"
@@ -133,17 +134,14 @@ type CatalogReader interface {
 		found bool, prefix catalog.ResolvedObjectPrefix, tbl catalog.TableDescriptor, idx catalog.Index,
 	)
 
-	// ReadObjectIDs looks up the IDs of all objects in a schema.
-	ReadObjectIDs(ctx context.Context, db catalog.DatabaseDescriptor, schema catalog.SchemaDescriptor) catalog.DescriptorIDSet
+	// GetAllSchemasInDatabase gets all schemas in a database.
+	GetAllSchemasInDatabase(ctx context.Context, database catalog.DatabaseDescriptor) nstree.Catalog
+
+	// GetAllObjectsInSchema gets all non-dropped objects in a schema.
+	GetAllObjectsInSchema(ctx context.Context, db catalog.DatabaseDescriptor, schema catalog.SchemaDescriptor) nstree.Catalog
 
 	// MustReadDescriptor looks up a descriptor by ID.
 	MustReadDescriptor(ctx context.Context, id descpb.ID) catalog.Descriptor
-
-	// MustGetSchemasForDatabase gets schemas associated with
-	// a database.
-	MustGetSchemasForDatabase(
-		ctx context.Context, database catalog.DatabaseDescriptor,
-	) map[descpb.ID]string
 }
 
 // TableReader implements functions for inspecting tables during the build phase,

pkg/sql/schemachanger/scdeps/build_deps.go

@@ -178,19 +178,15 @@ func (d *buildDeps) MayResolveIndex(
 	return found, prefix, tbl, idx
 }
 
-// ReadObjectIDs implements the scbuild.CatalogReader interface.
-func (d *buildDeps) ReadObjectIDs(
+// GetAllObjectsInSchema implements the scbuild.CatalogReader interface.
+func (d *buildDeps) GetAllObjectsInSchema(
 	ctx context.Context, db catalog.DatabaseDescriptor, schema catalog.SchemaDescriptor,
-) (ret catalog.DescriptorIDSet) {
-	c, err := d.descsCollection.GetObjectNamesAndIDs(ctx, d.txn, db, schema)
+) nstree.Catalog {
+	c, err := d.descsCollection.GetAllObjectsInSchema(ctx, d.txn, db, schema)
 	if err != nil {
 		panic(err)
 	}
-	_ = c.ForEachNamespaceEntry(func(e nstree.NamespaceEntry) error {
-		ret.Add(e.GetID())
-		return nil
-	})
-	return ret
+	return c
 }
 
 // ResolveType implements the scbuild.CatalogReader interface.
@@ -247,11 +243,11 @@ func (d *buildDeps) MustReadDescriptor(ctx context.Context, id descpb.ID) catalo
 	return desc
 }
 
-// MustGetSchemasForDatabase  implements the scbuild.CatalogReader interface.
-func (d *buildDeps) MustGetSchemasForDatabase(
+// GetAllSchemasInDatabase implements the scbuild.CatalogReader interface.
+func (d *buildDeps) GetAllSchemasInDatabase(
 	ctx context.Context, database catalog.DatabaseDescriptor,
-) map[descpb.ID]string {
-	schemas, err := d.descsCollection.GetSchemasForDatabase(ctx, d.txn, database)
+) nstree.Catalog {
+	schemas, err := d.descsCollection.GetAllSchemasInDatabase(ctx, d.txn, database)
 	if err != nil {
 		panic(err)
 	}

pkg/sql/schemachanger/scdeps/sctestdeps/BUILD.bazel

@@ -51,6 +51,7 @@ go_library(
         "//pkg/sql/types",
         "//pkg/testutils/sqlutils",
         "//pkg/util/hlc",
+        "//pkg/util/iterutil",
         "//pkg/util/log/eventpb",
         "//pkg/util/log/logpb",
         "//pkg/util/protoutil",

pkg/sql/schemachanger/scdeps/sctestdeps/test_deps.go

@@ -43,6 +43,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/iterutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logpb"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
@@ -295,42 +296,46 @@ func (s *TestState) MayResolveType(
 // MayResolveIndex implements the scbuild.CatalogReader interface.
 func (s *TestState) MayResolveIndex(
 	ctx context.Context, tableIndexName tree.TableIndexName,
-) (found bool, _ catalog.ResolvedObjectPrefix, _ catalog.TableDescriptor, _ catalog.Index) {
+) (
+	found bool,
+	prefix catalog.ResolvedObjectPrefix,
+	tbl catalog.TableDescriptor,
+	idx catalog.Index,
+) {
 	if tableIndexName.Table.Object() != "" {
-		prefix, tbl := s.MayResolveTable(ctx, *tableIndexName.Table.ToUnresolvedObjectName())
+		prefix, tbl = s.MayResolveTable(ctx, *tableIndexName.Table.ToUnresolvedObjectName())
 		if tbl == nil {
 			return false, catalog.ResolvedObjectPrefix{}, nil, nil
 		}
-		idx, err := tbl.FindNonDropIndexWithName(string(tableIndexName.Index))
-		if err != nil {
-			return false, catalog.ResolvedObjectPrefix{}, nil, nil
-		}
-		return true, prefix, tbl, idx
-	}
-
-	db, schema := s.mayResolvePrefix(tableIndexName.Table.ObjectNamePrefix)
-	prefix := catalog.ResolvedObjectPrefix{
-		ExplicitDatabase: true,
-		ExplicitSchema:   true,
-		Database:         db.(catalog.DatabaseDescriptor),
-		Schema:           schema.(catalog.SchemaDescriptor),
-	}
-	objectIDs := s.ReadObjectIDs(ctx, prefix.Database, prefix.Schema)
-	for _, objectID := range objectIDs.Ordered() {
-		desc, _ := s.mustReadImmutableDescriptor(objectID)
-		if desc == nil {
-			continue
-		}
-		tbl, ok := desc.(catalog.TableDescriptor)
-		if !ok {
-			continue
+		idx, _ = tbl.FindNonDropIndexWithName(string(tableIndexName.Index))
+	} else {
+		db, schema := s.mayResolvePrefix(tableIndexName.Table.ObjectNamePrefix)
+		prefix = catalog.ResolvedObjectPrefix{
+			ExplicitDatabase: true,
+			ExplicitSchema:   true,
+			Database:         db.(catalog.DatabaseDescriptor),
+			Schema:           schema.(catalog.SchemaDescriptor),
 		}
-		idx, _ := tbl.FindNonDropIndexWithName(string(tableIndexName.Index))
-		if idx != nil {
-			return true, prefix, tbl, idx
+		var objects nstree.Catalog
+		if db != nil && schema != nil {
+			objects = s.GetAllObjectsInSchema(ctx, prefix.Database, prefix.Schema)
 		}
+		_ = objects.ForEachDescriptor(func(desc catalog.Descriptor) error {
+			var ok bool
+			tbl, ok = desc.(catalog.TableDescriptor)
+			if !ok {
+				return nil
+			}
+			idx, _ = tbl.FindNonDropIndexWithName(string(tableIndexName.Index))
+			if idx != nil {
+				return iterutil.StopIteration()
+			}
+			return nil
+		})
+	}
+	if idx != nil {
+		return true, prefix, tbl, idx
 	}
-
 	return false, catalog.ResolvedObjectPrefix{}, nil, nil
 }
 
@@ -436,17 +441,18 @@ func (s *TestState) mayGetByName(
 	return desc
 }
 
-// ReadObjectIDs implements the scbuild.CatalogReader interface.
-func (s *TestState) ReadObjectIDs(
-	_ context.Context, db catalog.DatabaseDescriptor, schema catalog.SchemaDescriptor,
-) (ret catalog.DescriptorIDSet) {
-	_ = s.uncommitted.ForEachNamespaceEntry(func(e nstree.NamespaceEntry) error {
-		if e.GetParentID() == db.GetID() && e.GetParentSchemaID() == schema.GetID() {
-			ret.Add(e.GetID())
+// GetAllObjectsInSchema implements the scbuild.CatalogReader interface.
+func (s *TestState) GetAllObjectsInSchema(
+	ctx context.Context, db catalog.DatabaseDescriptor, schema catalog.SchemaDescriptor,
+) nstree.Catalog {
+	var ret nstree.MutableCatalog
+	_ = s.uncommitted.ForEachDescriptor(func(desc catalog.Descriptor) error {
+		if desc.GetParentSchemaID() == schema.GetID() {
+			ret.UpsertDescriptor(desc)
 		}
 		return nil
 	})
-	return ret
+	return ret.Catalog
 }
 
 // ResolveType implements the scbuild.CatalogReader interface.
@@ -530,19 +536,18 @@ func (s *TestState) CurrentDatabase() string {
 	return s.currentDatabase
 }
 
-// MustGetSchemasForDatabase implements the scbuild.CatalogReader interface.
-func (s *TestState) MustGetSchemasForDatabase(
-	ctx context.Context, database catalog.DatabaseDescriptor,
-) map[descpb.ID]string {
-	schemas := make(map[descpb.ID]string)
-	err := database.ForEachSchema(func(id descpb.ID, name string) error {
-		schemas[id] = name
+// GetAllSchemasInDatabase implements the scbuild.CatalogReader interface.
+func (s *TestState) GetAllSchemasInDatabase(
+	_ context.Context, database catalog.DatabaseDescriptor,
+) nstree.Catalog {
+	var ret nstree.MutableCatalog
+	_ = s.uncommitted.ForEachDescriptor(func(desc catalog.Descriptor) error {
+		if desc.GetParentID() == database.GetID() && desc.GetParentSchemaID() == descpb.InvalidID {
+			ret.UpsertDescriptor(desc)
+		}
 		return nil
 	})
-	if err != nil {
-		panic(err)
-	}
-	return schemas
+	return ret.Catalog
 }
 
 // MustReadDescriptor implements the scbuild.CatalogReader interface.