Merge #147817

147817: sql/catalog,schemachanger: allow UDFs in ON UPDATE and computed expressions r=rafiss a=rafiss

We already had the dependency tracking in place, but there was still a
block in place since we don't yet support UDFs in backfills. There were
just a few places relating to ON UPDATE expressions that needed to be
updated to emit the mutation event that modifies dependency back
references.

This patch removes the block within the schemachanger, so that now the
only thing that gets blocked is actually attempting to use a UDF in the
context of a backfill.

informs #87699
fixes #115346
Release note (sql change): Computed column expressions and ON UPDATE
expressions can now reference user-defined functions.

Co-authored-by: Rafi Shamim <[email protected]>

Author: craig[bot]

pkg/ccl/logictestccl/tests/3node-tenant/generated_test.go

@@ -156,12 +156,16 @@ func (p *planner) addColumnImpl(
 	}
 
 	// We need to allocate new ID for the created column in order to correctly
-	// assign sequence ownership.
+	// assign sequence ownership and function dependencies.
 	version := params.ExecCfg().Settings.Version.ActiveVersion(params.ctx)
 	if err := n.tableDesc.AllocateIDs(params.ctx, version); err != nil {
 		return err
 	}
 
+	if err := params.p.maybeUpdateFunctionReferencesForColumn(params.ctx, n.tableDesc, col); err != nil {
+		return err
+	}
+
 	// If the new column has a DEFAULT or an ON UPDATE expression that uses a
 	// sequence, add references between its descriptor and this column descriptor.
 	if err := cdd.ForEachTypedExpr(func(expr tree.TypedExpr, colExprKind tabledesc.ColExprKind) error {

pkg/ccl/logictestccl/tests/local-read-committed/generated_test.go

@@ -20,6 +20,7 @@ go_library(
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/catenumpb",
         "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/fetchpb",
         "//pkg/sql/catalog/schemaexpr",
         "//pkg/sql/catalog/tabledesc",

pkg/ccl/logictestccl/tests/local-repeatable-read/generated_test.go

@@ -18,11 +18,11 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/fetchpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemaexpr"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
-	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/row"
@@ -230,7 +230,7 @@ func (cb *ColumnBackfiller) InitForDistributedUse(
 	var defaultExprs, computedExprs []tree.TypedExpr
 	// Install type metadata in the target descriptors, as well as resolve any
 	// user defined types in the column expressions.
-	if err := flowCtx.Cfg.DB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+	if err := flowCtx.Cfg.DB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
 		resolver := flowCtx.NewTypeResolver(txn.KV())
 		// Hydrate all the types present in the table.
 		if err := typedesc.HydrateTypesInDescriptor(ctx, desc, &resolver); err != nil {
@@ -239,6 +239,7 @@ func (cb *ColumnBackfiller) InitForDistributedUse(
 		// Set up a SemaContext to type check the default and computed expressions.
 		semaCtx := tree.MakeSemaContext(&resolver)
 		semaCtx.UnsupportedTypeChecker = eval.NewUnsupportedTypeChecker(evalCtx.Settings.Version)
+		semaCtx.FunctionResolver = descs.NewDistSQLFunctionResolver(txn.Descriptors(), txn.KV())
 		var err error
 		defaultExprs, err = schemaexpr.MakeDefaultExprs(
 			ctx, cb.added, &transform.ExprTransformContext{}, evalCtx, &semaCtx,
@@ -758,7 +759,7 @@ func (ib *IndexBackfiller) InitForDistributedUse(
 
 	// Install type metadata in the target descriptors, as well as resolve any
 	// user defined types in partial index predicate expressions.
-	if err := flowCtx.Cfg.DB.Txn(ctx, func(ctx context.Context, txn isql.Txn) (err error) {
+	if err := flowCtx.Cfg.DB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) (err error) {
 		resolver := flowCtx.NewTypeResolver(txn.KV())
 		// Hydrate all the types present in the table.
 		if err = typedesc.HydrateTypesInDescriptor(ctx, desc, &resolver); err != nil {
@@ -767,6 +768,7 @@ func (ib *IndexBackfiller) InitForDistributedUse(
 		// Set up a SemaContext to type check the default and computed expressions.
 		semaCtx := tree.MakeSemaContext(&resolver)
 		semaCtx.UnsupportedTypeChecker = eval.NewUnsupportedTypeChecker(evalCtx.Settings.Version)
+		semaCtx.FunctionResolver = descs.NewDistSQLFunctionResolver(txn.Descriptors(), txn.KV())
 		// Convert any partial index predicate strings into expressions.
 		predicates, colExprs, referencedColumns, err = constructExprs(
 			ctx, desc, ib.added, ib.cols, ib.addedCols, ib.computedCols, evalCtx, &semaCtx,

pkg/sql/add_column.go

@@ -680,6 +680,12 @@ type TableDescriptor interface {
 		colID descpb.ColumnID,
 	) (DescriptorIDSet, error)
 
+	// GetAllReferencedFunctionIDsInIndex returns descriptor IDs of all user
+	// defined functions referenced in expressions used by this index.
+	GetAllReferencedFunctionIDsInIndex(
+		indexID descpb.IndexID,
+	) (DescriptorIDSet, error)
+
 	// ForeachDependedOnBy runs a function on all indexes, including those being
 	// added in the mutations.
 	ForeachDependedOnBy(f func(dep *descpb.TableDescriptor_Reference) error) error

pkg/sql/backfill/BUILD.bazel

@@ -385,8 +385,18 @@ func (desc *immutable) validateInboundTableRef(
 			return errors.AssertionFailedf("depended-on-by relation %q (%d) does not have an index with ID %d",
 				backRefTbl.GetName(), by.ID, idxID)
 		}
-		// TODO(chengxiong): add logic to validate reference in index expressions
-		// when UDF usage is allowed in indexes.
+		fnIDs, err := backRefTbl.GetAllReferencedFunctionIDsInIndex(idxID)
+		if err != nil {
+			return err
+		}
+		if fnIDs.Contains(desc.GetID()) {
+			foundInTable = true
+			continue
+		}
+		return errors.AssertionFailedf(
+			"index %d in depended-on-by relation %q (%d) does not have reference to function %q (%d)",
+			idxID, backRefTbl.GetName(), backRefTbl.GetID(), desc.GetName(), desc.GetID(),
+		)
 	}
 
 	for _, cstID := range by.ConstraintIDs {
@@ -852,6 +862,55 @@ func (desc *Mutable) RemovePolicyReference(id descpb.ID, policyID descpb.PolicyI
 	}
 }
 
+// AddIndexReference adds back reference to an index to the function.
+func (desc *Mutable) AddIndexReference(id descpb.ID, indexID descpb.IndexID) error {
+	for _, dep := range desc.DependsOn {
+		if dep == id {
+			return pgerror.Newf(pgcode.InvalidFunctionDefinition,
+				"cannot add dependency from descriptor %d to function %s (%d) because there will be a dependency cycle", id, desc.GetName(), desc.GetID(),
+			)
+		}
+	}
+	for i := range desc.DependedOnBy {
+		if desc.DependedOnBy[i].ID == id {
+			for _, prevID := range desc.DependedOnBy[i].IndexIDs {
+				if prevID == indexID {
+					return nil
+				}
+			}
+			desc.DependedOnBy[i].IndexIDs = append(desc.DependedOnBy[i].IndexIDs, indexID)
+			return nil
+		}
+	}
+	desc.DependedOnBy = append(
+		desc.DependedOnBy,
+		descpb.FunctionDescriptor_Reference{
+			ID:       id,
+			IndexIDs: []descpb.IndexID{indexID},
+		},
+	)
+	sort.Slice(desc.DependedOnBy, func(i, j int) bool {
+		return desc.DependedOnBy[i].ID < desc.DependedOnBy[j].ID
+	})
+	return nil
+}
+
+// RemoveIndexReference removes back reference to an index from the function.
+func (desc *Mutable) RemoveIndexReference(id descpb.ID, indexID descpb.IndexID) {
+	for i := range desc.DependedOnBy {
+		if desc.DependedOnBy[i].ID == id {
+			dep := &desc.DependedOnBy[i]
+			for j := range dep.IndexIDs {
+				if dep.IndexIDs[j] == indexID {
+					dep.IndexIDs = append(dep.IndexIDs[:j], dep.IndexIDs[j+1:]...)
+					desc.maybeRemoveTableReference(id)
+					return
+				}
+			}
+		}
+	}
+}
+
 // maybeRemoveTableReference removes a table's references from the function if
 // the column, index and constraint references are all empty. This function is
 // only used internally when removing an individual column, index or constraint

pkg/sql/backfill/backfill.go

@@ -15,12 +15,22 @@ var schemaExprContextAllowingUDF = map[tree.SchemaExprContext]clusterversion.Key
 	tree.CheckConstraintExpr:           clusterversion.MinSupported,
 	tree.ColumnDefaultExprInNewTable:   clusterversion.MinSupported,
 	tree.ColumnDefaultExprInSetDefault: clusterversion.MinSupported,
+
+	tree.PolicyUsingExpr:     clusterversion.V25_2,
+	tree.PolicyWithCheckExpr: clusterversion.V25_2,
+
+	tree.ColumnDefaultExprInAddColumn: clusterversion.V25_3,
+	tree.ColumnDefaultExprInNewView:   clusterversion.V25_3,
+	tree.ColumnOnUpdateExpr:           clusterversion.V25_3,
+	tree.ExpressionIndexElementExpr:   clusterversion.V25_3,
+	tree.IndexPredicateExpr:           clusterversion.V25_3,
+	tree.StoredComputedColumnExpr:     clusterversion.V25_3,
+	tree.VirtualComputedColumnExpr:    clusterversion.V25_3,
 }
 
 // MaybeFailOnUDFUsage returns an error if the given expression or any
 // sub-expression used a UDF unless it's explicitly listed as an allowed use
 // case.
-// TODO(chengxiong): remove this function when we start allowing UDF references.
 func MaybeFailOnUDFUsage(
 	expr tree.TypedExpr, exprContext tree.SchemaExprContext, version clusterversion.ClusterVersion,
 ) error {

pkg/sql/catalog/descriptor.go

@@ -26,7 +26,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
-	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/errors"
 )
 
@@ -80,13 +79,8 @@ func DequalifyAndValidateExprImpl(
 		return "", nil, colIDs, err
 	}
 
-	// TODO(87699): We can remove this once we have forward/backward references to
-	// functions. We skip this for policies because we have those references in
-	// place already.
-	if context != tree.PolicyUsingExpr && context != tree.PolicyWithCheckExpr {
-		if err := funcdesc.MaybeFailOnUDFUsage(typedExpr, context, version); err != nil {
-			return "", nil, colIDs, unimplemented.NewWithIssue(87699, "usage of user-defined function from relations not supported")
-		}
+	if err := funcdesc.MaybeFailOnUDFUsage(typedExpr, context, version); err != nil {
+		return "", nil, colIDs, err
 	}
 
 	// We need to do the rewrite here before the expression is serialized because