Merge pull request #131064 from mgartner/backport24.1-130725

release-24.1: opt: do not include virtual computed columns in lock columns

pkg/ccl/logictestccl/testdata/logic_test/read_committed

@@ -538,3 +538,28 @@ user testuser
 
 statement ok
 COMMIT
+
+subtest regression_130661
+
+statement ok
+CREATE TABLE t130661 (
+	id INT PRIMARY KEY NOT NULL,
+	i INT NOT NULL,
+	v INT AS (i + 1) VIRTUAL NOT NULL,
+	FAMILY (id),
+	FAMILY (i)
+)
+
+statement ok
+INSERT INTO t130661 VALUES (1, 10)
+
+statement ok
+BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED
+
+statement ok
+SELECT * FROM t130661 WHERE id = 1 FOR UPDATE
+
+statement ok
+COMMIT
+
+subtest end

pkg/sql/logictest/logic.go

@@ -3502,11 +3502,18 @@ func (t *logicTest) execStatement(stmt logicStatement) (bool, error) {
 	return t.finishExecStatement(stmt, execSQL, res, err)
 }
 
+var uniqueHashPattern = regexp.MustCompile(`UNIQUE.*USING\s+HASH`)
+
 func (t *logicTest) finishExecStatement(
 	stmt logicStatement, execSQL string, res gosql.Result, err error,
 ) (bool, error) {
 	if err == nil {
-		sqlutils.VerifyStatementPrettyRoundtrip(t.t(), stmt.sql)
+		// TODO(#65929, #107398): Roundtrips for unique, hash-sharded indexes do
+		// not work because only unique hash-sharded indexes are allowed, yet we
+		// format them as unique constraints.
+		if !uniqueHashPattern.MatchString(stmt.sql) {
+			sqlutils.VerifyStatementPrettyRoundtrip(t.t(), stmt.sql)
+		}
 	}
 	if err == nil && stmt.expectCount >= 0 {
 		var count int64
@@ -3589,7 +3596,12 @@ func (t *logicTest) execQuery(query logicQuery) error {
 
 func (t *logicTest) finishExecQuery(query logicQuery, rows *gosql.Rows, err error) error {
 	if err == nil {
-		sqlutils.VerifyStatementPrettyRoundtrip(t.t(), query.sql)
+		// TODO(#65929, #107398): Roundtrips for unique, hash-sharded indexes do
+		// not work because only unique hash-sharded indexes are allowed, yet we
+		// format them as unique constraints.
+		if !uniqueHashPattern.MatchString(query.sql) {
+			sqlutils.VerifyStatementPrettyRoundtrip(t.t(), query.sql)
+		}
 
 		// If expecting an error, then read all result rows, since some errors are
 		// only triggered after initial rows are returned.

pkg/sql/opt/exec/execbuilder/testdata/select_for_update_read_committed

@@ -22,6 +22,8 @@ query T
 EXPLAIN (OPT) SELECT aisle FROM supermarket WHERE person = 'matilda' FOR UPDATE
 ----
 lock supermarket
+ ├── key columns: person
+ ├── lock columns: (7-10)
  ├── locking: for-update,durability-guaranteed
  └── scan supermarket
       └── constraint: /1: [/'matilda' - /'matilda']
@@ -57,6 +59,8 @@ update supermarket
            └── subquery
                 └── project
                      └── lock supermarket
+                          ├── key columns: person
+                          ├── lock columns: (19-22)
                           ├── locking: for-update,durability-guaranteed
                           └── scan supermarket
                                └── constraint: /13: [/'matilda' - /'matilda']
@@ -119,6 +123,8 @@ SELECT aisle + 1 FROM s
 with &1 (s)
  ├── project
  │    └── lock supermarket
+ │         ├── key columns: person
+ │         ├── lock columns: (7-10)
  │         ├── locking: for-update,durability-guaranteed
  │         └── scan supermarket
  │              └── constraint: /1: [/'matilda' - /'matilda']
@@ -182,6 +188,8 @@ with &1 (names)
  ├── values
  │    └── ('matilda',)
  └── lock supermarket
+      ├── key columns: supermarket.person
+      ├── lock columns: (9-12)
       ├── locking: for-update,durability-guaranteed
       └── project
            └── inner-join (lookup supermarket)
@@ -254,6 +262,8 @@ SELECT aisle
   FOR UPDATE
 ----
 lock supermarket
+ ├── key columns: person
+ ├── lock columns: (7-10)
  ├── locking: for-update,durability-guaranteed
  └── project
       └── index-join supermarket
@@ -309,6 +319,8 @@ SELECT aisle
   FOR UPDATE
 ----
 lock supermarket
+ ├── key columns: person
+ ├── lock columns: (7-10)
  ├── locking: for-update,durability-guaranteed
  └── project
       └── inner-join (lookup supermarket)

pkg/sql/opt/memo/expr_format.go

@@ -702,6 +702,8 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 		}
 
 	case *LockExpr:
+		f.formatColList(tp, "key columns:", t.KeyCols, opt.ColSet{} /* notNullCols */)
+		tp.Childf("lock columns: %v", t.LockCols)
 		f.formatLocking(tp, t.Locking)
 
 	case *WithExpr:

pkg/sql/opt/norm/testdata/rules/limit

@@ -1598,6 +1598,8 @@ SELECT * FROM abcde WHERE d IS NULL LIMIT 1 FOR UPDATE
 ----
 lock abcde
  ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ ├── key columns: a:1
+ ├── lock columns: (8-12)
  ├── locking: for-update
  ├── cardinality: [0 - 1]
  ├── volatile, mutations
@@ -1619,6 +1621,8 @@ SELECT * FROM abcde WHERE d IS NULL OFFSET 1 FOR UPDATE
 ----
 lock abcde
  ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ ├── key columns: a:1
+ ├── lock columns: (8-12)
  ├── locking: for-update
  ├── volatile, mutations
  ├── key: (1)
@@ -1650,6 +1654,8 @@ limit
  ├── fd: ()-->(1-5)
  ├── lock abcde
  │    ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ │    ├── key columns: a:1
+ │    ├── lock columns: (8-12)
  │    ├── locking: for-update,skip-locked
  │    ├── volatile, mutations
  │    ├── key: (1)
@@ -1680,6 +1686,8 @@ offset
  ├── key: (1)
  ├── lock abcde
  │    ├── columns: a:1!null
+ │    ├── key columns: a:1
+ │    ├── lock columns: (8-12)
  │    ├── locking: for-update,skip-locked
  │    ├── volatile, mutations
  │    ├── key: (1)
@@ -1708,6 +1716,8 @@ SELECT * FROM abcde WHERE d IS NULL LIMIT 1 FOR UPDATE
 ----
 lock abcde
  ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ ├── key columns: a:1
+ ├── lock columns: (8-12)
  ├── locking: for-update,durability-guaranteed
  ├── cardinality: [0 - 1]
  ├── volatile, mutations
@@ -1729,6 +1739,8 @@ SELECT * FROM abcde WHERE d IS NULL OFFSET 1 FOR UPDATE
 ----
 lock abcde
  ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ ├── key columns: a:1
+ ├── lock columns: (8-12)
  ├── locking: for-update,durability-guaranteed
  ├── volatile, mutations
  ├── key: (1)
@@ -1760,6 +1772,8 @@ limit
  ├── fd: ()-->(1-5)
  ├── lock abcde
  │    ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ │    ├── key columns: a:1
+ │    ├── lock columns: (8-12)
  │    ├── locking: for-update,skip-locked,durability-guaranteed
  │    ├── volatile, mutations
  │    ├── key: (1)
@@ -1791,6 +1805,8 @@ offset
  ├── fd: ()-->(4), (1)-->(2,3,5)
  ├── lock abcde
  │    ├── columns: a:1!null b:2 c:3!null d:4 e:5
+ │    ├── key columns: a:1
+ │    ├── lock columns: (8-12)
  │    ├── locking: for-update,skip-locked,durability-guaranteed
  │    ├── volatile, mutations
  │    ├── key: (1)

pkg/sql/opt/optbuilder/locking.go

@@ -438,16 +438,18 @@ func (b *Builder) buildLock(lb *lockBuilder, locking opt.Locking, inScope *scope
 	newTabID := md.DuplicateTable(lb.table, b.factory.RemapCols)
 	newTab := md.Table(newTabID)
 	// Add remapped columns for the new table reference. For now we lock all
-	// column families of the primary index of the table, so include all ordinary
-	// and mutation columns.
+	// non-virtual columns of all families of the primary index of the table, so
+	// include all ordinary and mutation columns.
 	ordinals := tableOrdinals(newTab, columnKinds{
 		includeMutations: true,
 		includeSystem:    false,
 		includeInverted:  false,
 	})
 	var lockCols opt.ColSet
 	for _, ord := range ordinals {
-		lockCols.Add(newTabID.ColumnID(ord))
+		if !tab.Column(ord).IsVirtualComputed() {
+			lockCols.Add(newTabID.ColumnID(ord))
+		}
 	}
 	private := &memo.LockPrivate{
 		Table:     newTabID,

pkg/sql/opt/optbuilder/testdata/read-committed

@@ -0,0 +1,35 @@
+# Tests for read-committed isolation level.
+
+exec-ddl
+CREATE TABLE t130661 (
+	id INT PRIMARY KEY,
+	i INT NOT NULL,
+	v INT AS (i + 10) VIRTUAL NOT NULL
+)
+----
+
+# Regression test for #130661. The lock columns should not include virtual
+# computed columns (in this case the column with ID 8).
+build isolation=ReadCommitted
+SELECT * FROM t130661 WHERE id = 1 FOR UPDATE
+----
+lock t130661
+ ├── columns: id:1!null i:2!null v:3!null
+ ├── key columns: id:1
+ ├── lock columns: (6,7)
+ ├── locking: for-update,durability-guaranteed
+ └── project
+      ├── columns: id:1!null i:2!null v:3!null
+      └── select
+           ├── columns: id:1!null i:2!null v:3!null crdb_internal_mvcc_timestamp:4 tableoid:5
+           ├── project
+           │    ├── columns: v:3!null id:1!null i:2!null crdb_internal_mvcc_timestamp:4 tableoid:5
+           │    ├── scan t130661
+           │    │    ├── columns: id:1!null i:2!null crdb_internal_mvcc_timestamp:4 tableoid:5
+           │    │    └── computed column expressions
+           │    │         └── v:3
+           │    │              └── i:2 + 10
+           │    └── projections
+           │         └── i:2 + 10 [as=v:3]
+           └── filters
+                └── id:1 = 1