Fix jsonb replace and delete on scalars and empty structures

These operations now error out if attempted on scalars, and simply
return the input if attempted on empty arrays or objects. Along the way
we remove the unnecessary cloning of the input when it's known to be
unchanged. Regression tests covering these cases are added.

4 files changed, 165 insertions, 53 deletions

diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index 2d414551276..3dfdbf9d56f 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -3332,7 +3332,7 @@ jsonb_concat(PG_FUNCTION_ARGS)
 		out = JsonbValueToJsonb(res);
 	}
 
-	PG_RETURN_POINTER(out);
+	PG_RETURN_JSONB(out);
 }
 
 
@@ -3349,7 +3349,6 @@ jsonb_delete(PG_FUNCTION_ARGS)
 	text	   *key = PG_GETARG_TEXT_PP(1);
 	char	   *keyptr = VARDATA_ANY(key);
 	int			keylen = VARSIZE_ANY_EXHDR(key);
-	Jsonb	   *out = palloc(VARSIZE(in));
 	JsonbParseState *state = NULL;
 	JsonbIterator *it;
 	uint32		r;
@@ -3357,10 +3356,13 @@ jsonb_delete(PG_FUNCTION_ARGS)
 			   *res = NULL;
 	bool		skipNested = false;
 
-	SET_VARSIZE(out, VARSIZE(in));
+	if (JB_ROOT_IS_SCALAR(in))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("cannot delete from scalar")));
 
 	if (JB_ROOT_COUNT(in) == 0)
-		PG_RETURN_POINTER(out);
+		PG_RETURN_JSONB(in);
 
 	it = JsonbIteratorInit(&in->root);
 
@@ -3382,13 +3384,9 @@ jsonb_delete(PG_FUNCTION_ARGS)
 		res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 	}
 
-	if (res == NULL || (res->type == jbvArray && res->val.array.nElems == 0) ||
-		(res->type == jbvObject && res->val.object.nPairs == 0))
-		SET_VARSIZE(out, VARHDRSZ);
-	else
-		out = JsonbValueToJsonb(res);
+	Assert(res != NULL);
 
