feat(expr): add jsonb_pretty function (#13050)

Signed-off-by: Runji Wang <[email protected]>

proto/expr.proto

@@ -220,6 +220,7 @@ message ExprNode {
     JSONB_ARRAY_LENGTH = 603;
     IS_JSON = 604;
     JSONB_CAT = 605;
+    JSONB_PRETTY = 607;
 
     // Non-pure functions below (> 1000)
     // ------------------------

src/common/src/types/jsonb.rs

@@ -107,19 +107,6 @@ impl Ord for JsonbRef<'_> {
 
 impl crate::types::to_text::ToText for JsonbRef<'_> {
     fn write<W: std::fmt::Write>(&self, f: &mut W) -> std::fmt::Result {
-        struct FmtToIoUnchecked<F>(F);
-        impl<F: std::fmt::Write> std::io::Write for FmtToIoUnchecked<F> {
-            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-                let s = unsafe { std::str::from_utf8_unchecked(buf) };
-                self.0.write_str(s).map_err(|_| std::io::ErrorKind::Other)?;
-                Ok(buf.len())
-            }
-
-            fn flush(&mut self) -> std::io::Result<()> {
-                Ok(())
-            }
-        }
-
         // Use custom [`ToTextFormatter`] to serialize. If we are okay with the default, this can be
         // just `write!(f, "{}", self.0)`
         use serde::Serialize as _;
@@ -412,6 +399,16 @@ impl<'a> JsonbRef<'a> {
             .ok_or_else(|| format!("cannot deconstruct a jsonb {}", self.type_name()))?;
         Ok(object.iter().map(|(k, v)| (k, Self(v))))
     }
+
+    /// Pretty print the jsonb value to the given writer, with 4 spaces indentation.
+    pub fn pretty(self, f: &mut impl std::fmt::Write) -> std::fmt::Result {
+        use serde::Serialize;
+        use serde_json::ser::{PrettyFormatter, Serializer};
+
+        let mut ser =
+            Serializer::with_formatter(FmtToIoUnchecked(f), PrettyFormatter::with_indent(b"    "));
+        self.0.serialize(&mut ser).map_err(|_| std::fmt::Error)
+    }
 }
 
 /// A custom implementation for [`serde_json::ser::Formatter`] to match PostgreSQL, which adds extra
@@ -448,3 +445,18 @@ impl serde_json::ser::Formatter for ToTextFormatter {
         writer.write_all(b": ")
     }
 }
+
+/// A wrapper of [`std::fmt::Write`] to implement [`std::io::Write`].
+struct FmtToIoUnchecked<F>(F);
+
+impl<F: std::fmt::Write> std::io::Write for FmtToIoUnchecked<F> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let s = unsafe { std::str::from_utf8_unchecked(buf) };
+        self.0.write_str(s).map_err(|_| std::io::ErrorKind::Other)?;
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}

src/expr/impl/src/scalar/jsonb_info.rs

@@ -52,3 +52,24 @@ pub fn is_json_type(s: &str, t: &str) -> bool {
         }
     })
 }
+
+/// Converts the given JSON value to pretty-printed, indented text.
+///
+/// # Examples
+// TODO: enable docslt after sqllogictest supports multiline output
+/// ```text
+/// query T
+/// select jsonb_pretty('[{"f1":1,"f2":null}, 2]');
+/// ----
+/// [
+///     {
+///         "f1": 1,
+///         "f2": null
+///     },
+///     2
+/// ]
+/// ```
+#[function("jsonb_pretty(jsonb) -> varchar")]
+pub fn jsonb_pretty(v: JsonbRef<'_>, writer: &mut impl Write) {
+    v.pretty(writer).unwrap()
+}

src/frontend/src/binder/expr/function.rs

@@ -877,6 +877,7 @@ impl Binder {
                 ("jsonb_array_element_text", raw_call(ExprType::JsonbAccessStr)),
                 ("jsonb_typeof", raw_call(ExprType::JsonbTypeof)),
                 ("jsonb_array_length", raw_call(ExprType::JsonbArrayLength)),
+                ("jsonb_pretty", raw_call(ExprType::JsonbPretty)),
                 // Functions that return a constant value
                 ("pi", pi()),
                 // greatest and least

src/frontend/src/expr/pure.rs

@@ -177,6 +177,7 @@ impl ExprVisitor for ImpureAnalyzer {
             | expr_node::Type::JsonbAccessStr
             | expr_node::Type::JsonbTypeof
             | expr_node::Type::JsonbArrayLength
+            | expr_node::Type::JsonbPretty
             | expr_node::Type::IsJson
             | expr_node::Type::Sind
             | expr_node::Type::Cosd

src/risedevtool/src/bin/risedev-docslt.rs

@@ -53,11 +53,12 @@ fn extract_slt(filepath: &Path) -> Vec<SltBlock> {
             if !(line.starts_with("///") || line.starts_with("//!")) {
                 panic!("expect /// or //! at {}:{}", filepath.display(), i + 1);
             }
-            line = line[3..].trim();
-            if line == "```" {
+            line = &line[3..];
+            if line.trim() == "```" {
                 break;
             }
-            content += line;
+            // strip one leading space
+            content += line.strip_prefix(' ').unwrap_or(line);
             content += "\n";
         }
         blocks.push(SltBlock {

src/tests/regress/data/sql/jsonb.sql

@@ -1047,9 +1047,9 @@ SELECT '["a","b","c",[1,2],null]'::jsonb -> -6;
 --@ select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }');
 
 
---@ select jsonb_pretty('{"a": "test", "b": [1, 2, 3], "c": "test3", "d":{"dd": "test4", "dd2":{"ddd": "test5"}}}');
---@ select jsonb_pretty('[{"f1":1,"f2":null},2,null,[[{"x":true},6,7],8],3]');
---@ select jsonb_pretty('{"a":["b", "c"], "d": {"e":"f"}}');
+select jsonb_pretty('{"a": "test", "b": [1, 2, 3], "c": "test3", "d":{"dd": "test4", "dd2":{"ddd": "test5"}}}');
+select jsonb_pretty('[{"f1":1,"f2":null},2,null,[[{"x":true},6,7],8],3]');
+select jsonb_pretty('{"a":["b", "c"], "d": {"e":"f"}}');
 --@ 
 --@ select jsonb_concat('{"d": "test", "a": [1, 2]}', '{"g": "test2", "c": {"c1":1, "c2":2}}');
 --@