descs: push descriptor type hydration into the desc.Collection

Fixes cockroachdb#49484.
Preparation for cockroachdb#51385.

Up until now, the planner was responsible for installing user defined
type metadata in tables that contained user defined types. This was
slightly messy and caused leakage of responsibility regarding when
descriptors had user defined types vs when they didn't. This commit
pushes that responsibility into the descs.Collection. It also paves
the way for work to avoid copying ImmutableTableDescriptors that
contain user defined types every time that they need hydration.

Release note: None

pkg/sql/backfill.go

@@ -28,7 +28,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/resolver"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
@@ -655,21 +654,6 @@ func (sc *SchemaChanger) truncateIndexes(
 				if err != nil {
 					return err
 				}
-
-				// Hydrate types used in the retrieved table.
-				// TODO (rohany): This can be removed once table access from the
-				//  desc.Collection returns tables with hydrated types.
-				typLookup := func(ctx context.Context, id sqlbase.ID) (*tree.TypeName, sqlbase.TypeDescriptorInterface, error) {
-					return resolver.ResolveTypeDescByID(ctx, txn, sc.execCfg.Codec, id, tree.ObjectLookupFlags{})
-				}
-				if err := sqlbase.HydrateTypesInTableDescriptor(
-					ctx,
-					tableDesc.TableDesc(),
-					sqlbase.TypeLookupFunc(typLookup),
-				); err != nil {
-					return err
-				}
-
 				rd, err := row.MakeDeleter(
 					ctx,
 					txn,

pkg/sql/catalog/descs/collection.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
 )
 
@@ -162,7 +163,11 @@ func (tc *Collection) GetMutableTableDescriptor(
 	if !ok {
 		return nil, nil
 	}
-	return mutDesc, nil
+	hydrated, err := tc.hydrateTypesInTableDesc(ctx, txn, mutDesc)
+	if err != nil {
+		return nil, err
+	}
+	return hydrated.(*sqlbase.MutableTableDescriptor), nil
 }
 
 func (tc *Collection) getMutableObjectDescriptor(
@@ -288,7 +293,11 @@ func (tc *Collection) GetTableVersion(
 		}
 		return nil, nil
 	}
-	return table, nil
+	hydrated, err := tc.hydrateTypesInTableDesc(ctx, txn, table)
+	if err != nil {
+		return nil, err
+	}
+	return hydrated.(*sqlbase.ImmutableTableDescriptor), nil
 }
 
 func (tc *Collection) getObjectVersion(
@@ -434,7 +443,11 @@ func (tc *Collection) GetTableVersionByID(
 		return nil, sqlbase.NewUndefinedRelationError(
 			&tree.TableRef{TableID: int64(tableID)})
 	}
-	return table, nil
+	hydrated, err := tc.hydrateTypesInTableDesc(ctx, txn, table)
+	if err != nil {
+		return nil, err
+	}
+	return hydrated.(*sqlbase.ImmutableTableDescriptor), nil
 }
 
 func (tc *Collection) getDescriptorVersionByID(
@@ -506,7 +519,12 @@ func (tc *Collection) GetMutableTableVersionByID(
 	if err != nil {
 		return nil, err
 	}
-	return desc.(*sqlbase.MutableTableDescriptor), nil
+	table := desc.(*sqlbase.MutableTableDescriptor)
+	hydrated, err := tc.hydrateTypesInTableDesc(ctx, txn, table)
+	if err != nil {
+		return nil, err
+	}
+	return hydrated.(*sqlbase.MutableTableDescriptor), nil
 }
 
 func (tc *Collection) getMutableDescriptorByID(
@@ -528,6 +546,46 @@ func (tc *Collection) getMutableDescriptorByID(
 	return desc, nil
 }
 
+// hydrateTypesInTableDesc installs user defined type metadata in all types.T
+// present in the input TableDescriptor. It always returns the same type of
+// TableDescriptor that was passed in. It ensures that ImmutableTableDescriptors
+// are not modified during the process of metadata installation.
+func (tc *Collection) hydrateTypesInTableDesc(
+	ctx context.Context, txn *kv.Txn, desc sqlbase.TableDescriptorInterface,
+) (sqlbase.TableDescriptorInterface, error) {
+	getType := func(ctx context.Context, id sqlbase.ID) (*tree.TypeName, sqlbase.TypeDescriptorInterface, error) {
+		// TODO (rohany): Use the collection API's here.
+		return resolver.ResolveTypeDescByID(ctx, txn, tc.codec(), id, tree.ObjectLookupFlags{})
+	}
+	switch t := desc.(type) {
+	case *sqlbase.MutableTableDescriptor:
+		// It is safe to hydrate directly into MutableTableDescriptor since it is
+		// not shared.
+		return desc, sqlbase.HydrateTypesInTableDescriptor(ctx, t.TableDesc(), sqlbase.TypeLookupFunc(getType))
+	case *sqlbase.ImmutableTableDescriptor:
+		// ImmutableTableDescriptors need to be copied before hydration, because
+		// they are potentially read by multiple threads. If there aren't any user
+		// defined types in the descriptor, then return early.
+		if !t.ContainsUserDefinedTypes() {
+			return desc, nil
+		}
+
+		// TODO (rohany, ajwerner): Here we would look into the cached set of
+		//  hydrated table descriptors and potentially return without having to
+		//  make a copy. However, we could avoid hitting the cache if any of the
+		//  user defined types have been modified in this transaction.
+
+		// Make a copy of the underlying descriptor before hydration.
+		descBase := protoutil.Clone(t.TableDesc()).(*sqlbase.TableDescriptor)
+		if err := sqlbase.HydrateTypesInTableDescriptor(ctx, descBase, sqlbase.TypeLookupFunc(getType)); err != nil {
+			return nil, err
+		}
+		return sqlbase.NewImmutableTableDescriptor(*descBase), nil
+	default:
+		return desc, nil
+	}
+}
+
 // ReleaseSpecifiedLeases releases the leases for the descriptors with ids in
 // the passed slice. Errors are logged but ignored.
 func (tc *Collection) ReleaseSpecifiedLeases(ctx context.Context, descs []lease.IDVersion) {

pkg/sql/planner.go

@@ -490,13 +490,7 @@ func (p *planner) LookupTableByID(
 		}
 		return catalog.TableEntry{}, err
 	}
-	// TODO (rohany): This shouldn't be needed once the descs.Collection always
-	//  returns descriptors with hydrated types.
-	hydratedDesc, err := p.maybeHydrateTypesInDescriptor(ctx, table)
-	if err != nil {
-		return catalog.TableEntry{}, err
-	}
-	return catalog.TableEntry{Desc: hydratedDesc.(*sqlbase.ImmutableTableDescriptor)}, nil
+	return catalog.TableEntry{Desc: table}, nil
 }
 
 // TypeAsString enforces (not hints) that the given expression typechecks as a

pkg/sql/resolver.go

@@ -16,7 +16,6 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/resolver"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
@@ -25,7 +24,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
-	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
 )
 
