Add jsonb_insert

feodor · feodor · commit 0b62fd036e1a · 2016-04-06T19:25:00.000+03:00
It inserts a new value into an jsonb array at arbitrary position or
a new key to jsonb object.

Author: Dmitry Dolgov
Reviewers: Petr Jelinek, Vitaly Burovoy, Andrew Dunstan
diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
@@ -10902,6 +10902,9 @@ table2-mapping
   <indexterm>
    <primary>jsonb_set</primary>
   </indexterm>
+  <indexterm>
+   <primary>jsonb_insert</primary>
+  </indexterm>
   <indexterm>
    <primary>jsonb_pretty</primary>
   </indexterm>
@@ -11183,6 +11186,39 @@ table2-mapping
          </para><para><literal>[{"f1": 1, "f2": null, "f3": [2, 3, 4]}, 2]</literal>
         </para></entry>
        </row>
+      <row>
+       <entry>
+           <para><literal>
+           jsonb_insert(target jsonb, path text[], new_value jsonb, <optional><parameter>insert_after</parameter> <type>boolean</type></optional>)
+           </literal></para>
+       </entry>
+       <entry><para><type>jsonb</type></para></entry>
+       <entry>
+         Returns <replaceable>target</replaceable> with
+         <replaceable>new_value</replaceable> inserted. If
+         <replaceable>target</replaceable> section designated by
+         <replaceable>path</replaceable> is in a JSONB array,
+         <replaceable>new_value</replaceable> will be inserted before target or
+         after if <replaceable>insert_after</replaceable> is true (default is
+         <literal>false</>). If <replaceable>target</replaceable> section
+         designated by <replaceable>path</replaceable> is in JSONB object,
+         <replaceable>new_value</replaceable> will be inserted only if
+         <replaceable>target</replaceable> does not exist. As with the path
+         orientated operators, negative integers that appear in
+         <replaceable>path</replaceable> count from the end of JSON arrays.
+       </entry>
+       <entry>
+           <para><literal>
+               jsonb_insert('{"a": [0,1,2]}', '{a, 1}', '"new_value"')
+           </literal></para>
+           <para><literal>
+               jsonb_insert('{"a": [0,1,2]}', '{a, 1}', '"new_value"', true)
+           </literal></para>
+       </entry>
+       <entry><para><literal>{"a": [0, "new_value", 1, 2]}</literal>
+         </para><para><literal>{"a": [0, 1, "new_value", 2]}</literal>
+        </para></entry>
+       </row>
       <row>
        <entry><para><literal>jsonb_pretty(from_json jsonb)</literal>
          </para></entry>
@@ -11235,10 +11271,11 @@ table2-mapping
   <note>
     <para>
       All the items of the <literal>path</> parameter of <literal>jsonb_set</>
-      must be present in the <literal>target</>, unless
-      <literal>create_missing</> is true, in which case all but the last item
-      must be present. If these conditions are not met the <literal>target</>
-      is returned unchanged.
+      as well as <literal>jsonb_insert</> except the last item must be present
+      in the <literal>target</>. If <literal>create_missing</> is false, all
+      items of the <literal>path</> parameter of <literal>jsonb_set</> must be
+      present. If these conditions are not met the <literal>target</> is
+      returned unchanged.
     </para>
     <para>
       If the last path item is an object key, it will be created if it
diff --git a/src/backend/catalog/system_views.sql b/src/backend/catalog/system_views.sql
@@ -997,3 +997,11 @@ RETURNS text[]
 LANGUAGE INTERNAL
 STRICT IMMUTABLE
 AS 'parse_ident';
+
+CREATE OR REPLACE FUNCTION
+  jsonb_insert(jsonb_in jsonb, path text[] , replacement jsonb,
+            insert_after boolean DEFAULT false)
+RETURNS jsonb
+LANGUAGE INTERNAL
+STRICT IMMUTABLE
+AS 'jsonb_insert';
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
@@ -33,6 +33,15 @@
 #include "utils/memutils.h"
 #include "utils/typcache.h"
 
