Merge #144047

144047: sql: fix handle index expressions in declarative schema changer r=fqazi a=fqazi

Previously, the declarative schema changer did not correctly handle index expressions, and would end up falling back. This happened because when evaluating the expression no type was specified, so the logic would be quietly skipped because of fallback for an add column case. To address this, this patch fixes the logic to resolve the type correctly when adding columns for index expressions.

Fixes: #143891
Release note: None

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

pkg/ccl/schemachangerccl/backup_base_generated_test.go

@@ -838,31 +838,37 @@ 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) tree.Expr {
+func (b *builderState) ComputedColumnExpression(
+	tbl *scpb.Table, d *tree.ColumnTableDef, exprContext tree.SchemaExprContext,
+) (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, _, err := schemaexpr.ValidateComputedColumnExpression(
+	expr, typ, err := schemaexpr.ValidateComputedColumnExpression(
 		b.ctx,
 		b.descCache[tbl.TableID].desc.(catalog.TableDescriptor),
 		d,
 		&tn,
-		tree.ComputedColumnExprContext(d.IsVirtual()),
+		exprContext,
 		b.semaCtx,
 		b.clusterSettings.Version.ActiveVersion(b.ctx),
 	)
 	if err != nil {
 		// This may be referencing newly added columns, so cheat and return
 		// a not implemented error.
-		panic(errors.Wrapf(errors.WithSecondaryError(scerrors.NotImplementedError(d), err),
-			"computed column validation 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)
 	if err != nil {
 		panic(err)
 	}
-	return parsedExpr
+	return parsedExpr, typ
 }
 
 // PartialIndexPredicateExpression implements the scbuildstmt.TableHelpers interface.

pkg/sql/schemachanger/scbuild/builder_state.go

@@ -173,7 +173,8 @@ func alterTableAddColumn(
 		))
 	}
 	if desc.IsComputed() {
-		expr := b.WrapExpression(tbl.TableID, b.ComputedColumnExpression(tbl, d))
+		validExpr, _ := b.ComputedColumnExpression(tbl, d, tree.ComputedColumnExprContext(d.IsVirtual()))
+		expr := b.WrapExpression(tbl.TableID, validExpr)
 		if spec.colType.ElementCreationMetadata.In_24_3OrLater {
 			spec.compute = &scpb.ColumnComputeExpression{
 				TableID:    tbl.TableID,

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

@@ -860,20 +860,17 @@ func maybeCreateVirtualColumnForIndex(
 	// TODO(postamar): call addColumn instead of building AST.
 	d := &tree.ColumnTableDef{
 		Name: tree.Name(colName),
+		Type: types.AnyElement,
 	}
 	d.Computed.Computed = true
 	d.Computed.Virtual = true
 	d.Computed.Expr = expr
 	d.Nullable.Nullability = tree.Null
 	// Infer column type from expression.
 	{
-		replacedExpr := b.ComputedColumnExpression(tbl, d)
-		typedExpr, err := tree.TypeCheck(b, replacedExpr, b.SemaCtx(), types.AnyElement)
-		if err != nil {
-			panic(err)
-		}
-		d.Type = typedExpr.ResolvedType()
-		validateColumnIndexableType(typedExpr.ResolvedType())
+		_, columnType := b.ComputedColumnExpression(tbl, d, tree.ExpressionIndexElementExpr)
+		d.Type = columnType
+		validateColumnIndexableType(columnType)
 	}
 	alterTableAddColumn(b, tn, tbl, stmt, &tree.AlterTableAddColumn{ColumnDef: d})
 	// When a virtual column for an index expression gets added for CREATE INDEX

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

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"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/log/logpb"
 )
 
@@ -312,7 +313,7 @@ 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) tree.Expr
+	ComputedColumnExpression(tbl *scpb.Table, d *tree.ColumnTableDef, exprContext tree.SchemaExprContext) (tree.Expr, *types.T)
 
 	// PartialIndexPredicateExpression returns a validated partial predicate
 	// wrapped expression

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

@@ -97,10 +97,10 @@ DROP INDEX idx3 CASCADE
   {indexId: 6, name: idx3, tableId: 104}
 - [[IndexData:{DescID: 104, IndexID: 6}, ABSENT], PUBLIC]
   {indexId: 6, tableId: 104}
-- [[CheckConstraint:{DescID: 104, IndexID: 0, ConstraintID: 2, ReferencedColumnIDs: [5]}, ABSENT], PUBLIC]
-  {columnIds: [5], constraintId: 2, expr: 'crdb_internal_i_shard_16 IN (0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15)', fromHashShardedColumn: true, referencedColumnIds: [5], tableId: 104}
-- [[ConstraintWithoutIndexName:{DescID: 104, Name: check_crdb_internal_i_shard_16, ConstraintID: 2}, ABSENT], PUBLIC]
-  {constraintId: 2, name: check_crdb_internal_i_shard_16, tableId: 104}
+- [[CheckConstraint:{DescID: 104, IndexID: 0, ConstraintID: 3, ReferencedColumnIDs: [5]}, ABSENT], PUBLIC]
+  {columnIds: [5], constraintId: 3, expr: 'crdb_internal_i_shard_16 IN (0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15)', fromHashShardedColumn: true, referencedColumnIds: [5], tableId: 104}
+- [[ConstraintWithoutIndexName:{DescID: 104, Name: check_crdb_internal_i_shard_16, ConstraintID: 3}, ABSENT], PUBLIC]
+  {constraintId: 3, name: check_crdb_internal_i_shard_16, tableId: 104}
 - [[TableData:{DescID: 104, ReferencedDescID: 100}, PUBLIC], PUBLIC]
   {databaseId: 100, tableId: 104}
 - [[Namespace:{DescID: 105, Name: v, ReferencedDescID: 100}, ABSENT], PUBLIC]