sql: Make it possible to have custom function name resolution.

Introduce a `CustomFunctionDefinitionResolver` interface which allows
`tree.SearchPath` implementations to optionally implement custom
function defition resolution logic.

Release Notes: None

pkg/sql/catalog/schemaexpr/expr.go

@@ -345,7 +345,7 @@ func newNameResolver(
 // unresolved names replaced with IndexedVars.
 func (nr *nameResolver) resolveNames(expr tree.Expr) (tree.Expr, error) {
 	var v NameResolutionVisitor
-	return ResolveNamesUsingVisitor(&v, expr, nr.source, *nr.ivarHelper, nr.evalCtx.SessionData().SearchPath)
+	return ResolveNamesUsingVisitor(&v, expr, nr.source, *nr.ivarHelper)
 }
 
 // addColumn adds a new column to the nameResolver so that it can be resolved in

pkg/sql/catalog/schemaexpr/select_name_resolution.go

@@ -18,15 +18,13 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
-	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 )
 
 // NameResolutionVisitor is a tree.Visitor implementation used to
 // resolve the column names in an expression.
 type NameResolutionVisitor struct {
 	err        error
 	iVarHelper tree.IndexedVarHelper
-	searchPath sessiondata.SearchPath
 	resolver   colinfo.ColumnResolver
 }
 
