Add custom jsonpath item methods and functions

Author: Nikita Glukhov

src/backend/utils/adt/jsonpath.c

@@ -576,6 +576,42 @@ copyJsonPathItem(JsonPathContext *cxt, JsonPathItem *item, int level,
 			}
 			break;
 
+		case jpiMethod:
+		case jpiFunction:
+			{
+				int32		nargs = item->content.func.nargs;
+				int			offset;
+				int			i;
+
+				/* assign cache id */
+				appendBinaryStringInfo(buf, (const char *) &cxt->id, sizeof(cxt->id));
+				++cxt->id;
+
+				appendBinaryStringInfo(buf, (char *) &nargs, sizeof(nargs));
+				offset = buf->len;
+				appendStringInfoSpaces(buf, sizeof(int32) * nargs);
+
+				appendBinaryStringInfo(buf, (char *) &item->content.func.namelen,
+									   sizeof(item->content.func.namelen));
+				appendBinaryStringInfo(buf, item->content.func.name,
+									   item->content.func.namelen);
+				appendStringInfoChar(buf, '\0');
+
+				for (i = 0; i < nargs; i++)
+				{
+					JsonPathItem arg;
+					int32		argpos;
+
+					jspGetFunctionArg(item, i, &arg);
+
+					argpos = copyJsonPathItem(cxt, &arg, level, NULL, NULL);
+
+					*(int32 *) &buf->data[offset] = argpos - pos;
+					offset += sizeof(int32);
+				}
+			}
+			break;
+
 		default:
 			elog(ERROR, "Unknown jsonpath item type: %d", item->type);
 	}
@@ -768,6 +804,39 @@ flattenJsonPathParseItem(JsonPathContext *cxt, JsonPathParseItem *item,
 				offset += sizeof(int32);
 			}
 			break;
+		case jpiMethod:
+		case jpiFunction:
+			{
+				int32		nargs = list_length(item->value.func.args);
+				ListCell   *lc;
+				int			offset;
+
+				/* assign cache id */
+				appendBinaryStringInfo(buf, (const char *) &cxt->id, sizeof(cxt->id));
+				++cxt->id;
+
+				appendBinaryStringInfo(buf, (char *) &nargs, sizeof(nargs));
+				offset = buf->len;
+				appendStringInfoSpaces(buf, sizeof(int32) * nargs);
+
+				appendBinaryStringInfo(buf, (char *) &item->value.func.namelen,
+									   sizeof(item->value.func.namelen));
+				appendBinaryStringInfo(buf, item->value.func.name,
+									   item->value.func.namelen);
+				appendStringInfoChar(buf, '\0');
+
+				foreach(lc, item->value.func.args)
+				{
+					int32		argpos =
+						flattenJsonPathParseItem(cxt, lfirst(lc),
+												 nestingLevel + 1,
+												 insideArraySubscript);
+
+					*(int32 *) &buf->data[offset] = argpos - pos;
+					offset += sizeof(int32);
+				}
+			}
+			break;
 		case jpiNull:
 			break;
 		case jpiRoot:
@@ -1278,6 +1347,31 @@ printJsonPathItem(StringInfo buf, JsonPathItem *v, bool inKey,
 			if (printBracketes || jspHasNext(v))
 				appendStringInfoChar(buf, ')');
 			break;
+		case jpiMethod:
+		case jpiFunction:
+			if (v->type == jpiMethod)
+			{
+				jspGetMethodItem(v, &elem);
+				printJsonPathItem(buf, &elem, false,
+								  operationPriority(elem.type) <=
+								  operationPriority(v->type));
+				appendStringInfoChar(buf, '.');
+			}
+
+			escape_json(buf, v->content.func.name);
+			appendStringInfoChar(buf, '(');
+
+			for (i = v->type == jpiMethod ? 1 : 0; i < v->content.func.nargs; i++)
+			{
+				if (i > (v->type == jpiMethod ? 1 : 0))
+					appendBinaryStringInfo(buf, ", ", 2);
+
+				jspGetFunctionArg(v, i, &elem);
+				printJsonPathItem(buf, &elem, false, elem.type == jpiSequence);
+			}
+
+			appendStringInfoChar(buf, ')');
+			break;
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", v->type);
 	}
@@ -1371,11 +1465,13 @@ operationPriority(JsonPathItemType op)
 		case jpiDiv:
 		case jpiMod:
 			return 4;
