Add json[b]_extract_keys()

Nikita Glukhov · Nikita Glukhov · commit 5085a240b3ad · 2021-03-29T20:40:11.000+03:00
diff --git a/src/backend/utils/adt/json.c b/src/backend/utils/adt/json.c
@@ -57,6 +57,7 @@
 #define jsonb_object_field			json_object_field
 #define jsonb_object_field_text		json_object_field_text
 #define jsonb_object_keys			json_object_keys
+#define jsonb_extract_keys			json_extract_keys
 #define jsonb_populate_record		json_populate_record
 #define jsonb_populate_recordset	json_populate_recordset
 #define jsonb_pretty				json_pretty
diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
@@ -577,8 +577,8 @@ makeJsonLexContext(text *json, bool need_escapes)
  * limited in size to NAMEDATALEN and the number of keys is unlikely to
  * be so huge that it has major memory implications.
  */
-Datum
-jsonb_object_keys(PG_FUNCTION_ARGS)
+static Datum
+jsonb_extract_keys_internal(FunctionCallInfo fcinfo, bool outermost)
 {
 	FuncCallContext *funcctx;
 	OkeysState *state;
@@ -592,18 +592,25 @@ jsonb_object_keys(PG_FUNCTION_ARGS)
 		JsonbValue	v;
 		JsonbIteratorToken r;
 
-		if (JB_ROOT_IS_SCALAR(jb))
-			ereport(ERROR,
-					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-					 errmsg("cannot call %s on a scalar",
-							JSONB"_object_keys")));
-		else if (JB_ROOT_IS_ARRAY(jb))
-			ereport(ERROR,
-					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-					 errmsg("cannot call %s on an array",
-							JSONB"_object_keys")));
+		if (outermost)
+		{
+			if (JB_ROOT_IS_SCALAR(jb))
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						 errmsg("cannot call %s on a scalar",
+								JSONB"_object_keys")));
+			else if (JB_ROOT_IS_ARRAY(jb))
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						 errmsg("cannot call %s on an array",
+								JSONB"_object_keys")));
+		}
 
 		funcctx = SRF_FIRSTCALL_INIT();
+
+		if (!outermost && JB_ROOT_IS_SCALAR(jb))
+			SRF_RETURN_DONE(funcctx);
+
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
 		state = palloc(sizeof(OkeysState));
@@ -619,7 +626,7 @@ jsonb_object_keys(PG_FUNCTION_ARGS)
 
 		while ((r = JsonbIteratorNext(&it, &v, skipNested)) != WJB_DONE)
 		{
-			skipNested = true;
+			skipNested = outermost;
 
 			if (r == WJB_KEY)
 			{
@@ -655,6 +662,18 @@ jsonb_object_keys(PG_FUNCTION_ARGS)
 	SRF_RETURN_DONE(funcctx);
 }
 
+Datum
+jsonb_object_keys(PG_FUNCTION_ARGS)
+{
+	return jsonb_extract_keys_internal(fcinfo, true);
+}
+
+Datum
+jsonb_extract_keys(PG_FUNCTION_ARGS)
+{
+	return jsonb_extract_keys_internal(fcinfo, false);
+}
+
 #ifndef JSON_C
 /*
  * Report a JSON error.
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
@@ -8707,6 +8707,9 @@
 { oid => '3957', descr => 'get json object keys',
   proname => 'json_object_keys', prorows => '100', proretset => 't',
   prorettype => 'text', proargtypes => 'json', prosrc => 'json_object_keys' },
+{ oid => '8162', descr => 'extract all json object keys',
+  proname => 'json_extract_keys', prorows => '100', proretset => 't',
+  prorettype => 'text', proargtypes => 'json', prosrc => 'json_extract_keys' },
 { oid => '3958', descr => 'key value pairs of a json object',
   proname => 'json_each', prorows => '100', proretset => 't',
   prorettype => 'record', proargtypes => 'json',
@@ -9574,6 +9577,9 @@
 { oid => '3931', descr => 'get jsonb object keys',
   proname => 'jsonb_object_keys', prorows => '100', proretset => 't',
   prorettype => 'text', proargtypes => 'jsonb', prosrc => 'jsonb_object_keys' },
+{ oid => '8161', descr => 'extract all jsonb object keys',
+  proname => 'jsonb_extract_keys', prorows => '100', proretset => 't',
+  prorettype => 'text', proargtypes => 'jsonb', prosrc => 'jsonb_extract_keys' },
 { oid => '3208', descr => 'key value pairs of a jsonb object',
   proname => 'jsonb_each', prorows => '100', proretset => 't',
   prorettype => 'record', proargtypes => 'jsonb',
