scbuild: use ValidateComputedColumnExpressionWithLookup function

This validates the computed column by referencing the builder state
rather than the table descriptor. The table descriptor would not reflect
the schema changes that are being staged by the declarative schema
changer.

Release note: None

Author: rafiss

pkg/sql/schemachanger/scbuild/builder_state.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"sort"
 
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
@@ -838,29 +839,60 @@ func (b *builderState) WrapExpression(tableID catid.DescID, expr tree.Expr) *scp
 
 // ComputedColumnExpression implements the scbuildstmt.TableHelpers interface.
 func (b *builderState) ComputedColumnExpression(
-	tbl *scpb.Table, d *tree.ColumnTableDef, exprContext tree.SchemaExprContext,
+	tbl *scpb.Table,
+	d *tree.ColumnTableDef,
+	exprContext tree.SchemaExprContext,
+	getAllNonDropColumnsFn func() colinfo.ResultColumns,
+	columnLookupByNameFn schemaexpr.ColumnLookupFn,
 ) (tree.Expr, *types.T) {
 	_, _, ns := scpb.FindNamespace(b.QueryByID(tbl.TableID))
 	tn := tree.MakeTableNameFromPrefix(b.NamePrefix(tbl), tree.Name(ns.Name))
 	b.ensureDescriptor(tbl.TableID)
-	// TODO(postamar): this doesn't work when referencing newly added columns.
-	expr, typ, err := schemaexpr.ValidateComputedColumnExpression(
+
+	// In versions before 26.1, computed columns referencing newly added columns
+	// are not supported in the declarative schema changer. Fall back to the
+	// legacy schema changer for those cases.
+	if !b.clusterSettings.Version.IsActive(b.ctx, clusterversion.V26_1) {
+		// Use the old validation logic that doesn't support newly added columns.
+		expr, typ, err := schemaexpr.ValidateComputedColumnExpression(
+			b.ctx,
+			b.descCache[tbl.TableID].desc.(catalog.TableDescriptor),
+			d,
+			&tn,
+			exprContext,
+			b.semaCtx,
+			b.clusterSettings.Version.ActiveVersion(b.ctx),
+		)
+		if err != nil {
+			// This may be referencing newly added columns, so return a
+			// NotImplementedError to force fallback to the legacy schema changer.
+			if pgerror.GetPGCode(err) == pgcode.UndefinedColumn {
+				panic(errors.Wrapf(errors.WithSecondaryError(scerrors.NotImplementedError(d), err),
+					"computed column validation error"))
+			}
+			panic(err)
+		}
+		parsedExpr, err := parser.ParseExpr(expr)
+		if err != nil {
+			panic(err)
+		}
+		return parsedExpr, typ
+	}
+
+	// In 26.1+, we can validate computed columns that reference newly added
+	// columns by using the lookup functions to query the builder state.
+	expr, typ, err := schemaexpr.ValidateComputedColumnExpressionWithLookup(
 		b.ctx,
 		b.descCache[tbl.TableID].desc.(catalog.TableDescriptor),
 		d,
 		&tn,
 		exprContext,
 		b.semaCtx,
 		b.clusterSettings.Version.ActiveVersion(b.ctx),
+		getAllNonDropColumnsFn,
+		columnLookupByNameFn,
 	)
 	if err != nil {
-		// This may be referencing newly added columns, so cheat and return
-		// a not implemented error.
-		if pgerror.GetPGCode(err) == pgcode.UndefinedColumn {
-
-			panic(errors.Wrapf(errors.WithSecondaryError(scerrors.NotImplementedError(d), err),
-				"computed column validation error"))
-		}
 		panic(err)
 	}
 	parsedExpr, err := parser.ParseExpr(expr)

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_column.go

@@ -167,7 +167,15 @@ func alterTableAddColumn(
 		))
 	}
 	if desc.IsComputed() {
-		validExpr, _ := b.ComputedColumnExpression(tbl, d, tree.ComputedColumnExprContext(d.IsVirtual()))
+		validExpr, _ := b.ComputedColumnExpression(
+			tbl, d, tree.ComputedColumnExprContext(d.IsVirtual()),
+			func() colinfo.ResultColumns {
+				return getNonDropResultColumns(b, tbl.TableID)
+			},
+			func(columnName tree.Name) (exists, accessible, computed bool, id catid.ColumnID, typ *types.T) {
+				return columnLookupFn(b, tbl.TableID, columnName)
+			},
+		)
 		expr := b.WrapExpression(tbl.TableID, validExpr)
 		if spec.colType.ElementCreationMetadata.In_24_3OrLater {
 			spec.compute = &scpb.ColumnComputeExpression{

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/create_index.go

@@ -978,7 +978,15 @@ func maybeCreateVirtualColumnForIndex(
 	d.Nullable.Nullability = tree.Null
 	// Infer column type from expression.
 	{
-		_, columnType := b.ComputedColumnExpression(tbl, d, tree.ExpressionIndexElementExpr)
+		_, columnType := b.ComputedColumnExpression(
+			tbl, d, tree.ExpressionIndexElementExpr,
+			func() colinfo.ResultColumns {
+				return getNonDropResultColumns(b, tbl.TableID)
+			},
+			func(columnName tree.Name) (exists, accessible, computed bool, id catid.ColumnID, typ *types.T) {
+				return columnLookupFn(b, tbl.TableID, columnName)
+			},
+		)
 		d.Type = columnType
 		validateColumnIndexableType(columnType)
 	}

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/dependencies.go

@@ -15,8 +15,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"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/multiregion"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemaexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scdecomp"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
@@ -313,7 +315,11 @@ type TableHelpers interface {
 	// ComputedColumnExpression returns a validated computed column expression
 	// and its type.
 	// TODO(postamar): make this more low-level instead of consuming an AST
-	ComputedColumnExpression(tbl *scpb.Table, d *tree.ColumnTableDef, exprContext tree.SchemaExprContext) (tree.Expr, *types.T)
+	ComputedColumnExpression(
+		tbl *scpb.Table, d *tree.ColumnTableDef, exprContext tree.SchemaExprContext,
+		getAllNonDropColumnsFn func() colinfo.ResultColumns,
+		columnLookupByNameFn schemaexpr.ColumnLookupFn,
+	) (tree.Expr, *types.T)
 
 	// PartialIndexPredicateExpression returns a validated partial predicate
 	// wrapped expression