+		case jpiMethod:
+			return 5;
 		case jpiPlus:
 		case jpiMinus:
-			return 5;
-		default:
 			return 6;
+		default:
+			return 7;
 	}
 }
 
@@ -1483,6 +1579,15 @@ jspInitByBuffer(JsonPathItem *v, char *base, int32 pos)
 						 v->content.lambda.nparams);
 			read_int32(v->content.lambda.expr, base, pos);
 			break;
+		case jpiMethod:
+		case jpiFunction:
+			read_int32(v->content.func.id, base, pos);
+			read_int32(v->content.func.nargs, base, pos);
+			read_int32_n(v->content.func.args, base, pos,
+						 v->content.func.nargs);
+			read_int32(v->content.func.namelen, base, pos);
+			v->content.func.name = base + pos;
+			break;
 		case jpiNot:
 		case jpiExists:
 		case jpiIsUnknown:
@@ -1581,7 +1686,9 @@ jspGetNext(JsonPathItem *v, JsonPathItem *a)
 			   v->type == jpiArray ||
 			   v->type == jpiObject ||
 			   v->type == jpiLambda ||
-			   v->type == jpiArgument);
+			   v->type == jpiArgument ||
+			   v->type == jpiFunction ||
+			   v->type == jpiMethod);
 
 		if (a)
 			jspInitByBuffer(a, v->base, v->nextPos);
@@ -1717,6 +1824,24 @@ jspGetLambdaExpr(JsonPathItem *lambda, JsonPathItem *expr)
 	return expr;
 }
 
