opt: Fix UPDATE error when constraint is used with computed column

The following statement fails with an ambiguous column name error:

  CREATE TABLE t (a INT, b INT AS (a + 1) STORED, CONSTRAINT ck CHECK (b > 0));
  UPDATE t SET a = 123 WHERE a = 456;

The problem is that the input scope for the constraint expression
contains two versions of column "b", namely the column containing
the existing value of "b" and the column containing the updated value
of "b". The fix is to clear the name of the existing value so that the
constraint expression binds unambiguously to the updated value.

Fixes #34649

Release note: None

pkg/sql/opt/norm/prune_cols.go

@@ -128,6 +128,9 @@ func (c *CustomFuncs) NeededMutationFetchCols(
 			// It is possible to update a subset of families only for the primary
 			// index, and only when key columns are not being updated. Otherwise,
 			// all columns in the index must be fetched.
+			// TODO(andyk): It should be possible to not include columns that are
+			// being updated, since the existing value is not used. However, this
+			// would require execution support.
 			if i == cat.PrimaryIndex && !keyCols.Intersects(updateCols) {
 				addFamilyCols(updateCols)
 			} else {

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -374,6 +374,10 @@ func (mb *mutationBuilder) addSynthesizedCols(
 // a constraint violation error if the value of the column is false.
 func (mb *mutationBuilder) addCheckConstraintCols() {
 	if mb.tab.CheckCount() > 0 {
+		// Disambiguate names so that references in the constraint expression refer
+		// to the correct columns.
+		mb.disambiguateColumns()
+
 		mb.checkColList = make(opt.ColList, mb.tab.CheckCount())
 		projectionsScope := mb.outScope.replace()
 		projectionsScope.appendColumnsFromScope(mb.outScope)
@@ -396,6 +400,29 @@ func (mb *mutationBuilder) addCheckConstraintCols() {
 	}
 }
 
+// disambiguateColumns ranges over the scope and ensures that at most one column
+// has each table column name, and that name refers to the column with the final
+// value that the mutation applies.
+func (mb *mutationBuilder) disambiguateColumns() {
+	for i, n := 0, mb.tab.DeletableColumnCount(); i < n; i++ {
+		colName := mb.tab.Column(i).ColName()
+		colID := mb.mapToReturnColID(i)
+		for i := range mb.outScope.cols {
+			col := &mb.outScope.cols[i]
+			if col.name == colName {
+				if col.id == colID {
+					// Use table name, not alias name, since computed column
+					// expressions will not reference aliases.
+					col.table = *mb.tab.Name()
+				} else {
+					// Clear name so that it will never match.
+					col.clearName()
+				}
+			}
+		}
+	}
+}
+
 // makeMutationPrivate builds a MutationPrivate struct containing the table and
 // column metadata needed for the mutation operator.
 func (mb *mutationBuilder) makeMutationPrivate(needResults bool) *memo.MutationPrivate {
@@ -414,27 +441,41 @@ func (mb *mutationBuilder) makeMutationPrivate(needResults bool) *memo.MutationP
 		// can be non-zero.
 		private.ReturnCols = make(opt.ColList, mb.tab.DeletableColumnCount())
 		for i, n := 0, mb.tab.ColumnCount(); i < n; i++ {
-			// Map to columns in this order: upsert, update, fetch, insert.
-			switch {
-			case mb.upsertColList != nil && mb.upsertColList[i] != 0:
-				private.ReturnCols[i] = mb.upsertColList[i]
+			private.ReturnCols[i] = mb.mapToReturnColID(i)
+		}
+	}
 
-			case mb.updateColList != nil && mb.updateColList[i] != 0:
-				private.ReturnCols[i] = mb.updateColList[i]
+	return private
+}
 
-			case mb.fetchColList != nil && mb.fetchColList[i] != 0:
-				private.ReturnCols[i] = mb.fetchColList[i]
+// mapToReturnColID returns the ID of the input column that will provide the
+// value for the corresponding return column. Columns take priority in this
+// order:
+//
+//   upsert, update, fetch, insert
+//
+// If an upsert column is available, then it already combines an update/fetch
+// value with an insert value, so it takes priority. If an update column is
+// available, then it overrides any fetch value. Finally, the relative priority
+// of fetch and insert columns doesn't matter, since they're only used together
+// in the upsert case where an upsert column would be available.
+func (mb *mutationBuilder) mapToReturnColID(ord int) opt.ColumnID {
+	switch {
+	case mb.upsertColList != nil && mb.upsertColList[ord] != 0:
+		return mb.upsertColList[ord]
 
-			case mb.insertColList != nil && mb.insertColList[i] != 0:
-				private.ReturnCols[i] = mb.insertColList[i]
+	case mb.updateColList != nil && mb.updateColList[ord] != 0:
+		return mb.updateColList[ord]
 
-			default:
-				panic("could not find return column")
-			}
-		}
-	}
+	case mb.fetchColList != nil && mb.fetchColList[ord] != 0:
+		return mb.fetchColList[ord]
 
-	return private
+	case mb.insertColList != nil && mb.insertColList[ord] != 0:
+		return mb.insertColList[ord]
+
+	default:
+		panic("could not find return column")
+	}
 }
 
 // buildReturning wraps the input expression with a Project operator that

pkg/sql/opt/optbuilder/testdata/insert

@@ -67,16 +67,18 @@ CREATE TABLE checks (
     a INT PRIMARY KEY CHECK (a > 0),
     b INT,
     c INT,
-    CHECK (checks.b < c)
+    d INT AS (c + 1) STORED,
+    CHECK (checks.b < d)
 )
 ----
 TABLE checks
  ├── a int not null
  ├── b int
  ├── c int
+ ├── d int
  ├── INDEX primary
  │    └── a int not null
- ├── CHECK (checks.b < c)
+ ├── CHECK (checks.b < d)
  └── CHECK (a > 0)
 
 # Unknown target table.
@@ -1256,22 +1258,29 @@ INSERT INTO checks (a, b, c) VALUES (1, 2, 3)
 insert checks
  ├── columns: <none>
  ├── insert-mapping:
- │    ├──  column1:4 => a:1
- │    ├──  column2:5 => b:2
- │    └──  column3:6 => c:3
- ├── check columns: check1:7(bool) check2:8(bool)
+ │    ├──  column1:5 => a:1
+ │    ├──  column2:6 => b:2
+ │    ├──  column3:7 => c:3
+ │    └──  column8:8 => d:4
+ ├── check columns: check1:9(bool) check2:10(bool)
  └── project
-      ├── columns: check1:7(bool) check2:8(bool) column1:4(int) column2:5(int) column3:6(int)
-      ├── values
-      │    ├── columns: column1:4(int) column2:5(int) column3:6(int)
-      │    └── tuple [type=tuple{int, int, int}]
-      │         ├── const: 1 [type=int]
-      │         ├── const: 2 [type=int]
-      │         └── const: 3 [type=int]
+      ├── columns: check1:9(bool) check2:10(bool) column1:5(int) column2:6(int) column3:7(int) column8:8(int)
+      ├── project
+      │    ├── columns: column8:8(int) column1:5(int) column2:6(int) column3:7(int)
+      │    ├── values
+      │    │    ├── columns: column1:5(int) column2:6(int) column3:7(int)
+      │    │    └── tuple [type=tuple{int, int, int}]
+      │    │         ├── const: 1 [type=int]
+      │    │         ├── const: 2 [type=int]
+      │    │         └── const: 3 [type=int]
+      │    └── projections
+      │         └── plus [type=int]
+      │              ├── variable: column3 [type=int]
+      │              └── const: 1 [type=int]
       └── projections
            ├── lt [type=bool]
            │    ├── variable: column2 [type=int]
-           │    └── variable: column3 [type=int]
+           │    └── variable: column8 [type=int]
            └── gt [type=bool]
                 ├── variable: column1 [type=int]
                 └── const: 0 [type=int]
@@ -1283,20 +1292,27 @@ INSERT INTO checks SELECT a, b, c FROM abcde
 insert checks
  ├── columns: <none>
  ├── insert-mapping:
- │    ├──  abcde.a:4 => checks.a:1
- │    ├──  abcde.b:5 => checks.b:2
- │    └──  abcde.c:6 => checks.c:3
- ├── check columns: check1:10(bool) check2:11(bool)
+ │    ├──  abcde.a:5 => checks.a:1
+ │    ├──  abcde.b:6 => checks.b:2
+ │    ├──  abcde.c:7 => checks.c:3
+ │    └──  column11:11 => checks.d:4
+ ├── check columns: check1:12(bool) check2:13(bool)
  └── project
-      ├── columns: check1:10(bool) check2:11(bool) abcde.a:4(int!null) abcde.b:5(int) abcde.c:6(int)
+      ├── columns: check1:12(bool) check2:13(bool) abcde.a:5(int!null) abcde.b:6(int) abcde.c:7(int) column11:11(int)
       ├── project
-      │    ├── columns: abcde.a:4(int!null) abcde.b:5(int) abcde.c:6(int)
-      │    └── scan abcde
-      │         └── columns: abcde.a:4(int!null) abcde.b:5(int) abcde.c:6(int) d:7(int) e:8(int) rowid:9(int!null)
+      │    ├── columns: column11:11(int) abcde.a:5(int!null) abcde.b:6(int) abcde.c:7(int)
+      │    ├── project
+      │    │    ├── columns: abcde.a:5(int!null) abcde.b:6(int) abcde.c:7(int)
+      │    │    └── scan abcde
+      │    │         └── columns: abcde.a:5(int!null) abcde.b:6(int) abcde.c:7(int) abcde.d:8(int) e:9(int) rowid:10(int!null)
+      │    └── projections
+      │         └── plus [type=int]
+      │              ├── variable: abcde.c [type=int]
+      │              └── const: 1 [type=int]
       └── projections
            ├── lt [type=bool]
            │    ├── variable: abcde.b [type=int]
-           │    └── variable: abcde.c [type=int]
+           │    └── variable: column11 [type=int]
            └── gt [type=bool]
                 ├── variable: abcde.a [type=int]
                 └── const: 0 [type=int]