Merge #65100 #65283 #65332

65100: schemaexpr: remove unnecessary structs r=mgartner a=mgartner

#### schemaexpr: remove ComputedColumnValidator struct

This commit removes `schemaexpr.ComputedColumnValidator`. The struct's
methods did not store any state in the struct, so they have been
converted to functions without a receiver.

Release note: None

#### schemaexpr: remove IndexPredicateValidator struct

This commit removes `schemaexpr.IndexPredicateValidator`. The struct's
methods did not store any state in the struct, so they have been
converted to functions without a receiver.

Release note: None


65283: build: add pipefail to teamcity-diagram-generation.sh r=jlinder a=RichardJCai

Release note: None

65332: sql/opt: support implicit SELECT FOR UPDATE locking with nested renders r=nvanbenschoten a=nvanbenschoten

This commit adds support for implicit SELECT FOR UPDATE to a collection of UPDATE and UPSERT queries that previously did not use row-level locking because they included 2 or more nested `ProjExprs`. The pattern matching that enabled implicit SELECT FOR UPDATE was not handling this case correctly.

The most common case where this had an effect was with
```sql
INSERT INTO ... ON CONFLICT ... DO UPDATE SET
```
statements without a predicate. This class of statement always seems to include nested render expressions, and so it wasn't acquiring row-level locks during its initial row scan.

Release note (sql change): `INSERT INTO ... ON CONFLICT ... DO UPDATE SET` statements without predicates now acquire locks using the FOR UPDATE locking mode during their initial row scan, which improves performance for contended workloads. This behavior is configurable using the enable_implicit_select_for_update session variable and the sql.defaults.implicit_select_for_update.enabled cluster setting.

Co-authored-by: Marcus Gartner <[email protected]>
Co-authored-by: richardjcai <[email protected]>
Co-authored-by: Nathan VanBenschoten <[email protected]>

build/teamcity-diagram-generation.sh

@@ -1,5 +1,7 @@
 #!/usr/bin/env bash
 
+set -euo pipefail
+
 source "$(dirname "${0}")/teamcity-support.sh"
 
 tc_start_block "Get Railroad Jar"

pkg/sql/add_column.go

