1 files changed, 310 insertions, 40 deletions

diff --git a/src/backend/utils/adt/jsonpath_exec.c b/src/backend/utils/adt/jsonpath_exec.c
index eff3734b6ab..7811fa31e07 100644
--- a/src/backend/utils/adt/jsonpath_exec.c
+++ b/src/backend/utils/adt/jsonpath_exec.c
@@ -86,12 +86,16 @@ typedef struct JsonBaseObjectInfo
 	int			id;
 } JsonBaseObjectInfo;
 
+typedef int (*JsonPathVarCallback) (void *vars, char *varName, int varNameLen,
+									JsonbValue *val, JsonbValue *baseObject);
+
 /*
  * Context of jsonpath execution.
  */
 typedef struct JsonPathExecContext
 {
-	Jsonb	   *vars;			/* variables to substitute into jsonpath */
+	void	   *vars;			/* variables to substitute into jsonpath */
+	JsonPathVarCallback getVar;
 	JsonbValue *root;			/* for $ evaluation */
 	JsonbValue *current;		/* for @ evaluation */
 	JsonBaseObjectInfo baseObject;	/* "base object" for .keyvalue()
@@ -173,7 +177,8 @@ typedef JsonPathBool (*JsonPathPredicateCallback) (JsonPathItem *jsp,
 												   void *param);
 typedef Numeric (*BinaryArithmFunc) (Numeric num1, Numeric num2, bool *error);
 
-static JsonPathExecResult executeJsonPath(JsonPath *path, Jsonb *vars,
+static JsonPathExecResult executeJsonPath(JsonPath *path, void *vars,
+										  JsonPathVarCallback getVar,
 										  Jsonb *json, bool throwErrors,
 										  JsonValueList *result, bool useTz);
 static JsonPathExecResult executeItem(JsonPathExecContext *cxt,
@@ -225,7 +230,10 @@ static JsonPathExecResult appendBoolResult(JsonPathExecContext *cxt,
 static void getJsonPathItem(JsonPathExecContext *cxt, JsonPathItem *item,
 							JsonbValue *value);
 static void getJsonPathVariable(JsonPathExecContext *cxt,
-								JsonPathItem *variable, Jsonb *vars, JsonbValue *value);
+								JsonPathItem *variable, JsonbValue *value);
+static int getJsonPathVariableFromJsonb(void *varsJsonb, char *varName,
+										int varNameLen, JsonbValue *val,
+										JsonbValue *baseObject);
 static int	JsonbArraySize(JsonbValue *jb);
 static JsonPathBool executeComparison(JsonPathItem *cmp, JsonbValue *lv,
 									  JsonbValue *rv, void *p);
@@ -283,7 +291,8 @@ jsonb_path_exists_internal(FunctionCallInfo fcinfo, bool tz)
 		silent = PG_GETARG_BOOL(3);
 	}
 
-	res = executeJsonPath(jp, vars, jb, !silent, NULL, tz);
+	res = executeJsonPath(jp, vars, getJsonPathVariableFromJsonb,
+						  jb, !silent, NULL, tz);
 
 	PG_FREE_IF_COPY(jb, 0);
 	PG_FREE_IF_COPY(jp, 1);
@@ -338,7 +347,8 @@ jsonb_path_match_internal(FunctionCallInfo fcinfo, bool tz)
 		silent = PG_GETARG_BOOL(3);
 	}
 
-	(void) executeJsonPath(jp, vars, jb, !silent, &found, tz);
+	(void) executeJsonPath(jp, vars, getJsonPathVariableFromJsonb,
+						   jb, !silent, &found, tz);
 
 	PG_FREE_IF_COPY(jb, 0);
 	PG_FREE_IF_COPY(jp, 1);
@@ -416,7 +426,8 @@ jsonb_path_query_internal(FunctionCallInfo fcinfo, bool tz)
 		vars = PG_GETARG_JSONB_P_COPY(2);
 		silent = PG_GETARG_BOOL(3);
 
-		(void) executeJsonPath(jp, vars, jb, !silent, &found, tz);
+		(void) executeJsonPath(jp, vars, getJsonPathVariableFromJsonb,
+							   jb, !silent, &found, tz);
 
 		funcctx->user_fctx = JsonValueListGetList(&found);
 
@@ -463,7 +474,8 @@ jsonb_path_query_array_internal(FunctionCallInfo fcinfo, bool tz)
 	Jsonb	   *vars = PG_GETARG_JSONB_P(2);
 	bool		silent = PG_GETARG_BOOL(3);
 
-	(void) executeJsonPath(jp, vars, jb, !silent, &found, tz);
+	(void) executeJsonPath(jp, vars, getJsonPathVariableFromJsonb,
+						   jb, !silent, &found, tz);
 
 	PG_RETURN_JSONB_P(JsonbValueToJsonb(wrapItemsInArray(&found)));
 }
@@ -494,7 +506,8 @@ jsonb_path_query_first_internal(FunctionCallInfo fcinfo, bool tz)
 	Jsonb	   *vars = PG_GETARG_JSONB_P(2);
 	bool		silent = PG_GETARG_BOOL(3);
 
-	(void) executeJsonPath(jp, vars, jb, !silent, &found, tz);
+	(void) executeJsonPath(jp, vars, getJsonPathVariableFromJsonb,
+						   jb, !silent, &found, tz);
 
 	if (JsonValueListLength(&found) >= 1)
 		PG_RETURN_JSONB_P(JsonbValueToJsonb(JsonValueListHead(&found)));
@@ -536,8 +549,9 @@ jsonb_path_query_first_tz(PG_FUNCTION_ARGS)
  * In other case it tries to find all the satisfied result items.
  */
 static JsonPathExecResult
-executeJsonPath(JsonPath *path, Jsonb *vars, Jsonb *json, bool throwErrors,
-				JsonValueList *result, bool useTz)
+executeJsonPath(JsonPath *path, void *vars, JsonPathVarCallback getVar,
+				Jsonb *json, bool throwErrors, JsonValueList *result,
+				bool useTz)
 {
 	JsonPathExecContext cxt;
 	JsonPathExecResult res;
@@ -549,22 +563,16 @@ executeJsonPath(JsonPath *path, Jsonb *vars, Jsonb *json, bool throwErrors,
 	if (!JsonbExtractScalar(&json->root, &jbv))
 		JsonbInitBinary(&jbv, json);
 
-	if (vars && !JsonContainerIsObject(&vars->root))
-	{
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("\"vars\" argument is not an object"),
-				 errdetail("Jsonpath parameters should be encoded as key-value pairs of \"vars\" object.")));
-	}
-
 	cxt.vars = vars;
+	cxt.getVar = getVar;
 	cxt.laxMode = (path->header & JSONPATH_LAX) != 0;
 	cxt.ignoreStructuralErrors = cxt.laxMode;
 	cxt.root = &jbv;
 	cxt.current = &jbv;
 	cxt.baseObject.jbc = NULL;
 	cxt.baseObject.id = 0;
-	cxt.lastGeneratedObjectId = vars ? 2 : 1;
+	/* 1 + number of base objects in vars */
+	cxt.lastGeneratedObjectId = 1 + getVar(vars, NULL, 0, NULL, NULL);
 	cxt.innermostArraySize = -1;
 	cxt.throwErrors = throwErrors;
 	cxt.useTz = useTz;
@@ -2093,7 +2101,7 @@ getJsonPathItem(JsonPathExecContext *cxt, JsonPathItem *item,
 												 &value->val.string.len);
 			break;
 		case jpiVariable:
-			getJsonPathVariable(cxt, item, cxt->vars, value);
+			getJsonPathVariable(cxt, item, value);
 			return;
 		default:
 			elog(ERROR, "unexpected jsonpath item type");
@@ -2105,42 +2113,63 @@ getJsonPathItem(JsonPathExecContext *cxt, JsonPathItem *item,
  */
 static void
 getJsonPathVariable(JsonPathExecContext *cxt, JsonPathItem *variable,
-					Jsonb *vars, JsonbValue *value)
+					JsonbValue *value)
 {
 	char	   *varName;
 	int			varNameLength;
+	JsonbValue	baseObject;
+	int			baseObjectId;
+
+	Assert(variable->type == jpiVariable);
+	varName = jspGetString(variable, &varNameLength);
+
+	if (!cxt->vars ||
+		(baseObjectId = cxt->getVar(cxt->vars, varName, varNameLength, value,
+									&baseObject)) < 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+				 errmsg("could not find jsonpath variable \"%s\"",
+						pnstrdup(varName, varNameLength))));
+
+	if (baseObjectId > 0)
+		setBaseObject(cxt, &baseObject, baseObjectId);
+}
+
+static int
+getJsonPathVariableFromJsonb(void *varsJsonb, char *varName, int varNameLength,
+							 JsonbValue *value, JsonbValue *baseObject)
+{
+	Jsonb	   *vars = varsJsonb;
 	JsonbValue	tmp;
 	JsonbValue *v;
 
-	if (!vars)
+	if (!varName)
 	{
-		value->type = jbvNull;
-		return;
+		if (vars && !JsonContainerIsObject(&vars->root))
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("\"vars\" argument is not an object"),
+					 errdetail("Jsonpath parameters should be encoded as key-value pairs of \"vars\" object.")));
+		}
+
+		return vars ? 1 : 0;	/* count of base objects */
 	}
 
