schemaexpr: add iterColsWithLookupFn

This function allows the computed column validation to be used from the
declarative schema changer -- it uses a lookup function to access
columns rather than referring to the table descriptor directly.

Release note: None

Author: rafiss

pkg/sql/catalog/schemaexpr/column.go

@@ -221,6 +221,51 @@ func iterColDescriptors(
 	return err
 }
 
+// iterColsWithLookupFn iterates over the expression's variable columns and
+// calls f on each, providing the column information from the lookup function.
+//
+// If the expression references a column that does not exist according to the
+// lookup function, iterColsWithLookupFn errs with pgcode.UndefinedColumn.
+//
+// The column lookup function allows looking up columns both in the descriptor
+// or in declarative schema changer elements.
+func iterColsWithLookupFn(
+	rootExpr tree.Expr,
+	columnLookupFn ColumnLookupFn,
+	f func(columnName tree.Name, id catid.ColumnID, typ *types.T, isAccessible, isComputed bool) error,
+) error {
+	_, err := tree.SimpleVisit(rootExpr, func(expr tree.Expr) (recurse bool, newExpr tree.Expr, err error) {
+		vBase, ok := expr.(tree.VarName)
+		if !ok {
+			// Not a VarName, don't do anything to this node.
+			return true, expr, nil
+		}
+
+		v, err := vBase.NormalizeVarName()
+		if err != nil {
+			return false, nil, err
+		}
+
+		c, ok := v.(*tree.ColumnItem)
+		if !ok {
+			return true, expr, nil
+		}
+
+		colExists, colIsAccessible, isComputed, colID, colType := columnLookupFn(c.ColumnName)
+		if !colExists {
+			return false, nil, pgerror.Newf(pgcode.UndefinedColumn,
+				"column %q does not exist, referenced in %q", c.ColumnName, rootExpr.String())
+		}
+
+		if err := f(c.ColumnName, colID, colType, colIsAccessible, isComputed); err != nil {
+			return false, nil, err
+		}
+		return false, expr, err
+	})
+
+	return err
+}
+
 // dummyColumn represents a variable column that can type-checked. It is used
 // in validating check constraint and partial index predicate expressions. This
 // validation requires that the expression can be both both typed-checked and

pkg/sql/catalog/schemaexpr/computed_column.go

@@ -12,10 +12,12 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/parserutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/transform"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -36,6 +38,9 @@ import (
 //
 //   - It does not have a default value.
 //   - It does not reference other computed columns.
