Add JsonValueUnwrap(), JsonExtractScalar()

Nikita Glukhov · Nikita Glukhov · commit cc253f90b301 · 2020-11-24T18:49:53.000+03:00
diff --git a/src/backend/utils/adt/json_generic.c b/src/backend/utils/adt/json_generic.c
@@ -104,6 +104,52 @@ JsonValueCopy(JsonValue *res, const JsonValue *val)
 	return res;
 }
 
+JsonValue *
+JsonExtractScalar(JsonContainer *jc, JsonValue *scalar)
+{
+	JsonIterator   *it;
+	JsonValue		val;
+
+	Assert(JsonContainerIsScalar(jc));
+
+	it = JsonIteratorInit(jc);
+
+	if (JsonIteratorNext(&it, &val, false) != WJB_BEGIN_ARRAY ||
+		JsonIteratorNext(&it, scalar, false) != WJB_ELEM ||
+		JsonIteratorNext(&it, &val, false) != WJB_END_ARRAY)
+		elog(ERROR, "unexpected structure of scalar json container");
+
+	return scalar;
+}
+
+const JsonValue *
+JsonValueUnwrap(const JsonValue *val, JsonValue *valbuf)
+{
+	if (val->type == jbvBinary)
+	{
+		JsonContainer *jc = val->val.binary.data;
+
+		if (jc->ops == &jsonvContainerOps)
+		{
+			val = (JsonbValue *) jc->data;
+			Assert(val->type != jbvBinary);
+		}
+		else if (JsonContainerIsScalar(jc))
+		{
+			val = JsonExtractScalar(jc, valbuf);
+			Assert(IsAJsonbScalar(val));
+		}
+	}
+
+	if (val->type == jbvArray && val->val.array.rawScalar)
+	{
+		val = &val->val.array.elems[0];
+		Assert(IsAJsonbScalar(val));
+	}
+
+	return val;
+}
+
 static inline JsonValue *
 jsonFindKeyInObjectInternal(JsonContainer *obj, const char *key, int len,
 							bool last)
diff --git a/src/backend/utils/adt/jsonb.c b/src/backend/utils/adt/jsonb.c
@@ -943,15 +943,9 @@ datum_to_jsonb(Datum val, bool is_null, JsonbInState *result,
 											? DatumGetJsontP(val)
 											: DatumGetJsonbP(val);
 #endif
-					JsonbIterator *it;
-
-					it = JsonbIteratorInit(&jsonb->root);
-
 					if (JB_ROOT_IS_SCALAR(jsonb))
 					{
-						(void) JsonbIteratorNext(&it, &jb, true);
-						Assert(jb.type == jbvArray);
-						(void) JsonbIteratorNext(&it, &jb, true);
+						JsonExtractScalar(&jsonb->root, &jb);
 						break;
 					}
 
@@ -963,6 +957,7 @@ datum_to_jsonb(Datum val, bool is_null, JsonbInState *result,
 					else
 					{
 						JsonbIteratorToken type;
+						JsonbIterator *it = JsonbIteratorInit(&jsonb->root);
 
 						while ((type = JsonbIteratorNext(&it, &jb, false))
 							   != WJB_DONE)
diff --git a/src/backend/utils/adt/jsonb_util.c b/src/backend/utils/adt/jsonb_util.c
@@ -724,6 +724,16 @@ pushJsonbValueExt(JsonbParseState **pstate, JsonbIteratorToken seq,
 		jbval->type != jbvBinary || !unpackBinary)
 	{
 		/* drop through */
+		if (jbval && (seq == WJB_ELEM || seq == WJB_VALUE))
+			jbval = JsonValueUnwrap(jbval, &v);
+
+		return pushJsonbValueScalar(pstate, seq, jbval);
+	}
+
+	if (*pstate && JsonContainerIsScalar(jbval->val.binary.data))
+	{
+		jbval = JsonExtractScalar(jbval->val.binary.data, &v);
+		Assert(IsAJsonbScalar(jbval));
 		return pushJsonbValueScalar(pstate, seq, jbval);
 	}
 
diff --git a/src/include/utils/json_generic.h b/src/include/utils/json_generic.h
@@ -360,7 +360,9 @@ extern JsonValue *JsonToJsonValue(Json *json, JsonValue *jv);
 extern JsonValue *JsonValueUnpackBinary(const JsonValue *jbv);
 extern JsonContainer *JsonValueToContainer(const JsonValue *val);
 extern JsonValue *JsonValueCopy(JsonValue *res, const JsonValue *val);
+extern const JsonValue *JsonValueUnwrap(const JsonValue *val, JsonValue *buf);
 extern JsonContainer *JsonCopyFlat(JsonContainer *flatContainer);
+extern JsonValue *JsonExtractScalar(JsonContainer *jc, JsonValue *scalar);
 
 extern bool JsonbExtractScalar(JsonbContainer *jbc, JsonbValue *res);
 extern const char *JsonbTypeName(JsonbValue *jb);