+JsonPathItem *
+jspGetFunctionArg(JsonPathItem *func, int index, JsonPathItem *arg)
+{
+	Assert(func->type == jpiMethod || func->type == jpiFunction);
+	Assert(index < func->content.func.nargs);
+
+	jspInitByBuffer(arg, func->base, func->content.func.args[index]);
+
+	return arg;
+}
+
+JsonPathItem *
+jspGetMethodItem(JsonPathItem *method, JsonPathItem *arg)
+{
+	Assert(method->type == jpiMethod);
+	return jspGetFunctionArg(method, 0, arg);
+}
+
 static void
 checkJsonPathArgsMismatch(JsonPath *jp1, JsonPath *jp2)
 {

src/backend/utils/adt/jsonpath_exec.c

@@ -66,6 +66,7 @@
 #include "lib/stringinfo.h"
 #include "miscadmin.h"
 #include "nodes/nodeFuncs.h"
+#include "parser/parse_func.h"
 #include "regex/regex.h"
 #include "utils/builtins.h"
 #include "utils/date.h"
@@ -296,6 +297,9 @@ static int getJsonPathVariableFromJsonx(void *varsJsonb, bool isJsonb,
 static JsonPathBool executeComparison(JsonPathItem *cmp, JsonItem *lv,
 									  JsonItem *rv, void *p);
 static int	compareNumeric(Numeric a, Numeric b);
+static JsonPathExecResult executeFunction(JsonPathExecContext *cxt,
+				JsonPathItem *jsp, JsonItem *js,
+				JsonValueList *result /*, bool needBool */);
 
 static void JsonItemInitNull(JsonItem *item);
 static void JsonItemInitBool(JsonItem *item, bool val);
@@ -839,6 +843,7 @@ executeJsonPath(JsonPath *path, void *vars, JsonPathVarCallback getVar,
 
 	cxt.vars = vars;
 	cxt.getVar = getVar;
+	cxt.args = NULL;
 	cxt.laxMode = (path->header & JSONPATH_LAX) != 0;
 	cxt.ignoreStructuralErrors = cxt.laxMode;
 	cxt.root = &jsi;
@@ -1912,6 +1917,11 @@ executeItemOptUnwrapTarget(JsonPathExecContext *cxt, JsonPathItem *jsp,
 			elog(ERROR, "jsonpath lambda expression cannot be executed directly");
 			break;
 
+		case jpiMethod:
+		case jpiFunction:
+			res = executeFunction(cxt, jsp, jb, found);
+			break;
+
 		default:
 			elog(ERROR, "unrecognized jsonpath item type: %d", jsp->type);
 	}
@@ -2052,6 +2062,130 @@ executeItemOptUnwrapResult(JsonPathExecContext *cxt, JsonPathItem *jsp,
 	return executeItem(cxt, jsp, jb, found);
 }
 
+typedef struct JsonPathFuncCache
+{
+	FmgrInfo	finfo;
+	JsonPathItem *args;
+	void	  **argscache;
+} JsonPathFuncCache;
+
+static JsonPathFuncCache *
+prepareFunctionCache(JsonPathExecContext *cxt, JsonPathItem *jsp)
+{
+	MemoryContext oldcontext;
+	JsonPathFuncCache *cache = cxt->cache[jsp->content.func.id];
+	List	   *funcname = list_make1(makeString(jsp->content.func.name));
+	Oid			argtypes[] = {JSONPATH_FCXTOID};
+	Oid			funcid;
+	int32		i;
+
+	if (cache)
+		return cache;
+
+	funcid = LookupFuncName(funcname,
+							sizeof(argtypes) / sizeof(argtypes[0]), argtypes,
+							false);
+
+	if (get_func_rettype(funcid) != INT8OID)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATATYPE_MISMATCH),
+				 errmsg("return type of jsonpath item function %s is not %s",
+						 NameListToString(funcname), format_type_be(INT8OID))));
+
+	oldcontext = MemoryContextSwitchTo(cxt->cache_mcxt);
+
+	cache = cxt->cache[jsp->content.func.id] = palloc0(sizeof(*cache));
+
+	fmgr_info(funcid, &cache->finfo);
+
+	cache->args = palloc(sizeof(*cache->args) * jsp->content.func.nargs);
+	cache->argscache = palloc0(sizeof(*cache->argscache) *
+							   jsp->content.func.nargs);
+
+	for (i = 0; i < jsp->content.func.nargs; i++)
+		jspGetFunctionArg(jsp, i, &cache->args[i]);
+
+	MemoryContextSwitchTo(oldcontext);
+
+	return cache;
+}
+
+static JsonPathExecResult
+executeFunction(JsonPathExecContext *cxt, JsonPathItem *jsp, JsonItem *jb,
+				JsonValueList *result /*, bool needBool */)
+{
+	JsonPathFuncCache *cache = prepareFunctionCache(cxt, jsp);
+	JsonPathFuncContext fcxt;
+	JsonValueList tmpres = {0};
+	JsonValueListIterator tmpiter;
+	JsonPathExecResult res;
+	JsonItem   *jsi;
+
+	fcxt.cxt = cxt;
+	fcxt.funcname = jsp->content.func.name;
+	fcxt.jb = jb;
+	fcxt.result = jspHasNext(jsp) ? &tmpres : result;
+	fcxt.args = cache->args;
+	fcxt.argscache = cache->argscache;
+	fcxt.nargs = jsp->content.func.nargs;
+
+	if (jsp->type == jpiMethod)
+	{
+		JsonValueList items = {0};
+		JsonValueListIterator iter;
+
+		/* skip first item argument */
+		fcxt.args++;
+		fcxt.argscache++;
+		fcxt.nargs--;
+
+		res = executeItem(cxt, &cache->args[0], jb, &items);
+
+		if (jperIsError(res))
+			return res;
+
+		JsonValueListInitIterator(&items, &iter);
+
+		while ((jsi = JsonValueListNext(&items, &iter)))
+		{
+			fcxt.item = jsi;
+
+			res = (JsonPathExecResult)
+				DatumGetPointer(FunctionCall2(&cache->finfo,
+											  PointerGetDatum(&fcxt),
+											  PointerGetDatum(NULL)));
+			if (jperIsError(res))
+				return res;
+		}
+	}
+	else
+	{
+		fcxt.item = NULL;
+
+		res = (JsonPathExecResult)
+			DatumGetPointer(FunctionCall2(&cache->finfo,
+										  PointerGetDatum(&fcxt),
+										  PointerGetDatum(NULL)));
+		if (jperIsError(res))
+			return res;
+	}
+
+	if (!jspHasNext(jsp))
+		return res;
+
+	JsonValueListInitIterator(&tmpres, &tmpiter);
+
+	while ((jsi = JsonValueListNext(&tmpres, &tmpiter)))
+	{
+		res = executeNextItem(cxt, jsp, NULL, jsi, result, false /*, needBool FIXME */);
+
+		if (jperIsError(res))
+			return res;
+	}
+
+	return res;
+}
+
 /*
  * Same as executeItemOptUnwrapResult(), but with error suppression.
  */