From 09fd47f241859a716374f5b7926f758a5ca8fd3e Mon Sep 17 00:00:00 2001
From: "David G. Johnston" <David.G.Johnston@Gmail.com>
Date: Sat, 22 Mar 2025 09:45:49 -0700
Subject: [PATCH] Handler function lookups ignore non-handler functions

---
 src/backend/catalog/namespace.c   | 35 +++++++++++++++++++----
 src/backend/parser/parse_clause.c | 10 +------
 src/backend/parser/parse_func.c   | 47 +++++++++++++++++++++++--------
 src/include/catalog/namespace.h   |  7 +++++
 src/include/parser/parse_func.h   |  5 ++++
 5 files changed, 78 insertions(+), 26 deletions(-)

diff --git a/src/backend/catalog/namespace.c b/src/backend/catalog/namespace.c
index d97d632a7e..0cb66261a3 100644
--- a/src/backend/catalog/namespace.c
+++ b/src/backend/catalog/namespace.c
@@ -1115,9 +1115,20 @@ TypeIsVisibleExt(Oid typid, bool *is_missing)
 	return visible;
 }
 
-
+FuncCandidateList
+FuncnameGetCandidates(List *names, int nargs, List *argnames,
+					  bool expand_variadic, bool expand_defaults,
+					  bool include_out_arguments,
+					  bool missing_ok)
+{
+	return HandlerGetCandidates(names, nargs, argnames,
+								 expand_variadic, expand_defaults,
+								 include_out_arguments,
+								 NULL,
+								 missing_ok);
+}
 /*
- * FuncnameGetCandidates
+ * HandlerGetCandidates
  *		Given a possibly-qualified function name and argument count,
  *		retrieve a list of the possible matches.
  *
@@ -1184,14 +1195,20 @@ TypeIsVisibleExt(Oid typid, bool *is_missing)
  * The caller might end up discarding such an entry anyway, but if it selects
  * such an entry it should react as though the call were ambiguous.
  *
+ * The presence of a handlerkey further restricts the search to functions whose
+ * return data type matches the handlerkey.  These are pseudo-types known to
+ * the system are are never schema-qualified.
+ *
  * If missing_ok is true, an empty list (NULL) is returned if the name was
  * schema-qualified with a schema that does not exist.  Likewise if no
  * candidate is found for other reasons.
  */
 FuncCandidateList
-FuncnameGetCandidates(List *names, int nargs, List *argnames,
-					  bool expand_variadic, bool expand_defaults,
-					  bool include_out_arguments, bool missing_ok)
+HandlerGetCandidates(List *names, int nargs, List *argnames,
+					 bool expand_variadic, bool expand_defaults,
+					 bool include_out_arguments,
+					 Oid *handlerkey,
+					 bool missing_ok)
 {
 	FuncCandidateList resultList = NULL;
 	bool		any_special = false;
@@ -1374,6 +1391,14 @@ FuncnameGetCandidates(List *names, int nargs, List *argnames,
 				continue;
 		}
 
+		if (handlerkey)
+		{
+			{
+				if (procform->prorettype != handlerkey)
+					continue;
+			}
+		}
+
 		/*
 		 * We must compute the effective argument list so that we can easily
 		 * compare it to earlier results.  We waste a palloc cycle if it gets
diff --git a/src/backend/parser/parse_clause.c b/src/backend/parser/parse_clause.c
index 2e64fcae7b..dbea8c2fda 100644
--- a/src/backend/parser/parse_clause.c
+++ b/src/backend/parser/parse_clause.c
@@ -924,7 +924,7 @@ transformRangeTableSample(ParseState *pstate, RangeTableSample *rts)
 	 */
 	funcargtypes[0] = INTERNALOID;
 
-	handlerOid = LookupFuncName(rts->method, 1, funcargtypes, true);
+	handlerOid = LookupHandlerName(rts->method, 1, funcargtypes, (Oid *) TSM_HANDLEROID, true);
 
 	/* we want error to complain about no-such-method, not no-such-function */
 	if (!OidIsValid(handlerOid))
@@ -934,14 +934,6 @@ transformRangeTableSample(ParseState *pstate, RangeTableSample *rts)
 						NameListToString(rts->method)),
 				 parser_errposition(pstate, rts->location)));
 
