From 6f59b763a37703be52560692c91bb16c9545a4cd Mon Sep 17 00:00:00 2001
From: Florents Tselai <florents.tselai@gmail.com>
Date: Sun, 15 Sep 2024 04:02:19 +0300
Subject: [PATCH v2 1/3] Add first working version for jsonpath
 $.replace(string, string) function.

---
 src/backend/utils/adt/jsonpath.c             | 21 ++++++
 src/backend/utils/adt/jsonpath_exec.c        | 73 ++++++++++++++++++++
 src/backend/utils/adt/jsonpath_gram.y        | 24 ++++++-
 src/backend/utils/adt/jsonpath_scan.l        |  1 +
 src/include/utils/jsonpath.h                 |  1 +
 src/test/regress/expected/jsonb_jsonpath.out | 39 +++++++++++
 src/test/regress/expected/jsonpath.out       |  6 ++
 src/test/regress/sql/jsonb_jsonpath.sql      | 11 +++
 src/test/regress/sql/jsonpath.sql            |  1 +
 9 files changed, 174 insertions(+), 3 deletions(-)

diff --git a/src/backend/utils/adt/jsonpath.c b/src/backend/utils/adt/jsonpath.c
index 0f691bc5f0..3daf2e2d81 100644
--- a/src/backend/utils/adt/jsonpath.c
+++ b/src/backend/utils/adt/jsonpath.c
@@ -298,6 +298,7 @@ flattenJsonPathParseItem(StringInfo buf, int *result, struct Node *escontext,
 		case jpiMod:
 		case jpiStartsWith:
 		case jpiDecimal:
+		case jpiReplaceFunc:
 			{
 				/*
 				 * First, reserve place for left/right arg's positions, then
@@ -831,6 +832,21 @@ printJsonPathItem(StringInfo buf, JsonPathItem *v, bool inKey,
 			}
 			appendStringInfoChar(buf, ')');
 			break;
+		case jpiReplaceFunc:
+			appendStringInfoString(buf, ".replace(");
+			if (v->content.args.left)
+			{
+				jspGetLeftArg(v, &elem);
+				printJsonPathItem(buf, &elem, false, false);
+			}
+			if (v->content.args.right)
+			{
+				appendStringInfoChar(buf, ',');
+				jspGetRightArg(v, &elem);
+				printJsonPathItem(buf, &elem, false, false);
+			}
+			appendStringInfoChar(buf, ')');
+			break;
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", v->type);
 	}
@@ -906,6 +922,8 @@ jspOperationName(JsonPathItemType type)
 			return "number";
 		case jpiStringFunc:
 			return "string";
+		case jpiReplaceFunc:
+			return "replace";
 		case jpiTime:
 			return "time";
 		case jpiTimeTz:
@@ -1041,6 +1059,7 @@ jspInitByBuffer(JsonPathItem *v, char *base, int32 pos)
 		case jpiMod:
 		case jpiStartsWith:
 		case jpiDecimal:
+		case jpiReplaceFunc:
 			read_int32(v->content.args.left, base, pos);
 			read_int32(v->content.args.right, base, pos);
 			break;
@@ -1149,6 +1168,7 @@ jspGetNext(JsonPathItem *v, JsonPathItem *a)
 			   v->type == jpiInteger ||
 			   v->type == jpiNumber ||
 			   v->type == jpiStringFunc ||
+			   v->type == jpiReplaceFunc ||
 			   v->type == jpiTime ||
 			   v->type == jpiTimeTz ||
 			   v->type == jpiTimestamp ||
@@ -1501,6 +1521,7 @@ jspIsMutableWalker(JsonPathItem *jpi, struct JsonPathMutableContext *cxt)
 			case jpiInteger:
 			case jpiNumber:
 			case jpiStringFunc:
+			case jpiReplaceFunc:
 				status = jpdsNonDateTime;
 				break;
 
diff --git a/src/backend/utils/adt/jsonpath_exec.c b/src/backend/utils/adt/jsonpath_exec.c
index 1184cba983..3d23e44d34 100644
--- a/src/backend/utils/adt/jsonpath_exec.c
+++ b/src/backend/utils/adt/jsonpath_exec.c
@@ -1659,6 +1659,79 @@ 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);
+
+				if (unwrap && JsonbType(jb) == jbvArray)
+					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 (jsp->content.args.left)
+				{
+					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);
+
+						replace_res = DirectFunctionCall3Coll(replace_text,
+							C_COLLATION_OID,
+							CStringGetTextDatum(tmp),
+							PointerGetDatum(from),
+							PointerGetDatum(to));
+
+					}
+
+					res = jperOk;
+				}
+
+				if (res == jperNotFound)
+					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);
+
+				res = executeNextItem(cxt, jsp, NULL, jb, found, true);
+			}
+			break;
+
 
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", jsp->type);
diff --git a/src/backend/utils/adt/jsonpath_gram.y b/src/backend/utils/adt/jsonpath_gram.y
index 8733a0eac6..c33b73e98b 100644
--- a/src/backend/utils/adt/jsonpath_gram.y
+++ b/src/backend/utils/adt/jsonpath_gram.y
@@ -83,7 +83,7 @@ static bool makeItemLikeRegex(JsonPathParseItem *expr,
 %token	<str>		ABS_P SIZE_P TYPE_P FLOOR_P DOUBLE_P CEILING_P KEYVALUE_P
 %token	<str>		DATETIME_P
 %token	<str>		BIGINT_P BOOLEAN_P DATE_P DECIMAL_P INTEGER_P NUMBER_P
-%token	<str>		STRINGFUNC_P TIME_P TIME_TZ_P TIMESTAMP_P TIMESTAMP_TZ_P
+%token	<str>		STRINGFUNC_P TIME_P TIME_TZ_P TIMESTAMP_P TIMESTAMP_TZ_P REPLACEFUNC_P
 
 %type	<result>	result
 
@@ -91,9 +91,9 @@ static bool makeItemLikeRegex(JsonPathParseItem *expr,
 					any_path accessor_op key predicate delimited_predicate
 					index_elem starts_with_initial expr_or_predicate
 					datetime_template opt_datetime_template csv_elem
-					datetime_precision opt_datetime_precision
+					datetime_precision opt_datetime_precision replace_arg_elem
 
-%type	<elems>		accessor_expr csv_list opt_csv_list
+%type	<elems>		accessor_expr csv_list opt_csv_list replace_arg_list
 
 %type	<indexs>	index_list
 
@@ -266,6 +266,16 @@ accessor_op:
 						 errmsg("invalid input syntax for type %s", "jsonpath"),
 						 errdetail(".decimal() can only have an optional precision[,scale].")));
 		}
+	| '.' REPLACEFUNC_P '(' replace_arg_list ')'
+		{
+			if (list_length($4) == 2)
+				$$ = makeItemBinary(jpiReplaceFunc, linitial($4), lsecond($4));
+			else
+				ereturn(escontext, false,
+						(errcode(ERRCODE_SYNTAX_ERROR),
+						 errmsg("invalid input syntax for type %s", "jsonpath"),
+						 errdetail(".replace() accepts two arguments.")));
+		}
 	| '.' DATETIME_P '(' opt_datetime_template ')'
 		{ $$ = makeItemUnary(jpiDatetime, $4); }
 	| '.' TIME_P '(' opt_datetime_precision ')'
@@ -315,6 +325,14 @@ opt_datetime_template:
 	| /* EMPTY */					{ $$ = NULL; }
 	;
 
+replace_arg_elem:
+	STRING_P						{ $$ = makeItemString(&$1); }
+	;
+
+replace_arg_list:
+	replace_arg_elem							{ $$ = list_make1($1); }
+	| replace_arg_list ',' replace_arg_elem	{ $$ = lappend($1, $3); }
+	;
 key:
 	key_name						{ $$ = makeItemKey(&$1); }
 	;
diff --git a/src/backend/utils/adt/jsonpath_scan.l b/src/backend/utils/adt/jsonpath_scan.l
index 7acda77837..890522b323 100644
--- a/src/backend/utils/adt/jsonpath_scan.l
+++ b/src/backend/utils/adt/jsonpath_scan.l
@@ -423,6 +423,7 @@ static const JsonPathKeyword keywords[] = {
 	{ 7, false,	CEILING_P,	"ceiling"},
 	{ 7, false,	DECIMAL_P,	"decimal"},
 	{ 7, false,	INTEGER_P,	"integer"},
+	{ 7, false, REPLACEFUNC_P, "replace"},
 	{ 7, false,	TIME_TZ_P,	"time_tz"},
 	{ 7, false,	UNKNOWN_P,	"unknown"},
 	{ 8, false,	DATETIME_P,	"datetime"},
diff --git a/src/include/utils/jsonpath.h b/src/include/utils/jsonpath.h
index ee35698d08..cbb35f02d3 100644
--- a/src/include/utils/jsonpath.h
+++ b/src/include/utils/jsonpath.h
@@ -115,6 +115,7 @@ typedef enum JsonPathItemType
 	jpiTimeTz,					/* .time_tz() item method */
 	jpiTimestamp,				/* .timestamp() item method */
 	jpiTimestampTz,				/* .timestamp_tz() item method */
+	jpiReplaceFunc,				/* .replace() item method */
 } JsonPathItemType;
 
 /* XQuery regex mode flags for LIKE_REGEX predicate */
diff --git a/src/test/regress/expected/jsonb_jsonpath.out b/src/test/regress/expected/jsonb_jsonpath.out
index acdf7e436f..2a1b2afff0 100644
--- a/src/test/regress/expected/jsonb_jsonpath.out
+++ b/src/test/regress/expected/jsonb_jsonpath.out
@@ -2719,6 +2719,45 @@ select jsonb_path_query('"2023-08-15 12:34:56"', '$.timestamp().string()');
 (1 row)
 
 rollback;
+-- Test .replace()
+select jsonb_path_query('null', '$.replace("x", "bye")');
+ERROR:  jsonpath item method .replace() can only be applied to a string
+select jsonb_path_query('null', '$.replace("x", "bye")', silent => true);
+ jsonb_path_query 
+------------------
+(0 rows)
+
+select jsonb_path_query('[]', '$.replace("x", "bye")');
+ jsonb_path_query 
+------------------
+(0 rows)
+
+select jsonb_path_query('{}', '$.replace("x", "bye")');
+ERROR:  jsonpath item method .replace() can only be applied to a string
+select jsonb_path_query('[]', 'strict $.replace("x", "bye")', silent => true);
+ jsonb_path_query 
+------------------
+(0 rows)
+
+select jsonb_path_query('{}', '$.replace("x", "bye")', silent => true);
+ jsonb_path_query 
+------------------
+(0 rows)
+
+select jsonb_path_query('1.23', '$.replace("x", "bye")');
+ERROR:  jsonpath item method .replace() can only be applied to a string
+select jsonb_path_query('"hello world"', '$.replace("hello","bye")');
+ jsonb_path_query 
+------------------
+ "bye world"
+(1 row)
+
+select jsonb_path_query('"hello world"', '$.replace("hello","bye") starts with "bye"');
+ jsonb_path_query 
+------------------
+ true
+(1 row)
+
 -- Test .time()
 select jsonb_path_query('null', '$.time()');
 ERROR:  jsonpath item method .time() can only be applied to a string
diff --git a/src/test/regress/expected/jsonpath.out b/src/test/regress/expected/jsonpath.out
index fd9bd755f5..871b5e76c7 100644
--- a/src/test/regress/expected/jsonpath.out
+++ b/src/test/regress/expected/jsonpath.out
@@ -435,6 +435,12 @@ select '$.string()'::jsonpath;
  $.string()
 (1 row)
 
+select '$.replace("hello","bye")'::jsonpath;
+         jsonpath         
+--------------------------
+ $.replace("hello","bye")
+(1 row)
+
 select '$.time()'::jsonpath;
  jsonpath 
 ----------
diff --git a/src/test/regress/sql/jsonb_jsonpath.sql b/src/test/regress/sql/jsonb_jsonpath.sql
index da3f7969ca..c1b8491387 100644
--- a/src/test/regress/sql/jsonb_jsonpath.sql
+++ b/src/test/regress/sql/jsonb_jsonpath.sql
@@ -619,6 +619,17 @@ select jsonb_path_query('"2023-08-15 12:34:56 +5:30"', '$.timestamp_tz().string(
 select jsonb_path_query('"2023-08-15 12:34:56"', '$.timestamp().string()');
 rollback;
 
+-- Test .replace()
+select jsonb_path_query('null', '$.replace("x", "bye")');
+select jsonb_path_query('null', '$.replace("x", "bye")', silent => true);
+select jsonb_path_query('[]', '$.replace("x", "bye")');
+select jsonb_path_query('{}', '$.replace("x", "bye")');
+select jsonb_path_query('[]', 'strict $.replace("x", "bye")', silent => true);
+select jsonb_path_query('{}', '$.replace("x", "bye")', silent => true);
+select jsonb_path_query('1.23', '$.replace("x", "bye")');
+select jsonb_path_query('"hello world"', '$.replace("hello","bye")');
+select jsonb_path_query('"hello world"', '$.replace("hello","bye") starts with "bye"');
+
 -- Test .time()
 select jsonb_path_query('null', '$.time()');
 select jsonb_path_query('true', '$.time()');
diff --git a/src/test/regress/sql/jsonpath.sql b/src/test/regress/sql/jsonpath.sql
index 61a5270d4e..81704303bd 100644
--- a/src/test/regress/sql/jsonpath.sql
+++ b/src/test/regress/sql/jsonpath.sql
@@ -78,6 +78,7 @@ select '$.boolean()'::jsonpath;
 select '$.date()'::jsonpath;
 select '$.decimal(4,2)'::jsonpath;
 select '$.string()'::jsonpath;
+select '$.replace("hello","bye")'::jsonpath;
 select '$.time()'::jsonpath;
 select '$.time(6)'::jsonpath;
 select '$.time_tz()'::jsonpath;
-- 
2.39.3 (Apple Git-146)

