From df9483dd22413a053584377bf4ec5f14a3ae28ab Mon Sep 17 00:00:00 2001
From: Florents Tselai <florents.tselai@gmail.com>
Date: Tue, 17 Sep 2024 01:17:59 +0300
Subject: [PATCH v2 3/3] Fix logic for jpiReplaceFunc. Also make it acceptable
 in jspGetLeftArg and jspGetRightArg.

---
 src/backend/utils/adt/jsonpath.c      |  6 +-
 src/backend/utils/adt/jsonpath_exec.c | 83 ++++++++++++---------------
 2 files changed, 41 insertions(+), 48 deletions(-)

diff --git a/src/backend/utils/adt/jsonpath.c b/src/backend/utils/adt/jsonpath.c
index 3daf2e2d81..3a5c7a00cf 100644
--- a/src/backend/utils/adt/jsonpath.c
+++ b/src/backend/utils/adt/jsonpath.c
@@ -1199,7 +1199,8 @@ jspGetLeftArg(JsonPathItem *v, JsonPathItem *a)
 		   v->type == jpiDiv ||
 		   v->type == jpiMod ||
 		   v->type == jpiStartsWith ||
-		   v->type == jpiDecimal);
+		   v->type == jpiDecimal ||
+		   v->type == jpiReplaceFunc);
 
 	jspInitByBuffer(a, v->base, v->content.args.left);
 }
@@ -1221,7 +1222,8 @@ jspGetRightArg(JsonPathItem *v, JsonPathItem *a)
 		   v->type == jpiDiv ||
 		   v->type == jpiMod ||
 		   v->type == jpiStartsWith ||
-		   v->type == jpiDecimal);
+		   v->type == jpiDecimal ||
+		   v->type == jpiReplaceFunc);
 
 	jspInitByBuffer(a, v->base, v->content.args.right);
 }
diff --git a/src/backend/utils/adt/jsonpath_exec.c b/src/backend/utils/adt/jsonpath_exec.c
index 3d23e44d34..4f0fa3d6b7 100644
--- a/src/backend/utils/adt/jsonpath_exec.c
+++ b/src/backend/utils/adt/jsonpath_exec.c
@@ -1659,74 +1659,65 @@ executeItemOptUnwrapTarget(JsonPathExecContext *cxt, JsonPathItem *jsp,
 				res = executeNextItem(cxt, jsp, NULL, jb, found, true);
 			}
 			break;
+
 		case jpiReplaceFunc:
 			{
 				JsonbValue	jbv;
-				Datum		replace_res;
-				char		*tmp = NULL;
-
-				/*
-				 * Value is not necessarily null-terminated, so we do
-				 * pnstrdup() here.
-				 */
-				tmp = pnstrdup(jb->val.string.val,
-							   jb->val.string.len);
+				char		*replacedTxt;
+				char		*txt = NULL;
 
 				if (unwrap && JsonbType(jb) == jbvArray)
-					return executeItemUnwrapTargetArray(cxt, jsp, jb, found,
-														false);
+					return executeItemUnwrapTargetArray(cxt, jsp, jb, found,false);
 
-				/* TODO: probably need ERRCODE for that? */
-				if (!(jb = getScalar(jb, jbvString)))
-				RETURN_ERROR(ereport(ERROR,
-									 (errcode(ERRCODE_INVALID_ARGUMENT_FOR_SQL_JSON_DATETIME_FUNCTION),
-									  errmsg("jsonpath item method .%s() can only be applied to a string",
-											 jspOperationName(jsp->type)))));
+				if (jb->type == jbvString) {
+					/* Value is not necessarily null-terminated, so we do pnstrdup() here. */
+					txt = pnstrdup(jb->val.string.val, jb->val.string.len);
 
-				if (jsp->content.args.left)
+					res = jperOk;
+				}
+
+				if (res == jperNotFound) {
+					/* TODO: probably need ERRCODE for that? */
+					RETURN_ERROR(ereport(ERROR,
+										 (errcode(ERRCODE_INVALID_ARGUMENT_FOR_SQL_JSON_DATETIME_FUNCTION),
+										  errmsg("jsonpath item method .%s() can only be applied to a string",
+												 jspOperationName(jsp->type)))));
+				}
+				if (jsp->content.args.left && jsp->content.args.right)
 				{
-					text		*from, *to;
 					char		*from_str, *to_str;
 					int			from_len, to_len;
-					ErrorSaveContext escontext = {T_ErrorSaveContext};
 
 					jspGetLeftArg(jsp, &elem);
 					if (elem.type != jpiString)
 						elog(ERROR, "invalid jsonpath item type for .replace() from");
 
 					from_str = jspGetString(&elem, &from_len);
-					from = cstring_to_text_with_len(from_str, from_len);
-
-					if (jsp->content.args.right)
-					{
-						jspGetRightArg(jsp, &elem);
-						if (elem.type != jpiString)
-							elog(ERROR, "invalid jsonpath item type for .replace() to");
 
-						to_str = jspGetString(&elem, &to_len);
-						to = cstring_to_text_with_len(to_str, to_len);
+					jspGetRightArg(jsp, &elem);
+					if (elem.type != jpiString)
+						elog(ERROR, "invalid jsonpath item type for .replace() to");
 
-						replace_res = DirectFunctionCall3Coll(replace_text,
-							C_COLLATION_OID,
-							CStringGetTextDatum(tmp),
-							PointerGetDatum(from),
-							PointerGetDatum(to));
+					to_str = jspGetString(&elem, &to_len);
 
-					}
+					replacedTxt = TextDatumGetCString(DirectFunctionCall3Coll(replace_text,
+						C_COLLATION_OID,
+						CStringGetTextDatum(txt),
+						CStringGetTextDatum(from_str),
+						CStringGetTextDatum(to_str)));
 
-					res = jperOk;
+					jb = &jbv;
+					jb->type = jbvString;
+					jb->val.string.val = replacedTxt;
+					jb->val.string.len = strlen(replacedTxt);
 				}
-
-				if (res == jperNotFound)
+				else {
 					RETURN_ERROR(ereport(ERROR,
-										 (errcode(ERRCODE_NON_NUMERIC_SQL_JSON_ITEM),
-										  errmsg("jsonpath item method .%s() accepts two string arguments",
-												 jspOperationName(jsp->type)))));
-
-				jb = &jbv;
-				jb->type = jbvString;
-				jb->val.string.val = VARDATA_ANY(replace_res);
-				jb->val.string.len = VARSIZE_ANY_EXHDR(replace_res);
+											 (errcode(ERRCODE_INVALID_ARGUMENT_FOR_SQL_JSON_DATETIME_FUNCTION),
+											  errmsg("jsonpath item method .%s() can only be applied to a boolean, string, numeric, or datetime value",
+													 jspOperationName(jsp->type)))));
+					break;
+				}
 
 				res = executeNextItem(cxt, jsp, NULL, jb, found, true);
 			}
-- 
2.39.3 (Apple Git-146)