+//   - It does not reference inaccessible columns.
+//   - It does not depend on the region column name if the table is REGIONAL BY
+//     ROW and uses a foreign key to populate the region column.
 //
 // TODO(mgartner): Add unit tests for Validate.
 func ValidateComputedColumnExpression(
@@ -47,69 +52,53 @@ func ValidateComputedColumnExpression(
 	semaCtx *tree.SemaContext,
 	version clusterversion.ClusterVersion,
 ) (serializedExpr string, _ *types.T, _ error) {
-	if d.HasDefaultExpr() {
-		return "", nil, pgerror.Newf(
-			pgcode.InvalidTableDefinition,
-			"%s cannot have default values",
-			context,
-		)
-	}
-
-	var depColIDs catalog.TableColSet
-	// First, check that no column in the expression is an inaccessible or
-	// computed column.
-	err := iterColDescriptors(desc, d.Computed.Expr, func(c catalog.Column) error {
-		if c.IsInaccessible() {
-			return pgerror.Newf(
-				pgcode.UndefinedColumn,
-				"column %q is inaccessible and cannot be referenced in a computed column expression",
-				c.GetName(),
-			)
-		}
-		if c.IsComputed() {
-			return pgerror.Newf(
-				pgcode.InvalidTableDefinition,
-				"%s expression cannot reference computed columns",
-				context,
-			)
+	// Create helper functions from the descriptor and delegate to the
+	// lookup-based validation function.
+	getAllNonDropColumnsFn := func() colinfo.ResultColumns {
+		cols := desc.NonDropColumns()
+		ret := make(colinfo.ResultColumns, len(cols))
+		for i, col := range cols {
+			ret[i] = colinfo.ResultColumn{
+				Name:           col.GetName(),
+				Typ:            col.GetType(),
+				Hidden:         col.IsHidden(),
+				TableID:        desc.GetID(),
+				PGAttributeNum: uint32(col.GetPGAttributeNum()),
+			}
 		}
-		depColIDs.Add(c.GetID())
-
-		return nil
-	})
-	if err != nil {
-		return "", nil, err
+		return ret
 	}
 
-	// Resolve the type of the computed column expression.
-	defType, err := tree.ResolveType(ctx, d.Type, semaCtx.GetTypeResolver())
-	if err != nil {
-		return "", nil, err
-	}
+	columnLookupFn := makeColumnLookupFnForTableDesc(desc)
 
-	// Check that the type of the expression is of type defType and that there
-	// are no variable expressions (besides dummyColumnItems) and no impure
-	// functions. In order to safely serialize user defined types and their
-	// members, we need to serialize the typed expression here.
-	expr, typ, _, err := DequalifyAndValidateExpr(
+	serializedExpr, typ, err := ValidateComputedColumnExpressionWithLookup(
 		ctx,
 		desc,
-		d.Computed.Expr,
-		defType,
+		d,
+		tn,
 		context,
 		semaCtx,
-		volatility.Immutable,
-		tn,
 		version,
+		getAllNonDropColumnsFn,
+		columnLookupFn,
 	)
 	if err != nil {
 		return "", nil, err
 	}
 
+	var depColIDs catalog.TableColSet
+	if err := iterColDescriptors(desc, d.Computed.Expr, func(c catalog.Column) error {
+		depColIDs.Add(c.GetID())
+		return nil
+	}); err != nil {
+		return "", nil, err
+	}
+
 	// Virtual computed columns must not refer to mutation columns because it
 	// would not be safe in the case that the mutation column was being
 	// backfilled and the virtual computed column value needed to be computed
 	// for the purpose of writing to a secondary index.
+	// This check is specific to the legacy schema changer.
 	if d.IsVirtual() {
 		var mutationColumnNames []string
 		var err error
@@ -144,6 +133,86 @@ func ValidateComputedColumnExpression(
 		}
 	}
 
+	return serializedExpr, typ, nil
+}
+
+// ValidateComputedColumnExpressionWithLookup verifies that an expression is a
+// valid computed column expression using a column lookup function. It returns
+// the serialized expression and its type if valid, and an error otherwise.
+//
+// This is similar to ValidateComputedColumnExpression but uses a ColumnLookupFn
+// instead of a catalog.TableDescriptor, allowing it to work with declarative
+// schema changer elements.
+func ValidateComputedColumnExpressionWithLookup(
+	ctx context.Context,
+	desc catalog.TableDescriptor,
+	d *tree.ColumnTableDef,
+	tn *tree.TableName,
+	context tree.SchemaExprContext,
+	semaCtx *tree.SemaContext,
+	version clusterversion.ClusterVersion,
+	getAllNonDropColumnsFn func() colinfo.ResultColumns,
+	columnLookupFn ColumnLookupFn,
+) (serializedExpr string, _ *types.T, _ error) {
+	if d.HasDefaultExpr() {
+		return "", nil, pgerror.Newf(
+			pgcode.InvalidTableDefinition,
+			"%s cannot have default values",
+			context,
+		)
+	}
+
+	var depColIDs catalog.TableColSet
+	// First, check that no column in the expression is an inaccessible or
+	// computed column.
+	err := iterColsWithLookupFn(d.Computed.Expr, columnLookupFn,
+		func(columnName tree.Name, id catid.ColumnID, typ *types.T, isAccessible, isComputed bool) error {
+			if !isAccessible {
+				return pgerror.Newf(
+					pgcode.UndefinedColumn,
+					"column %q is inaccessible and cannot be referenced in a computed column expression",
+					columnName,
+				)
+			}
+			if isComputed {
+				return pgerror.Newf(
+					pgcode.InvalidTableDefinition,
+					"%s expression cannot reference computed columns",
+					context,
+				)
+			}
+			depColIDs.Add(id)
+			return nil
+		})
+	if err != nil {
+		return "", nil, err
+	}
+
+	// Resolve the type of the computed column expression.
+	defType, err := tree.ResolveType(ctx, d.Type, semaCtx.GetTypeResolver())
+	if err != nil {
+		return "", nil, err
+	}
+
+	// Check that the type of the expression is of type defType and that there
+	// are no variable expressions (besides dummyColumnItems) and no impure
+	// functions. We use DequalifyAndValidateExprImpl with the lookup function.
+	expr, typ, _, err := DequalifyAndValidateExprImpl(
+		ctx,
+		d.Computed.Expr,
+		defType,
+		context,
+		semaCtx,
+		volatility.Immutable,
+		tn,
+		version,
+		getAllNonDropColumnsFn,
+		columnLookupFn,
+	)
+	if err != nil {
+		return "", nil, err
+	}
+
 	// If this is a REGIONAL BY ROW table using a foreign key to populate the
 	// region column, we need to check that the expression does not reference
 	// the region column. This is because the values of every (possibly computed)