+/* Operations available for setPath */
+#define JB_PATH_NOOP					0x0000
+#define JB_PATH_CREATE					0x0001
+#define JB_PATH_DELETE					0x0002
+#define JB_PATH_INSERT_BEFORE			0x0004
+#define JB_PATH_INSERT_AFTER			0x0008
+#define JB_PATH_CREATE_OR_INSERT \
+	(JB_PATH_INSERT_BEFORE | JB_PATH_INSERT_AFTER | JB_PATH_CREATE)
+
 /* semantic action functions for json_object_keys */
 static void okeys_object_field_start(void *state, char *fname, bool isnull);
 static void okeys_array_start(void *state);
@@ -130,14 +139,14 @@ static JsonbValue *IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 static JsonbValue *setPath(JsonbIterator **it, Datum *path_elems,
 		bool *path_nulls, int path_len,
 		JsonbParseState **st, int level, Jsonb *newval,
-		bool create);
+		int op_type);
 static void setPathObject(JsonbIterator **it, Datum *path_elems,
 			  bool *path_nulls, int path_len, JsonbParseState **st,
 			  int level,
-			  Jsonb *newval, uint32 npairs, bool create);
+			  Jsonb *newval, uint32 npairs, int op_type);
 static void setPathArray(JsonbIterator **it, Datum *path_elems,
 			 bool *path_nulls, int path_len, JsonbParseState **st,
-			 int level, Jsonb *newval, uint32 nelems, bool create);
+			 int level, Jsonb *newval, uint32 nelems, int op_type);
 static void addJsonbToParseState(JsonbParseState **jbps, Jsonb *jb);
 
 /* state for json_object_keys */
@@ -3544,7 +3553,7 @@ jsonb_set(PG_FUNCTION_ARGS)
 	it = JsonbIteratorInit(&in->root);
 
 	res = setPath(&it, path_elems, path_nulls, path_len, &st,
-				  0, newval, create);
+				  0, newval, create ? JB_PATH_CREATE : JB_PATH_NOOP);
 
 	Assert(res != NULL);
 
@@ -3588,7 +3597,52 @@ jsonb_delete_path(PG_FUNCTION_ARGS)
 
 	it = JsonbIteratorInit(&in->root);
 
