Merge #36518

36518: sql: validate computed columns during ALTER TABLE r=vivekmenezes a=mjibson

Previously these were only validated during CREATE, allowing the creation
of incorrect columns. validateComputedColumn has removed its dependency
on the tree.CreateTable struct so that the alter code can use it. Because
of this, all FKs must be computed prior to the computed validate checks.

Fixes #36036

Release note (bug fix): Correctly validate computed columns during ALTER
TABLE ADD COLUMN.

Co-authored-by: Matt Jibson <[email protected]>

pkg/sql/alter_table.go

@@ -169,6 +169,12 @@ func (n *alterTableNode) startExec(params runParams) error {
 				}
 			}
 
+			if d.IsComputed() {
+				if err := validateComputedColumn(n.tableDesc, d, &params.p.semaCtx); err != nil {
+					return err
+				}
+			}
+
 		case *tree.AlterTableAddConstraint:
 			info, err := n.tableDesc.GetConstraintInfo(params.ctx, nil)
 			if err != nil {

pkg/sql/create_table.go

@@ -1188,13 +1188,7 @@ func MakeTableDesc(
 	for _, def := range n.Defs {
 		switch d := def.(type) {
 		case *tree.ColumnTableDef:
-			// Now that we have all the other columns set up, we can validate any
-			// computed columns.
-			if d.IsComputed() {
-				if err := validateComputedColumn(&desc, n, d, semaCtx); err != nil {
-					return desc, err
-				}
-			}
+			// Check after all ResolveFK calls.
 
 		case *tree.IndexTableDef, *tree.UniqueConstraintTableDef, *tree.FamilyTableDef:
 			// Pass, handled above.
@@ -1210,10 +1204,23 @@ func MakeTableDesc(
 			if err := ResolveFK(ctx, txn, fkResolver, &desc, d, affected, NewTable); err != nil {
 				return desc, err
 			}
+
 		default:
 			return desc, errors.Errorf("unsupported table def: %T", def)
 		}
 	}
+	// Now that we have all the other columns set up, we can validate
+	// any computed columns.
+	for _, def := range n.Defs {
+		switch d := def.(type) {
+		case *tree.ColumnTableDef:
+			if d.IsComputed() {
+				if err := validateComputedColumn(&desc, d, semaCtx); err != nil {
+					return desc, err
+				}
+			}
+		}
+	}
 
 	// AllocateIDs mutates its receiver. `return desc, desc.AllocateIDs()`
 	// happens to work in gc, but does not work in gccgo.
@@ -1397,10 +1404,7 @@ func iterColDescriptorsInExpr(
 // validateComputedColumn checks that a computed column satisfies a number of
 // validity constraints, for instance, that it typechecks.
 func validateComputedColumn(
-	desc *sqlbase.MutableTableDescriptor,
-	t *tree.CreateTable,
-	d *tree.ColumnTableDef,
-	semaCtx *tree.SemaContext,
+	desc *sqlbase.MutableTableDescriptor, d *tree.ColumnTableDef, semaCtx *tree.SemaContext,
 ) error {
 	if d.HasDefaultExpr() {
 		return pgerror.NewError(
@@ -1426,24 +1430,28 @@ func validateComputedColumn(
 	// TODO(justin,bram): allow depending on columns like this. We disallow it
 	// for now because cascading changes must hook into the computed column
 	// update path.
-	for _, def := range t.Defs {
-		switch c := def.(type) {
-		case *tree.ColumnTableDef:
-			if _, ok := dependencies[string(c.Name)]; !ok {
+	if err := desc.ForeachNonDropIndex(func(idx *sqlbase.IndexDescriptor) error {
+		for _, name := range idx.ColumnNames {
+			if _, ok := dependencies[name]; !ok {
 				// We don't depend on this column.
 				continue
 			}
-			for _, action := range []tree.ReferenceAction{
-				c.References.Actions.Update,
-				c.References.Actions.Delete,
+			for _, action := range []sqlbase.ForeignKeyReference_Action{
+				idx.ForeignKey.OnDelete,
+				idx.ForeignKey.OnUpdate,
 			} {
 				switch action {
-				case tree.Cascade, tree.SetNull, tree.SetDefault:
+				case sqlbase.ForeignKeyReference_CASCADE,
+					sqlbase.ForeignKeyReference_SET_NULL,
+					sqlbase.ForeignKeyReference_SET_DEFAULT:
 					return pgerror.NewError(pgerror.CodeInvalidTableDefinitionError,
 						"computed columns cannot reference non-restricted FK columns")
 				}
 			}
 		}
+		return nil
+	}); err != nil {
+		return err
 	}
 
 	// Replace column references with typed dummies to allow typechecking.

pkg/sql/logictest/testdata/logic_test/computed

@@ -351,6 +351,16 @@ CREATE TABLE y (
   r INT AS (1) STORED REFERENCES x
 )
 
+# Regression test for #36036.
+statement ok
+CREATE TABLE tt (i INT8 AS (1) STORED)
+
+statement error variable sub-expressions are not allowed in computed column
+ALTER TABLE tt ADD COLUMN c STRING AS ((SELECT NULL)) STORED
+
+statement error computed columns cannot reference other computed columns
+ALTER TABLE tt ADD COLUMN c INT8 AS (i) STORED
+
 # Composite FK.
 
 statement ok

pkg/sql/logictest/testdata/logic_test/schema_change_in_txn

@@ -739,7 +739,7 @@ statement ok
 ALTER TABLE customers ADD j INT AS (i-1) STORED
 
 statement ok
-ALTER TABLE customers ADD COLUMN d INT DEFAULT 15, ADD COLUMN e INT AS (d + j) STORED
+ALTER TABLE customers ADD COLUMN d INT DEFAULT 15, ADD COLUMN e INT AS (d + (i-1)) STORED
 
 statement ok
 COMMIT
@@ -755,7 +755,7 @@ query TT
 SELECT status, description FROM [SHOW JOBS]
 WHERE job_type = 'SCHEMA CHANGE' ORDER BY job_id DESC LIMIT 1
 ----
-succeeded  ALTER TABLE test.public.customers ADD COLUMN i INT8 DEFAULT 5;ALTER TABLE test.public.customers ADD COLUMN j INT8 AS (i - 1) STORED;ALTER TABLE test.public.customers ADD COLUMN d INT8 DEFAULT 15, ADD COLUMN e INT8 AS (d + j) STORED
+succeeded  ALTER TABLE test.public.customers ADD COLUMN i INT8 DEFAULT 5;ALTER TABLE test.public.customers ADD COLUMN j INT8 AS (i - 1) STORED;ALTER TABLE test.public.customers ADD COLUMN d INT8 DEFAULT 15, ADD COLUMN e INT8 AS (d + (i - 1)) STORED
 
 # VALIDATE CONSTRAINT will not hang when executed in the same txn as
 # a schema change in the same txn #32118