@@ -98,11 +96,9 @@ func ResolveNamesUsingVisitor(
 	expr tree.Expr,
 	source *colinfo.DataSourceInfo,
 	ivarHelper tree.IndexedVarHelper,
-	searchPath sessiondata.SearchPath,
 ) (tree.Expr, error) {
 	*v = NameResolutionVisitor{
 		iVarHelper: ivarHelper,
-		searchPath: searchPath,
 		resolver: colinfo.ColumnResolver{
 			Source: source,
 		},
@@ -119,13 +115,10 @@ func (v *NameResolutionVisitor) Err() error {
 
 // MakeNameResolutionVisitor returns initialized name resolution visitor.
 func MakeNameResolutionVisitor(
-	source *colinfo.DataSourceInfo,
-	ivarHelper tree.IndexedVarHelper,
-	searchPath sessiondata.SearchPath,
+	source *colinfo.DataSourceInfo, ivarHelper tree.IndexedVarHelper,
 ) NameResolutionVisitor {
 	return NameResolutionVisitor{
 		iVarHelper: ivarHelper,
-		searchPath: searchPath,
 		resolver: colinfo.ColumnResolver{
 			Source: source,
 		},

pkg/sql/select_name_resolution.go

@@ -27,5 +27,5 @@ func (p *planner) resolveNames(
 	if expr == nil {
 		return nil, nil
 	}
-	return schemaexpr.ResolveNamesUsingVisitor(&p.nameResolutionVisitor, expr, source, ivarHelper, p.SessionData().SearchPath)
+	return schemaexpr.ResolveNamesUsingVisitor(&p.nameResolutionVisitor, expr, source, ivarHelper)
 }

pkg/sql/sem/tree/function_name_test.go

@@ -25,18 +25,25 @@ import (
 func TestResolveFunction(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
+
+	searchPath := func() tree.SearchPath {
+		sp := sessiondata.MakeSearchPath([]string{"pg_catalog"})
+		return &sp
+	}()
+
 	testCases := []struct {
-		in, out string
-		err     string
+		in, out          string
+		err              string
+		customSearchPath tree.SearchPath
 	}{
-		{`count`, `count`, ``},
-		{`pg_catalog.pg_typeof`, `pg_typeof`, ``},
+		{in: `count`, out: `count`},
+		{in: `pg_catalog.pg_typeof`, out: `pg_typeof`},
 
-		{`foo`, ``, `unknown function: foo`},
-		{`""`, ``, `invalid function name: ""`},
+		{in: `foo`, err: `unknown function: foo`},
+		{in: `foo`, out: `count`, customSearchPath: &customResolver{SearchPath: searchPath}},
+		{in: `""`, err: `invalid function name: ""`},
 	}
 
-	searchPath := sessiondata.MakeSearchPath([]string{"pg_catalog"})
 	for _, tc := range testCases {
 		stmt, err := parser.ParseOne("SELECT " + tc.in + "(1)")
 		if err != nil {
@@ -47,7 +54,11 @@ func TestResolveFunction(t *testing.T) {
 			t.Fatalf("%s does not parse to a tree.FuncExpr", tc.in)
 		}
 		q := f.Func
-		_, err = q.Resolve(&searchPath)
+		sp := searchPath
+		if tc.customSearchPath != nil {
+			sp = tc.customSearchPath
+		}
+		_, err = q.Resolve(sp)
 		if tc.err != "" {
 			if !testutils.IsError(err, tc.err) {
 				t.Fatalf("%s: expected %s, but found %v", tc.in, tc.err, err)
@@ -62,3 +73,21 @@ func TestResolveFunction(t *testing.T) {
 		}
 	}
 }
+
+type customResolver struct {
+	tree.SearchPath
+}
+
+var _ tree.CustomFunctionDefinitionResolver = (*customResolver)(nil)
+
+// Resolve implements tree.CustomFunctionDefinitionResolver
+func (r customResolver) Resolve(name string) *tree.FunctionDefinition {
+	if name == "foo" {
+		name = "count"
+	}
+	fn, found := tree.FunDefs[name]
+	if found {
+		return fn
+	}
+	return nil
+}

pkg/sql/sem/tree/name_resolution.go

@@ -139,6 +139,14 @@ type SearchPath interface {
 	IterateSearchPath(func(schema string) error) error
 }
 
+// CustomFunctionDefinitionResolver is an interface providing custom
+// function definition resolution functionality.
+type CustomFunctionDefinitionResolver interface {
+	// Resolve resolves function with specified name, and returns
+	// non-nil function definition if resolved successfully.
+	Resolve(name string) *FunctionDefinition
+}
+
 // EmptySearchPath is a SearchPath with no members.
 var EmptySearchPath SearchPath = emptySearchPath{}
 
@@ -148,6 +156,21 @@ func (emptySearchPath) IterateSearchPath(func(string) error) error {
 	return nil
 }
 
+// getFunctionDefinitionResolver returns a function which
+// is used to resolve function definition.
+func getFunctionDefinitionResolver(sp SearchPath) func(name string) *FunctionDefinition {
+	if custom, ok := sp.(CustomFunctionDefinitionResolver); ok {
+		return custom.Resolve
+	}
+	return func(name string) *FunctionDefinition {
+		fn, ok := FunDefs[name]
+		if ok {
+			return fn
+		}
+		return nil
+	}
+}
+
 // ResolveFunction transforms an UnresolvedName to a FunctionDefinition.
 //
 // Function resolution currently takes a "short path" using the
@@ -170,12 +193,14 @@ func (n *UnresolvedName) ResolveFunction(searchPath SearchPath) (*FunctionDefini
 			"invalid function name: %s", n)
 	}
 
+	resolveFn := getFunctionDefinitionResolver(searchPath)
+
 	// We ignore the catalog part. Like explained above, we currently
 	// only support functions in virtual schemas, which always exist
 	// independently of the database/catalog prefix.
 	function, prefix := n.Parts[0], n.Parts[1]
 
-	if d, ok := FunDefs[function]; ok && prefix == "" {
+	if d := resolveFn(function); d != nil && prefix == "" {
 		// Fast path: return early.
 		return d, nil
 	}
@@ -191,35 +216,33 @@ func (n *UnresolvedName) ResolveFunction(searchPath SearchPath) (*FunctionDefini
 		// If the user specified public, it may be from a PostgreSQL extension.
 		// Double check the function definition allows resolution on the public
 		// schema, and resolve as such if appropriate.
-		if d, ok := FunDefs[function]; ok && d.AvailableOnPublicSchema {
+		if d := resolveFn(function); d != nil && d.AvailableOnPublicSchema {
 			return d, nil
 		}
 	}
 
 	if prefix != "" {
 		fullName = prefix + "." + function
 	}
-	def, ok := FunDefs[fullName]
-	if !ok {
-		found := false
+	def := resolveFn(fullName)
+	if def == nil {
 		if prefix == "" {
 			// The function wasn't qualified, so we must search for it via
 			// the search path first.
 			if err := searchPath.IterateSearchPath(func(alt string) error {
 				fullName = alt + "." + function
-				if def, ok = FunDefs[fullName]; ok {
-					found = true
+				if def = resolveFn(fullName); def != nil {
 					return iterutil.StopIteration()
 				}
 				return nil
 			}); err != nil {
 				return nil, err
 			}
 		}
-		if !found {
+		if def == nil {
 			extraMsg := ""
 			// Try a little harder.
-			if rdef, ok := FunDefs[strings.ToLower(function)]; ok {
+			if rdef := resolveFn(strings.ToLower(function)); rdef != nil {
 				extraMsg = fmt.Sprintf(", but %s() exists", rdef.Name)
 			}
 			return nil, pgerror.Newf(