-	res = setPath(&it, path_elems, path_nulls, path_len, &st, 0, NULL, false);
+	res = setPath(&it, path_elems, path_nulls, path_len, &st,
+				  0, NULL, JB_PATH_DELETE);
+
+	Assert(res != NULL);
+
+	PG_RETURN_JSONB(JsonbValueToJsonb(res));
+}
+
+/*
+ * SQL function jsonb_insert(jsonb, text[], jsonb, boolean)
+ *
+ */
+Datum
+jsonb_insert(PG_FUNCTION_ARGS)
+{
+	Jsonb	   *in = PG_GETARG_JSONB(0);
+	ArrayType  *path = PG_GETARG_ARRAYTYPE_P(1);
+	Jsonb	   *newval = PG_GETARG_JSONB(2);
+	bool		after = PG_GETARG_BOOL(3);
+	JsonbValue *res = NULL;
+	Datum	   *path_elems;
+	bool	   *path_nulls;
+	int			path_len;
+	JsonbIterator *it;
+	JsonbParseState *st = NULL;
+
+	if (ARR_NDIM(path) > 1)
+		ereport(ERROR,
+				(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 set path in scalar")));
+
+	deconstruct_array(path, TEXTOID, -1, false, 'i',
+					  &path_elems, &path_nulls, &path_len);
+
+	if (path_len == 0)
+		PG_RETURN_JSONB(in);
+
+	it = JsonbIteratorInit(&in->root);
+
+	res = setPath(&it, path_elems, path_nulls, path_len, &st, 0, newval,
+				  after ? JB_PATH_INSERT_AFTER : JB_PATH_INSERT_BEFORE);
 
 	Assert(res != NULL);
 
@@ -3707,18 +3761,23 @@ IteratorConcat(JsonbIterator **it1, JsonbIterator **it2,
 }
 
 /*
- * Do most of the heavy work for jsonb_set
+ * Do most of the heavy work for jsonb_set/jsonb_insert
+ *
+ * If JB_PATH_DELETE bit is set in op_type, the element is to be removed.
  *
- * If newval is null, the element is to be removed.
+ * If any bit mentioned in JB_PATH_CREATE_OR_INSERT is set in op_type,
+ * we create the new value if the key or array index does not exist.
  *
- * If create is true, we create the new value if the key or array index
- * does not exist. All path elements before the last must already exist
- * whether or not create is true, or nothing is done.
+ * Bits JB_PATH_INSERT_BEFORE and JB_PATH_INSERT_AFTER in op_type
+ * behave as JB_PATH_CREATE if new value is inserted in JsonbObject.
+ *
+ * All path elements before the last must already exist
+ * whatever bits in op_type are set, or nothing is done.
  */
 static JsonbValue *
 setPath(JsonbIterator **it, Datum *path_elems,
 		bool *path_nulls, int path_len,
-		JsonbParseState **st, int level, Jsonb *newval, bool create)
+		JsonbParseState **st, int level, Jsonb *newval, int op_type)
 {
 	JsonbValue	v;
 	JsonbIteratorToken r;
@@ -3739,15 +3798,15 @@ setPath(JsonbIterator **it, Datum *path_elems,
 		case WJB_BEGIN_ARRAY:
 			(void) pushJsonbValue(st, r, NULL);
 			setPathArray(it, path_elems, path_nulls, path_len, st, level,
-						 newval, v.val.array.nElems, create);
+						 newval, v.val.array.nElems, op_type);
 			r = JsonbIteratorNext(it, &v, false);
 			Assert(r == WJB_END_ARRAY);
 			res = pushJsonbValue(st, r, NULL);
 			break;
 		case WJB_BEGIN_OBJECT:
 			(void) pushJsonbValue(st, r, NULL);
 			setPathObject(it, path_elems, path_nulls, path_len, st, level,
-						  newval, v.val.object.nPairs, create);
+						  newval, v.val.object.nPairs, op_type);
 			r = JsonbIteratorNext(it, &v, true);
 			Assert(r == WJB_END_OBJECT);
 			res = pushJsonbValue(st, r, NULL);
@@ -3771,7 +3830,7 @@ setPath(JsonbIterator **it, Datum *path_elems,
 static void
 setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 			  int path_len, JsonbParseState **st, int level,
-			  Jsonb *newval, uint32 npairs, bool create)
+			  Jsonb *newval, uint32 npairs, int op_type)
 {
 	JsonbValue	v;
 	int			i;
@@ -3782,7 +3841,8 @@ setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 		done = true;
 
 	/* empty object is a special case for create */
-	if ((npairs == 0) && create && (level == path_len - 1))
+	if ((npairs == 0) && (op_type & JB_PATH_CREATE_OR_INSERT) &&
+		(level == path_len - 1))
 	{
 		JsonbValue	newkey;
 
@@ -3807,8 +3867,19 @@ setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 		{
 			if (level == path_len - 1)
 			{
-				r = JsonbIteratorNext(it, &v, true);	/* skip */
-				if (newval != NULL)
+				/*
+				 * called from jsonb_insert(), it forbids redefining
+				 * an existsing value
+				 */
+				if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_INSERT_AFTER))
+					ereport(ERROR,
+							(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+							 errmsg("cannot replace existing key"),
+							 errhint("Try using the function jsonb_set "
+									 "to replace key value.")));
+
+				r = JsonbIteratorNext(it, &v, true); /* skip value */
+				if (!(op_type & JB_PATH_DELETE))
 				{
 					(void) pushJsonbValue(st, WJB_KEY, &k);
 					addJsonbToParseState(st, newval);
@@ -3819,12 +3890,13 @@ setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 			{
 				(void) pushJsonbValue(st, r, &k);
 				setPath(it, path_elems, path_nulls, path_len,
-						st, level + 1, newval, create);
+						st, level + 1, newval, op_type);
 			}
 		}
 		else
 		{
-			if (create && !done && level == path_len - 1 && i == npairs - 1)
+			if ((op_type & JB_PATH_CREATE_OR_INSERT) && !done &&
+				level == path_len - 1 && i == npairs - 1)
 			{
 				JsonbValue	newkey;
 
@@ -3865,7 +3937,7 @@ setPathObject(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 static void
 setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 			 int path_len, JsonbParseState **st, int level,
-			 Jsonb *newval, uint32 nelems, bool create)
+			 Jsonb *newval, uint32 nelems, int op_type)
 {
 	JsonbValue	v;
 	int			idx,
@@ -3909,7 +3981,8 @@ setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 	 * what the idx value is
 	 */
 
-	if ((idx == INT_MIN || nelems == 0) && create && (level == path_len - 1))
+	if ((idx == INT_MIN || nelems == 0) && (level == path_len - 1) &&
+		(op_type & JB_PATH_CREATE_OR_INSERT))
 	{
 		Assert(newval != NULL);
 		addJsonbToParseState(st, newval);
@@ -3926,14 +3999,26 @@ setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 			if (level == path_len - 1)
 			{
 				r = JsonbIteratorNext(it, &v, true);	/* skip */
-				if (newval != NULL)
+
+				if (op_type & (JB_PATH_INSERT_BEFORE | JB_PATH_CREATE))
+					addJsonbToParseState(st, newval);
+
+				/*
+				 * We should keep current value only in case of
+				 * JB_PATH_INSERT_BEFORE or JB_PATH_INSERT_AFTER
+				 * because otherwise it should be deleted or replaced
+				 */
+				if (op_type & (JB_PATH_INSERT_AFTER | JB_PATH_INSERT_BEFORE))
+					(void) pushJsonbValue(st, r, &v);
+
+				if (op_type & JB_PATH_INSERT_AFTER)
 					addJsonbToParseState(st, newval);
 
 				done = true;
 			}
 			else
 				(void) setPath(it, path_elems, path_nulls, path_len,
-							   st, level + 1, newval, create);
+							   st, level + 1, newval, op_type);
 		}
 		else
 		{
@@ -3958,7 +4043,8 @@ setPathArray(JsonbIterator **it, Datum *path_elems, bool *path_nulls,
 				}
 			}
 
-			if (create && !done && level == path_len - 1 && i == nelems - 1)
+			if (op_type & JB_PATH_CREATE_OR_INSERT && !done &&
+				level == path_len - 1 && i == nelems - 1)
 			{
 				addJsonbToParseState(st, newval);
 			}
diff --git a/src/include/catalog/catversion.h b/src/include/catalog/catversion.h
@@ -53,6 +53,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	201604053
+#define CATALOG_VERSION_NO	201604061
 
 #endif
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
@@ -4869,6 +4869,8 @@ DATA(insert OID = 3305 (  jsonb_set    PGNSP PGUID 12 1 0 0 0 f f f f t f i s 4
 DESCR("Set part of a jsonb");
 DATA(insert OID = 3306 (  jsonb_pretty	   PGNSP PGUID 12 1 0 0 0 f f f f t f i s 1 0 25 "3802" _null_ _null_ _null_ _null_ _null_ jsonb_pretty _null_ _null_ _null_ ));
 DESCR("Indented text from jsonb");
+DATA(insert OID = 3579 (  jsonb_insert    PGNSP PGUID 12 1 0 0 0 f f f f t f i s 4 0 3802 "3802 1009 3802 16" _null_ _null_ _null_ _null_ _null_ jsonb_insert _null_ _null_ _null_ ));
+DESCR("Insert value into a jsonb");
 /* txid */
 DATA(insert OID = 2939 (  txid_snapshot_in			PGNSP PGUID 12 1  0 0 0 f f f f t f i s 1 0 2970 "2275" _null_ _null_ _null_ _null_ _null_ txid_snapshot_in _null_ _null_ _null_ ));
 DESCR("I/O");
diff --git a/src/include/utils/jsonb.h b/src/include/utils/jsonb.h
@@ -408,6 +408,9 @@ extern Datum jsonb_delete_path(PG_FUNCTION_ARGS);
 /* replacement */
 extern Datum jsonb_set(PG_FUNCTION_ARGS);
 
+/* insert after or before (for arrays) */
+extern Datum jsonb_insert(PG_FUNCTION_ARGS);
+
 /* Support functions */
 extern uint32 getJsonbOffset(const JsonbContainer *jc, int index);
 extern uint32 getJsonbLength(const JsonbContainer *jc, int index);
diff --git a/src/test/regress/expected/jsonb.out b/src/test/regress/expected/jsonb.out
diff --git a/src/test/regress/sql/jsonb.sql b/src/test/regress/sql/jsonb.sql