@@ -129,14 +127,6 @@ func (p *planner) LookupObject(
 	lookupFlags.CommonLookupFlags = p.CommonLookupFlags(false /* required */)
 	objDesc, err := sc.GetObjectDesc(ctx, p.txn, p.ExecCfg().Settings, p.ExecCfg().Codec, dbName, scName, tbName, lookupFlags)
 
-	// The returned object may contain types.T that need hydrating.
-	if objDesc != nil {
-		objDesc, err = p.maybeHydrateTypesInDescriptor(ctx, objDesc)
-		if err != nil {
-			return false, nil, err
-		}
-	}
-
 	return objDesc != nil, objDesc, err
 }
 
@@ -205,38 +195,6 @@ func (p *planner) ResolveTypeByID(ctx context.Context, id uint32) (*types.T, err
 	return desc.MakeTypesT(ctx, name, p)
 }
 
-// maybeHydrateTypesInDescriptor hydrates any types.T's in the input descriptor.
-// TODO (rohany): Once we lease types, this should be pushed down into the
-//  leased object collection.
-func (p *planner) maybeHydrateTypesInDescriptor(
-	ctx context.Context, objDesc catalog.Descriptor,
-) (catalog.Descriptor, error) {
-	// As of now, only {Mutable,Immutable}TableDescriptor have types.T that
-	// need to be hydrated.
-	switch t := objDesc.(type) {
-	case *sqlbase.MutableTableDescriptor:
-		// MutableTableDescriptors are safe to modify in place.
-		if err := sqlbase.HydrateTypesInTableDescriptor(ctx, t.TableDesc(), p); err != nil {
-			return nil, err
-		}
-		return objDesc, nil
-	case *sqlbase.ImmutableTableDescriptor:
-		// ImmutableTableDescriptors need to be copied before hydration. If there
-		// aren't any user defined types in the descriptor, then just return.
-		if !t.ContainsUserDefinedTypes() {
-			return objDesc, nil
-		}
-		// Make a copy of the underlying TableDescriptor.
-		desc := protoutil.Clone(t.TableDesc()).(*sqlbase.TableDescriptor)
-		if err := sqlbase.HydrateTypesInTableDescriptor(ctx, desc, p); err != nil {
-			return nil, err
-		}
-		return sqlbase.NewImmutableTableDescriptor(*desc), nil
-	default:
-		return objDesc, nil
-	}
-}
-
 // ObjectLookupFlags is part of the resolver.SchemaResolver interface.
 func (p *planner) ObjectLookupFlags(required, requireMutable bool) tree.ObjectLookupFlags {
 	return tree.ObjectLookupFlags{