schemachanger: track partial index dependencies on functions

This commit adds the plumbing necessary to find referenced UDF IDs in a
partial index predicate, and updates the declarative schema changer to
add back-references.

Release note (sql change): Partial indexes can now reference
user-defined function.

Author: rafiss

pkg/sql/catalog/descriptor.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/catalog/funcdesc/func_desc.go

@@ -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/catalog/tabledesc/structured.go

@@ -458,12 +458,20 @@ func (desc *wrapper) GetAllReferencedFunctionIDs() (catalog.DescriptorIDSet, err
 	for i := range desc.Triggers {
 		ret = ret.Union(catalog.MakeDescriptorIDSet(desc.Triggers[i].DependsOnRoutines...))
 	}
-	// Add deps from policies
+	// Add deps from policies.
 	for i := range desc.Policies {
 		ret = ret.Union(catalog.MakeDescriptorIDSet(desc.Policies[i].DependsOnFunctions...))
 	}
-	// TODO(chengxiong): add logic to extract references from indexes when UDFs
-	// are allowed in them.
+	// Add deps from partial indexes.
+	for _, idx := range desc.AllIndexes() {
+		if idx.IsPartial() {
+			ids, err := schemaexpr.GetUDFIDsFromExprStr(idx.GetPredicate())
+			if err != nil {
+				return catalog.DescriptorIDSet{}, err
+			}
+			ret = ret.Union(ids)
+		}
+	}
 	return ret.Union(catalog.MakeDescriptorIDSet(desc.DependsOnFunctions...)), nil
 }
 
@@ -496,6 +504,26 @@ func (desc *wrapper) GetAllReferencedFunctionIDsInTrigger(
 	return fnIDs
 }
 
+// GetAllReferencedFunctionIDsInIndex implements the TableDescriptor interface.
+func (desc *wrapper) GetAllReferencedFunctionIDsInIndex(
+	indexID descpb.IndexID,
+) (fnIDs catalog.DescriptorIDSet, err error) {
+	idx := catalog.FindIndexByID(desc, indexID)
+	if idx == nil {
+		return catalog.DescriptorIDSet{}, nil
+	}
+
+	var ret catalog.DescriptorIDSet
+	if idx.IsPartial() {
+		ids, err := schemaexpr.GetUDFIDsFromExprStr(idx.GetPredicate())
+		if err != nil {
+			return catalog.DescriptorIDSet{}, err
+		}
+		ret = ret.Union(ids)
+	}
+	return ret, nil
+}
+
 // GetAllReferencedFunctionIDsInPolicy implements the TableDescriptor interface.
 func (desc *wrapper) GetAllReferencedFunctionIDsInPolicy(
 	policyID descpb.PolicyID,