-	PG_RETURN_POINTER(out);
+	PG_RETURN_JSONB(JsonbValueToJsonb(res));
 }
 
 /*
@@ -3403,7 +3401,6 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 {
 	Jsonb	   *in = PG_GETARG_JSONB(0);
 	int			idx = PG_GETARG_INT32(1);
-	Jsonb	   *out = palloc(VARSIZE(in));
 	JsonbParseState *state = NULL;
 	JsonbIterator *it;
 	uint32		r,
@@ -3412,11 +3409,13 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 	JsonbValue	v,
 			   *res = NULL;
 
+	if (JB_ROOT_IS_SCALAR(in))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("cannot delete from scalar")));
+
 	if (JB_ROOT_COUNT(in) == 0)
-	{
-		memcpy(out, in, VARSIZE(in));
-		PG_RETURN_POINTER(out);
-	}
+		PG_RETURN_JSONB(in);
 
 	it = JsonbIteratorInit(&in->root);
 
@@ -3435,10 +3434,7 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 	}
 
 	if (idx >= n)
-	{
-		memcpy(out, in, VARSIZE(in));
-		PG_RETURN_POINTER(out);
-	}
+		PG_RETURN_JSONB(in);
 
 	pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 
@@ -3457,13 +3453,9 @@ jsonb_delete_idx(PG_FUNCTION_ARGS)
 		res = pushJsonbValue(&state, r, r < WJB_BEGIN_ARRAY ? &v : NULL);
 	}
 
-	if (res == NULL || (res->type == jbvArray && res->val.array.nElems == 0) ||
-		(res->type == jbvObject && res->val.object.nPairs == 0))
-		SET_VARSIZE(out, VARHDRSZ);
-	else
-		out = JsonbValueToJsonb(res);
+	Assert (res != NULL);
 
-	PG_RETURN_POINTER(out);
+	PG_RETURN_JSONB(JsonbValueToJsonb(res));
 }
 
 /*
@@ -3475,7 +3467,6 @@ jsonb_replace(PG_FUNCTION_ARGS)
 	Jsonb	   *in = PG_GETARG_JSONB(0);
 	ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
 	Jsonb	   *newval = PG_GETARG_JSONB(2);
-	Jsonb	   *out = palloc(VARSIZE(in) + VARSIZE(newval));
 	JsonbValue *res = NULL;
 	Datum	   *path_elems;
 	bool	   *path_nulls;
@@ -3488,31 +3479,27 @@ jsonb_replace(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
 				 errmsg("wrong number of array subscripts")));
 
+	if (JB_ROOT_IS_SCALAR(in))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("cannot replace path in scalar")));
+
 	if (JB_ROOT_COUNT(in) == 0)
-	{
-		memcpy(out, in, VARSIZE(in));
-		PG_RETURN_POINTER(out);
-	}
+		PG_RETURN_JSONB(in);
 
 	deconstruct_array(path, TEXTOID, -1, false, 'i',
 					  &path_elems, &path_nulls, &path_len);
 
 	if (path_len == 0)
-	{
-		memcpy(out, in, VARSIZE(in));
-		PG_RETURN_POINTER(out);
-	}
+		PG_RETURN_JSONB(in);
 
 	it = JsonbIteratorInit(&in->root);
 
 	res = replacePath(&it, path_elems, path_nulls, path_len, &st, 0, newval);
 
-	if (res == NULL)
-		SET_VARSIZE(out, VARHDRSZ);
-	else
-		out = JsonbValueToJsonb(res);
+	Assert (res != NULL);
 
-	PG_RETURN_POINTER(out);
+	PG_RETURN_JSONB(JsonbValueToJsonb(res));
 }
 
 
@@ -3524,7 +3511,6 @@ jsonb_delete_path(PG_FUNCTION_ARGS)
 {
 	Jsonb	   *in = PG_GETARG_JSONB(0);
 	ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
-	Jsonb	   *out = palloc(VARSIZE(in));
 	JsonbValue *res = NULL;
 	Datum	   *path_elems;
 	bool	   *path_nulls;
@@ -3537,31 +3523,27 @@ jsonb_delete_path(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_ARRAY_SUBSCRIPT_ERROR),
 				 errmsg("wrong number of array subscripts")));
 
+	if (JB_ROOT_IS_SCALAR(in))
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("cannot delete path in scalar")));
+
 	if (JB_ROOT_COUNT(in) == 0)
-	{
-		memcpy(out, in, VARSIZE(in));
-		PG_RETURN_POINTER(out);
-	}
+		PG_RETURN_JSONB(in);
 
 	deconstruct_array(path, TEXTOID, -1, false, 'i',
 					  &path_elems, &path_nulls, &path_len);
 
 	if (path_len == 0)
-	{
-		memcpy(out, in, VARSIZE(in));
-		PG_RETURN_POINTER(out);
-	}
+		PG_RETURN_JSONB(in);
 
 	it = JsonbIteratorInit(&in->root);
 
 	res = replacePath(&it, path_elems, path_nulls, path_len, &st, 0, NULL);
 
-	if (res == NULL)
-		SET_VARSIZE(out, VARHDRSZ);
-	else
-		out = JsonbValueToJsonb(res);
+	Assert (res != NULL);
 
-	PG_RETURN_POINTER(out);
+	PG_RETURN_JSONB(JsonbValueToJsonb(res));
 }
 
 
diff --git a/src/test/regress/expected/jsonb.out b/src/test/regress/expected/jsonb.out
index 83201fb6ca5..c589cd1522a 100644
--- a/src/test/regress/expected/jsonb.out
+++ b/src/test/regress/expected/jsonb.out
@@ -3175,3 +3175,60 @@ select '{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb - '{d
  {"a": 1, "b": [1, 2], "c": {"1": 2}, "d": {"1": [3]}, "n": null}
 (1 row)
 
+-- empty structure and error conditions for delete and replace
+select '"a"'::jsonb - 'a'; -- error
+ERROR:  cannot delete from scalar
+select '{}'::jsonb - 'a';
+ ?column? 
+----------
+ {}
+(1 row)
+
+select '[]'::jsonb - 'a';
+ ?column? 
+----------
+ []
+(1 row)
+
+select '"a"'::jsonb - 1; -- error
+ERROR:  cannot delete from scalar
+select '{}'::jsonb -  1 ;
+ ?column? 
+----------
+ {}
+(1 row)
+
+select '[]'::jsonb - 1;
+ ?column? 
+----------
+ []
+(1 row)
+
+select '"a"'::jsonb - '{a}'::text[]; -- error
+ERROR:  cannot delete path in scalar
+select '{}'::jsonb - '{a}'::text[];
+ ?column? 
+----------
+ {}
+(1 row)
+
+select '[]'::jsonb - '{a}'::text[];
+ ?column? 
+----------
+ []
+(1 row)
+
+select jsonb_replace('"a"','{a}','"b"'); --error
+ERROR:  cannot replace path in scalar
+select jsonb_replace('{}','{a}','"b"');
+ jsonb_replace 
+---------------
+ {}
+(1 row)
+
+select jsonb_replace('[]','{1}','"b"');
+ jsonb_replace 
+---------------
+ []
+(1 row)
+
diff --git a/src/test/regress/expected/jsonb_1.out b/src/test/regress/expected/jsonb_1.out
index c8cb513391a..c4b51e56e37 100644
--- a/src/test/regress/expected/jsonb_1.out
+++ b/src/test/regress/expected/jsonb_1.out
@@ -3175,3 +3175,60 @@ select '{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb - '{d
  {"a": 1, "b": [1, 2], "c": {"1": 2}, "d": {"1": [3]}, "n": null}
 (1 row)
 
+-- empty structure and error conditions for delete and replace
+select '"a"'::jsonb - 'a'; -- error
+ERROR:  cannot delete from scalar
+select '{}'::jsonb - 'a';
+ ?column? 
+----------
+ {}
+(1 row)
+
+select '[]'::jsonb - 'a';
+ ?column? 
+----------
+ []
+(1 row)
+
+select '"a"'::jsonb - 1; -- error
+ERROR:  cannot delete from scalar
+select '{}'::jsonb -  1 ;
+ ?column? 
+----------
+ {}
+(1 row)
+
+select '[]'::jsonb - 1;
+ ?column? 
+----------
+ []
+(1 row)
+
+select '"a"'::jsonb - '{a}'::text[]; -- error
+ERROR:  cannot delete path in scalar
+select '{}'::jsonb - '{a}'::text[];
+ ?column? 
+----------
+ {}
+(1 row)
+
+select '[]'::jsonb - '{a}'::text[];
+ ?column? 
+----------
+ []
+(1 row)
+
+select jsonb_replace('"a"','{a}','"b"'); --error
+ERROR:  cannot replace path in scalar
+select jsonb_replace('{}','{a}','"b"');
+ jsonb_replace 
+---------------
+ {}
+(1 row)
+
+select jsonb_replace('[]','{1}','"b"');
+ jsonb_replace 
+---------------
+ []
+(1 row)
+
diff --git a/src/test/regress/sql/jsonb.sql b/src/test/regress/sql/jsonb.sql
index 808da9c0270..382a7fb6ae0 100644
--- a/src/test/regress/sql/jsonb.sql
+++ b/src/test/regress/sql/jsonb.sql
@@ -767,3 +767,19 @@ select jsonb_delete('{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'
 select '{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb - '{n}'::text[];
 select '{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb - '{b,-1}'::text[];
 select '{"n":null, "a":1, "b":[1,2], "c":{"1":2}, "d":{"1":[2,3]}}'::jsonb - '{d,1,0}'::text[];
+
+
+-- empty structure and error conditions for delete and replace
+
+select '"a"'::jsonb - 'a'; -- error
+select '{}'::jsonb - 'a';
+select '[]'::jsonb - 'a';
+select '"a"'::jsonb - 1; -- error
+select '{}'::jsonb -  1 ;
+select '[]'::jsonb - 1;
+select '"a"'::jsonb - '{a}'::text[]; -- error
+select '{}'::jsonb - '{a}'::text[];
+select '[]'::jsonb - '{a}'::text[];
+select jsonb_replace('"a"','{a}','"b"'); --error
+select jsonb_replace('{}','{a}','"b"');
+select jsonb_replace('[]','{1}','"b"');