-	Assert(variable->type == jpiVariable);
-	varName = jspGetString(variable, &varNameLength);
 	tmp.type = jbvString;
 	tmp.val.string.val = varName;
 	tmp.val.string.len = varNameLength;
 
 	v = findJsonbValueFromContainer(&vars->root, JB_FOBJECT, &tmp);
 
-	if (v)
-	{
-		*value = *v;
-		pfree(v);
-	}
-	else
-	{
-		ereport(ERROR,
-				(errcode(ERRCODE_UNDEFINED_OBJECT),
-				 errmsg("could not find jsonpath variable \"%s\"",
-						pnstrdup(varName, varNameLength))));
-	}
+	if (!v)
+		return -1;
 
-	JsonbInitBinary(&tmp, vars);
-	setBaseObject(cxt, &tmp, 1);
+	*value = *v;
+	pfree(v);
+
+	JsonbInitBinary(baseObject, vars);
+	return 1;
 }
 
 /**************** Support functions for JsonPath execution *****************/
@@ -2797,3 +2826,244 @@ compareDatetime(Datum val1, Oid typid1, Datum val2, Oid typid2,
 
 	return DatumGetInt32(DirectFunctionCall2(cmpfunc, val1, val2));
 }
+
+/********************Interface to pgsql's executor***************************/
+
+bool
+JsonPathExists(Datum jb, JsonPath *jp, List *vars, bool *error)
+{
+	JsonPathExecResult res = executeJsonPath(jp, vars, EvalJsonPathVar,
+											 DatumGetJsonbP(jb), !error, NULL,
+											 true);
+
+	Assert(error || !jperIsError(res));
+
+	if (error && jperIsError(res))
+		*error = true;
+
+	return res == jperOk;
+}
+
+Datum
+JsonPathQuery(Datum jb, JsonPath *jp, JsonWrapper wrapper, bool *empty,
+			  bool *error, List *vars)
+{
+	JsonbValue *first;
+	bool		wrap;
+	JsonValueList found = {0};
+	JsonPathExecResult res PG_USED_FOR_ASSERTS_ONLY;
+	int			count;
+
+	res = executeJsonPath(jp, vars, EvalJsonPathVar, DatumGetJsonbP(jb), !error,
+						  &found, true);
+
+	Assert(error || !jperIsError(res));
+
+	if (error && jperIsError(res))
+	{
+		*error = true;
+		*empty = false;
+		return (Datum) 0;
+	}
+
+	count = JsonValueListLength(&found);
+
+	first = count ? JsonValueListHead(&found) : NULL;
+
+	if (!first)
+		wrap = false;
+	else if (wrapper == JSW_NONE)
+		wrap = false;
+	else if (wrapper == JSW_UNCONDITIONAL)
+		wrap = true;
+	else if (wrapper == JSW_CONDITIONAL)
+		wrap = count > 1 ||
+			IsAJsonbScalar(first) ||
+			(first->type == jbvBinary &&
+			 JsonContainerIsScalar(first->val.binary.data));
+	else
+	{
+		elog(ERROR, "unrecognized json wrapper %d", wrapper);
+		wrap = false;
+	}
+
+	if (wrap)
+		return JsonbPGetDatum(JsonbValueToJsonb(wrapItemsInArray(&found)));
+
+	if (count > 1)
+	{
+		if (error)
+		{
+			*error = true;
+			return (Datum) 0;
+		}
+
+		ereport(ERROR,
+				(errcode(ERRCODE_MORE_THAN_ONE_SQL_JSON_ITEM),
+				 errmsg("JSON path expression in JSON_QUERY should return "
+						"singleton item without wrapper"),
+				 errhint("use WITH WRAPPER clause to wrap SQL/JSON item "
+						 "sequence into array")));
+	}
+
+	if (first)
+		return JsonbPGetDatum(JsonbValueToJsonb(first));
+
+	*empty = true;
+	return PointerGetDatum(NULL);
+}
+
+JsonbValue *
+JsonPathValue(Datum jb, JsonPath *jp, bool *empty, bool *error, List *vars)
+{
+	JsonbValue   *res;
+	JsonValueList found = { 0 };
+	JsonPathExecResult jper PG_USED_FOR_ASSERTS_ONLY;
+	int			count;
+
+	jper = executeJsonPath(jp, vars, EvalJsonPathVar, DatumGetJsonbP(jb), !error,
+						   &found, true);
+
+	Assert(error || !jperIsError(jper));
+
+	if (error && jperIsError(jper))
+	{
+		*error = true;
+		*empty = false;
+		return NULL;
+	}
+
+	count = JsonValueListLength(&found);
+
+	*empty = !count;
+
+	if (*empty)
+		return NULL;
+
+	if (count > 1)
+	{
+		if (error)
+		{
+			*error = true;
+			return NULL;
+		}
+
+		ereport(ERROR,
+				(errcode(ERRCODE_MORE_THAN_ONE_SQL_JSON_ITEM),
+				 errmsg("JSON path expression in JSON_VALUE should return "
+						"singleton scalar item")));
+	}
+
+	res = JsonValueListHead(&found);
+
+	if (res->type == jbvBinary &&
+		JsonContainerIsScalar(res->val.binary.data))
+		JsonbExtractScalar(res->val.binary.data, res);
+
+	if (!IsAJsonbScalar(res))
+	{
+		if (error)
+		{
+			*error = true;
+			return NULL;
+		}
+
+		ereport(ERROR,
+				(errcode(ERRCODE_SQL_JSON_SCALAR_REQUIRED),
+				 errmsg("JSON path expression in JSON_VALUE should return "
+						"singleton scalar item")));
+	}
+
+	if (res->type == jbvNull)
+		return NULL;
+
+	return res;
+}
+
+static void
+JsonbValueInitNumericDatum(JsonbValue *jbv, Datum num)
+{
+	jbv->type = jbvNumeric;
+	jbv->val.numeric = DatumGetNumeric(num);
+}
+
+void
+JsonItemFromDatum(Datum val, Oid typid, int32 typmod, JsonbValue *res)
+{
+	switch (typid)
+	{
+		case BOOLOID:
+			res->type = jbvBool;
+			res->val.boolean = DatumGetBool(val);
+			break;
+		case NUMERICOID:
+			JsonbValueInitNumericDatum(res, val);
+			break;
+		case INT2OID:
+			JsonbValueInitNumericDatum(res, DirectFunctionCall1(int2_numeric, val));
+			break;
+		case INT4OID:
+			JsonbValueInitNumericDatum(res, DirectFunctionCall1(int4_numeric, val));
+			break;
+		case INT8OID:
+			JsonbValueInitNumericDatum(res, DirectFunctionCall1(int8_numeric, val));
+			break;
+		case FLOAT4OID:
+			JsonbValueInitNumericDatum(res, DirectFunctionCall1(float4_numeric, val));
+			break;
+		case FLOAT8OID:
+			JsonbValueInitNumericDatum(res, DirectFunctionCall1(float8_numeric, val));
+			break;
+		case TEXTOID:
+		case VARCHAROID:
+			res->type = jbvString;
+			res->val.string.val = VARDATA_ANY(val);
+			res->val.string.len = VARSIZE_ANY_EXHDR(val);
+			break;
+		case DATEOID:
+		case TIMEOID:
+		case TIMETZOID:
+		case TIMESTAMPOID:
+		case TIMESTAMPTZOID:
+			res->type = jbvDatetime;
+			res->val.datetime.value = val;
+			res->val.datetime.typid = typid;
+			res->val.datetime.typmod = typmod;
+			res->val.datetime.tz = 0;
+			break;
+		case JSONBOID:
+			{
+				JsonbValue *jbv = res;
+				Jsonb	   *jb = DatumGetJsonbP(val);
+
+				if (JsonContainerIsScalar(&jb->root))
+				{
+					bool		res PG_USED_FOR_ASSERTS_ONLY;
+
+					res = JsonbExtractScalar(&jb->root, jbv);
+					Assert(res);
+				}
+				else
+					JsonbInitBinary(jbv, jb);
+				break;
+			}
+		case JSONOID:
+			{
+				text	   *txt = DatumGetTextP(val);
+				char	   *str = text_to_cstring(txt);
+				Jsonb	   *jb =
+					DatumGetJsonbP(DirectFunctionCall1(jsonb_in,
+													   CStringGetDatum(str)));
+
+				pfree(str);
+
+				JsonItemFromDatum(JsonbPGetDatum(jb), JSONBOID, -1, res);
+				break;
+			}
+		default:
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("only bool, numeric and text types could be "
+							"casted to supported jsonpath types.")));
+	}
+}