@@ -157,13 +157,9 @@ func (p *planner) addColumnImpl(
 	}
 
 	if d.IsComputed() {
-		computedColValidator := schemaexpr.MakeComputedColumnValidator(
-			params.ctx,
-			n.tableDesc,
-			&params.p.semaCtx,
-			tn,
+		serializedExpr, err := schemaexpr.ValidateComputedColumnExpression(
+			params.ctx, n.tableDesc, d, tn, params.p.SemaCtx(),
 		)
-		serializedExpr, err := computedColValidator.Validate(d)
 		if err != nil {
 			return err
 		}

pkg/sql/alter_table.go

@@ -497,13 +497,7 @@ func (n *alterTableNode) startExec(params runParams) error {
 			}
 
 			// We cannot remove this column if there are computed columns that use it.
-			computedColValidator := schemaexpr.MakeComputedColumnValidator(
-				params.ctx,
-				n.tableDesc,
-				&params.p.semaCtx,
-				tn,
-			)
-			if err := computedColValidator.ValidateNoDependents(colToDrop); err != nil {
+			if err := schemaexpr.ValidateColumnHasNoDependents(n.tableDesc, colToDrop); err != nil {
 				return err
 			}
 

pkg/sql/catalog/schemaexpr/computed_column.go

@@ -26,29 +26,7 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
-// ComputedColumnValidator validates that an expression is a valid computed
-// column. See Validate for more details.
-type ComputedColumnValidator struct {
-	ctx       context.Context
-	desc      catalog.TableDescriptor
-	semaCtx   *tree.SemaContext
-	tableName *tree.TableName
-}
-
-// MakeComputedColumnValidator returns an ComputedColumnValidator struct that
-// can be used to validate computed columns. See Validate for more details.
-func MakeComputedColumnValidator(
-	ctx context.Context, desc catalog.TableDescriptor, semaCtx *tree.SemaContext, tn *tree.TableName,
-) ComputedColumnValidator {
-	return ComputedColumnValidator{
-		ctx:       ctx,
-		desc:      desc,
-		semaCtx:   semaCtx,
-		tableName: tn,
-	}
-}
-
-// Validate verifies that an expression is a valid computed column expression.
+// ValidateComputedColumnExpression verifies that an expression is a valid computed column expression.
 // It returns the serialized expression if valid, and an error otherwise.
 //
 // A computed column expression is valid if all of the following are true:
@@ -57,8 +35,12 @@ func MakeComputedColumnValidator(
 //   - It does not reference other computed columns.
 //
 // TODO(mgartner): Add unit tests for Validate.
-func (v *ComputedColumnValidator) Validate(
+func ValidateComputedColumnExpression(
+	ctx context.Context,
+	desc catalog.TableDescriptor,
 	d *tree.ColumnTableDef,
+	tn *tree.TableName,
+	semaCtx *tree.SemaContext,
 ) (serializedExpr string, _ error) {
 	if d.HasDefaultExpr() {
 		return "", pgerror.New(
@@ -69,7 +51,7 @@ func (v *ComputedColumnValidator) Validate(
 
 	var depColIDs catalog.TableColSet
 	// First, check that no column in the expression is a computed column.
-	err := iterColDescriptors(v.desc, d.Computed.Expr, func(c catalog.Column) error {
+	err := iterColDescriptors(desc, d.Computed.Expr, func(c catalog.Column) error {
 		if c.IsComputed() {
 			return pgerror.New(pgcode.InvalidTableDefinition,
 				"computed columns cannot reference other computed columns")
@@ -85,7 +67,7 @@ func (v *ComputedColumnValidator) Validate(
 	// 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.
-	if err := v.desc.ForeachOutboundFK(func(fk *descpb.ForeignKeyConstraint) error {
+	if err := desc.ForeachOutboundFK(func(fk *descpb.ForeignKeyConstraint) error {
 		for _, id := range fk.OriginColumnIDs {
 			if !depColIDs.Contains(id) {
 				// We don't depend on this column.
@@ -110,7 +92,7 @@ func (v *ComputedColumnValidator) Validate(
 	}
 
 	// Resolve the type of the computed column expression.
-	defType, err := tree.ResolveType(v.ctx, d.Type, v.semaCtx.GetTypeResolver())
+	defType, err := tree.ResolveType(ctx, d.Type, semaCtx.GetTypeResolver())
 	if err != nil {
 		return "", err
 	}
@@ -120,14 +102,14 @@ func (v *ComputedColumnValidator) Validate(
 	// functions. In order to safely serialize user defined types and their
 	// members, we need to serialize the typed expression here.
 	expr, _, err := DequalifyAndValidateExpr(
-		v.ctx,
-		v.desc,
+		ctx,
+		desc,
 		d.Computed.Expr,
 		defType,
 		"computed column",
-		v.semaCtx,
+		semaCtx,
 		tree.VolatilityImmutable,
-		v.tableName,
+		tn,
 	)
 	if err != nil {
 		return "", err
@@ -145,7 +127,7 @@ func (v *ComputedColumnValidator) Validate(
 				return
 			}
 			var col catalog.Column
-			if col, err = v.desc.FindColumnWithID(colID); err != nil {
+			if col, err = desc.FindColumnWithID(colID); err != nil {
 				err = errors.WithAssertionFailure(err)
 				return
 			}
@@ -167,12 +149,12 @@ func (v *ComputedColumnValidator) Validate(
 	return expr, nil
 }
 
-// ValidateNoDependents verifies that the input column is not dependent on a
-// computed column. The function errs if any existing computed columns or
-// computed columns being added reference the given column.
-// TODO(mgartner): Add unit tests for ValidateNoDependents.
-func (v *ComputedColumnValidator) ValidateNoDependents(col catalog.Column) error {
-	for _, c := range v.desc.NonDropColumns() {
+// ValidateColumnHasNoDependents verifies that the input column has no dependent
+// computed columns. It returns an error if any existing or ADD mutation
+// computed columns reference the given column.
+// TODO(mgartner): Add unit tests.
+func ValidateColumnHasNoDependents(desc catalog.TableDescriptor, col catalog.Column) error {
+	for _, c := range desc.NonDropColumns() {
 		if !c.IsComputed() {
 			continue
 		}
@@ -183,7 +165,7 @@ func (v *ComputedColumnValidator) ValidateNoDependents(col catalog.Column) error
 			return errors.WithAssertionFailure(err)
 		}
 
-		err = iterColDescriptors(v.desc, expr, func(colVar catalog.Column) error {
+		err = iterColDescriptors(desc, expr, func(colVar catalog.Column) error {
 			if colVar.GetID() == col.GetID() {
 				return pgerror.Newf(
 					pgcode.InvalidColumnReference,

pkg/sql/catalog/schemaexpr/partial_index.go

@@ -21,35 +21,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 )
 
-// IndexPredicateValidator validates that an expression is a valid partial index
-// predicate. See Validate for more details.
-type IndexPredicateValidator struct {
-	ctx       context.Context
-	tableName tree.TableName
-	desc      catalog.TableDescriptor
-	semaCtx   *tree.SemaContext
-}
-
-// MakeIndexPredicateValidator returns an IndexPredicateValidator struct that
-// can be used to validate partial index predicates. See Validate for more
-// details.
-func MakeIndexPredicateValidator(
-	ctx context.Context,
-	tableName tree.TableName,
-	desc catalog.TableDescriptor,
-	semaCtx *tree.SemaContext,
-) IndexPredicateValidator {
-	return IndexPredicateValidator{
-		ctx:       ctx,
-		tableName: tableName,
-		desc:      desc,
-		semaCtx:   semaCtx,
-	}
-}
-
-// Validate verifies that an expression is a valid partial index predicate. If
-// the expression is valid, it returns the serialized expression with the
-// columns dequalified.
+// ValidatePartialIndexPredicate verifies that an expression is a valid partial
+// index predicate. If the expression is valid, it returns the serialized
+// expression with the columns dequalified.
 //
 // A predicate expression is valid if all of the following are true:
 //
@@ -59,16 +33,22 @@ func MakeIndexPredicateValidator(
 //   - It does not include non-immutable, aggregate, window, or set returning
 //     functions.
 //
-func (v *IndexPredicateValidator) Validate(e tree.Expr) (string, error) {
+func ValidatePartialIndexPredicate(
+	ctx context.Context,
+	desc catalog.TableDescriptor,
+	e tree.Expr,
+	tn *tree.TableName,
+	semaCtx *tree.SemaContext,
+) (string, error) {
 	expr, _, err := DequalifyAndValidateExpr(
-		v.ctx,
-		v.desc,
+		ctx,
+		desc,
 		e,
 		types.Bool,
 		"index predicate",
-		v.semaCtx,
+		semaCtx,
 		tree.VolatilityImmutable,
-		&v.tableName,
+		tn,
 	)
 	if err != nil {
 		return "", err

pkg/sql/catalog/schemaexpr/partial_index_test.go

@@ -38,8 +38,6 @@ func TestIndexPredicateValidator_Validate(t *testing.T) {
 		[]testCol{{"c", types.String}},
 	)
 
-	validator := schemaexpr.MakeIndexPredicateValidator(ctx, tn, desc, &semaCtx)
-
 	testData := []struct {
 		expr          string
 		expectedValid bool
@@ -91,7 +89,9 @@ func TestIndexPredicateValidator_Validate(t *testing.T) {
 				t.Fatalf("%s: unexpected error: %s", d.expr, err)
 			}
 
-			deqExpr, err := validator.Validate(expr)
+			deqExpr, err := schemaexpr.ValidatePartialIndexPredicate(
+				ctx, desc, expr, &tn, &semaCtx,
+			)
 
 			if !d.expectedValid {
 				if err == nil {

pkg/sql/catalog/tabledesc/table.go

@@ -104,6 +104,11 @@ func MakeColumnDefDescs(
 	}
 
 	if d.IsComputed() {
+		// Note: We do not validate the computed column expression here because
+		// it may reference columns that have not yet been added to a table
+		// descriptor. Callers must validate the expression with
+		// schemaexpr.ValidateComputedColumnExpression once all possible
+		// reference columns are part of the table descriptor.
 		s := tree.Serialize(d.Computed.Expr)
 		col.ComputeExpr = &s
 	}

pkg/sql/create_index.go

@@ -242,8 +242,9 @@ func MakeIndexDescriptor(
 	}
 
 	if n.Predicate != nil {
-		idxValidator := schemaexpr.MakeIndexPredicateValidator(params.ctx, n.Table, tableDesc, &params.p.semaCtx)
-		expr, err := idxValidator.Validate(n.Predicate)
+		expr, err := schemaexpr.ValidatePartialIndexPredicate(
+			params.ctx, tableDesc, n.Predicate, &n.Table, params.p.SemaCtx(),
+		)
 		if err != nil {
 			return nil, err
 		}

pkg/sql/create_table.go

@@ -1879,7 +1879,6 @@ func NewTableDesc(
 		return newColumns, nil
 	}
 
-	idxValidator := schemaexpr.MakeIndexPredicateValidator(ctx, n.Table, &desc, semaCtx)
 	for _, def := range n.Defs {
 		switch d := def.(type) {
 		case *tree.ColumnTableDef, *tree.LikeTableDef:
@@ -1964,7 +1963,9 @@ func NewTableDesc(
 				}
 			}
 			if d.Predicate != nil {
-				expr, err := idxValidator.Validate(d.Predicate)
+				expr, err := schemaexpr.ValidatePartialIndexPredicate(
+					ctx, &desc, d.Predicate, &n.Table, semaCtx,
+				)
 				if err != nil {
 					return nil, err
 				}
@@ -2061,7 +2062,9 @@ func NewTableDesc(
 				}
 			}
 			if d.Predicate != nil {
-				expr, err := idxValidator.Validate(d.Predicate)
+				expr, err := schemaexpr.ValidatePartialIndexPredicate(
+					ctx, &desc, d.Predicate, &n.Table, semaCtx,
+				)
 				if err != nil {
 					return nil, err
 				}
@@ -2290,17 +2293,13 @@ func NewTableDesc(
 
 	// Now that we have all the other columns set up, we can validate
 	// any computed columns.
-	computedColValidator := schemaexpr.MakeComputedColumnValidator(
-		ctx,
-		&desc,
-		semaCtx,
-		&n.Table,
-	)
 	for _, def := range n.Defs {
 		switch d := def.(type) {
 		case *tree.ColumnTableDef:
 			if d.IsComputed() {
-				serializedExpr, err := computedColValidator.Validate(d)
+				serializedExpr, err := schemaexpr.ValidateComputedColumnExpression(
+					ctx, &desc, d, &n.Table, semaCtx,
+				)
 				if err != nil {
 					return nil, err
 				}

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -1122,12 +1122,10 @@ func (b *Builder) shouldApplyImplicitLockingToUpdateInput(upd *memo.UpdateExpr)
 
 	// Try to match the Update's input expression against the pattern:
 	//
-	//   [Project] [IndexJoin] Scan
+	//   [Project]* [IndexJoin] Scan
 	//
 	input := upd.Input
-	if proj, ok := input.(*memo.ProjectExpr); ok {
-		input = proj.Input
-	}
+	input = unwrapProjectExprs(input)
 	if idxJoin, ok := input.(*memo.IndexJoinExpr); ok {
 		input = idxJoin.Input
 	}
@@ -1138,22 +1136,17 @@ func (b *Builder) shouldApplyImplicitLockingToUpdateInput(upd *memo.UpdateExpr)
 // tryApplyImplicitLockingToUpsertInput determines whether or not the builder
 // should apply a FOR UPDATE row-level locking mode to the initial row scan of
 // an UPSERT statement.
-//
-// TODO(nvanbenschoten): implement this method to match on appropriate Upsert
-// expression trees and apply a row-level locking mode.
 func (b *Builder) shouldApplyImplicitLockingToUpsertInput(ups *memo.UpsertExpr) bool {
 	if !b.evalCtx.SessionData.ImplicitSelectForUpdate {
 		return false
 	}
 
 	// Try to match the Upsert's input expression against the pattern:
 	//
-	//   [Project] (LeftJoin Scan | LookupJoin) [Project] Values
+	//   [Project]* (LeftJoin Scan | LookupJoin) [Project]* Values
 	//
 	input := ups.Input
-	if proj, ok := input.(*memo.ProjectExpr); ok {
-		input = proj.Input
-	}
+	input = unwrapProjectExprs(input)
 	switch join := input.(type) {
 	case *memo.LeftJoinExpr:
 		if _, ok := join.Right.(*memo.ScanExpr); !ok {
@@ -1167,9 +1160,7 @@ func (b *Builder) shouldApplyImplicitLockingToUpsertInput(ups *memo.UpsertExpr)
 	default:
 		return false
 	}
-	if proj, ok := input.(*memo.ProjectExpr); ok {
-		input = proj.Input
-	}
+	input = unwrapProjectExprs(input)
 	_, ok := input.(*memo.ValuesExpr)
 	return ok
 }
@@ -1183,3 +1174,12 @@ func (b *Builder) shouldApplyImplicitLockingToUpsertInput(ups *memo.UpsertExpr)
 func (b *Builder) shouldApplyImplicitLockingToDeleteInput(del *memo.DeleteExpr) bool {
 	return false
 }
+
+// unwrapProjectExprs unwraps zero or more nested ProjectExprs. It returns the
+// first non-ProjectExpr in the chain, or the input if it is not a ProjectExpr.
+func unwrapProjectExprs(input memo.RelExpr) memo.RelExpr {
+	if proj, ok := input.(*memo.ProjectExpr); ok {
+		return unwrapProjectExprs(proj.Input)
+	}
+	return input
+}