-	/* check that handler has correct return type */
-	if (get_func_rettype(handlerOid) != TSM_HANDLEROID)
-		ereport(ERROR,
-				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
-				 errmsg("function %s must return type %s",
-						NameListToString(rts->method), "tsm_handler"),
-				 parser_errposition(pstate, rts->location)));
-
 	/* OK, run the handler to get TsmRoutine, for argument type info */
 	tsm = GetTsmRoutine(handlerOid);
 
diff --git a/src/backend/parser/parse_func.c b/src/backend/parser/parse_func.c
index 583bbbf232..42cff0f14c 100644
--- a/src/backend/parser/parse_func.c
+++ b/src/backend/parser/parse_func.c
@@ -50,7 +50,9 @@ static Node *ParseComplexProjection(ParseState *pstate, const char *funcname,
 									Node *first_arg, int location);
 static Oid	LookupFuncNameInternal(ObjectType objtype, List *funcname,
 								   int nargs, const Oid *argtypes,
-								   bool include_out_arguments, bool missing_ok,
+								   bool include_out_arguments,
+								   Oid *handlerkey,
+								   bool missing_ok,
 								   FuncLookupError *lookupError);
 
 
@@ -2048,7 +2050,9 @@ func_signature_string(List *funcname, int nargs,
 static Oid
 LookupFuncNameInternal(ObjectType objtype, List *funcname,
 					   int nargs, const Oid *argtypes,
-					   bool include_out_arguments, bool missing_ok,
+					   bool include_out_arguments,
+					   Oid *handlerkey,
+					   bool missing_ok,
 					   FuncLookupError *lookupError)
 {
 	Oid			result = InvalidOid;
@@ -2061,8 +2065,8 @@ LookupFuncNameInternal(ObjectType objtype, List *funcname,
 	*lookupError = FUNCLOOKUP_NOSUCHFUNC;
 
 	/* Get list of candidate objects */
-	clist = FuncnameGetCandidates(funcname, nargs, NIL, false, false,
-								  include_out_arguments, missing_ok);
+	clist = HandlerGetCandidates(funcname, nargs, NIL, false, false,
+								  include_out_arguments, handlerkey, missing_ok);
 
 	/* Scan list for a match to the arg types (if specified) and the objtype */
 	for (; clist != NULL; clist = clist->next)
@@ -2118,8 +2122,13 @@ LookupFuncNameInternal(ObjectType objtype, List *funcname,
 	return result;
 }
 
+Oid
+LookupFuncName(List *funcname, int nargs, const Oid *argtypes, bool missing_ok)
+{
+	return LookupHandlerName(funcname, nargs, argtypes, NULL, missing_ok);
+}
 /*
- * LookupFuncName
+ * LookupHandlerName
  *
  * Given a possibly-qualified function name and optionally a set of argument
  * types, look up the function.  Pass nargs == -1 to indicate that the number
@@ -2139,16 +2148,20 @@ LookupFuncNameInternal(ObjectType objtype, List *funcname,
  * Only functions will be found; procedures will be ignored even if they
  * match the name and argument types.  (However, we don't trouble to reject
  * aggregates or window functions here.)
+ *
+ * The presence of a handlerkey further restricts the search to functions whose
+ * return data type matches the handlerkey.  These are pseudo-types known to
+ * the system are are never schema-qualified.
  */
 Oid
-LookupFuncName(List *funcname, int nargs, const Oid *argtypes, bool missing_ok)
+LookupHandlerName(List *funcname, int nargs, const Oid *argtypes, Oid *handlerkey, bool missing_ok)
 {
 	Oid			funcoid;
 	FuncLookupError lookupError;
 
 	funcoid = LookupFuncNameInternal(OBJECT_FUNCTION,
 									 funcname, nargs, argtypes,
-									 false, missing_ok,
+									 false, handlerkey, missing_ok,
 									 &lookupError);
 
 	if (OidIsValid(funcoid))
@@ -2187,10 +2200,16 @@ LookupFuncName(List *funcname, int nargs, const Oid *argtypes, bool missing_ok)
 	return InvalidOid;			/* Keep compiler quiet */
 }
 
+Oid
+LookupFuncWithArgs(ObjectType objtype, ObjectWithArgs *func, bool missing_ok)
+{
+	return LookupHandlerWithArgs(objtype, func, NULL, missing_ok);
+}
+
 /*
- * LookupFuncWithArgs
+ * LookupHandlerWithArgs
  *
- * Like LookupFuncName, but the argument types are specified by an
+ * Like LookupHandlerName, but the argument types are specified by an
  * ObjectWithArgs node.  Also, this function can check whether the result is a
  * function, procedure, or aggregate, based on the objtype argument.  Pass
  * OBJECT_ROUTINE to accept any of them.
@@ -2201,9 +2220,13 @@ LookupFuncName(List *funcname, int nargs, const Oid *argtypes, bool missing_ok)
  * When missing_ok is true we don't generate any error for missing objects and
  * return InvalidOid.  Other types of errors can still be raised, regardless
  * of the value of missing_ok.
+ *
+ * The presence of a handlerkey further restricts the search to functions whose
+ * return data type matches the handlerkey.  These are pseudo-types known to
+ * the system are are never schema-qualified.
  */
 Oid
-LookupFuncWithArgs(ObjectType objtype, ObjectWithArgs *func, bool missing_ok)
+LookupHandlerWithArgs(ObjectType objtype, ObjectWithArgs *func, Oid *handlerkey, bool missing_ok)
 {
 	Oid			argoids[FUNC_MAX_ARGS];
 	int			argcount;
@@ -2269,7 +2292,7 @@ LookupFuncWithArgs(ObjectType objtype, ObjectWithArgs *func, bool missing_ok)
 	 */
 	oid = LookupFuncNameInternal(func->args_unspecified ? objtype : OBJECT_ROUTINE,
 								 func->objname, nargs, argoids,
-								 false, missing_ok,
+								 false, handlerkey, missing_ok,
 								 &lookupError);
 
 	/*
@@ -2315,7 +2338,7 @@ LookupFuncWithArgs(ObjectType objtype, ObjectWithArgs *func, bool missing_ok)
 			/* For objtype == OBJECT_PROCEDURE, we can ignore non-procedures */
 			poid = LookupFuncNameInternal(objtype, func->objname,
 										  argcount, argoids,
-										  true, missing_ok,
+										  true, handlerkey, missing_ok,
 										  &lookupError);
 
 			/* Combine results, handling ambiguity */
diff --git a/src/include/catalog/namespace.h b/src/include/catalog/namespace.h
index 8c7ccc69a3..98a703b742 100644
--- a/src/include/catalog/namespace.h
+++ b/src/include/catalog/namespace.h
@@ -103,6 +103,13 @@ extern FuncCandidateList FuncnameGetCandidates(List *names,
 											   bool expand_defaults,
 											   bool include_out_arguments,
 											   bool missing_ok);
+extern FuncCandidateList HandlerGetCandidates(List *names,
+											  int nargs, List *argnames,
+											  bool expand_variadic,
+											  bool expand_defaults,
+											  bool include_out_arguments,
+											  Oid *handlerkey,
+											  bool missing_ok);
 extern bool FunctionIsVisible(Oid funcid);
 
 extern Oid	OpernameGetOprid(List *names, Oid oprleft, Oid oprright);
diff --git a/src/include/parser/parse_func.h b/src/include/parser/parse_func.h
index a6f24b83d8..47f76cb4eb 100644
--- a/src/include/parser/parse_func.h
+++ b/src/include/parser/parse_func.h
@@ -68,6 +68,11 @@ extern Oid	LookupFuncName(List *funcname, int nargs, const Oid *argtypes,
 extern Oid	LookupFuncWithArgs(ObjectType objtype, ObjectWithArgs *func,
 							   bool missing_ok);
 
+extern Oid	LookupHandlerName(List *funcname, int nargs, const Oid *argtypes,
+								Oid *handlerkey, bool missing_ok);
+extern Oid	LookupHandlerWithArgs(ObjectType objtype, ObjectWithArgs *func,
+									Oid *handlerkey, bool missing_ok);
+
 extern void check_srf_call_placement(ParseState *pstate, Node *last_srf,
 									 int location);
 
-- 
2.34.1

