pretty-printing for INSERT statements and VALUES clause

- Refactor DisplayCommaSeparated to accept any iterable type.
- Update indented_list to work with iterators.
- Improve formatting in the Insert display implementation

Author: lovasoa

src/ast/dml.rs

@@ -581,28 +581,32 @@ impl Display for Insert {
             )?;
         }
         if !self.columns.is_empty() {
-            write!(f, "({}) ", display_comma_separated(&self.columns))?;
+            write!(f, "({})", display_comma_separated(&self.columns))?;
+            SpaceOrNewline.fmt(f)?;
         }
         if let Some(ref parts) = self.partitioned {
             if !parts.is_empty() {
-                write!(f, "PARTITION ({}) ", display_comma_separated(parts))?;
+                write!(f, "PARTITION ({})", display_comma_separated(parts))?;
+                SpaceOrNewline.fmt(f)?;
             }
         }
         if !self.after_columns.is_empty() {
-            write!(f, "({}) ", display_comma_separated(&self.after_columns))?;
+            write!(f, "({})", display_comma_separated(&self.after_columns))?;
+            SpaceOrNewline.fmt(f)?;
         }
 
         if let Some(settings) = &self.settings {
-            write!(f, "SETTINGS {} ", display_comma_separated(settings))?;
+            write!(f, "SETTINGS {}", display_comma_separated(settings))?;
+            SpaceOrNewline.fmt(f)?;
         }
 
         if let Some(source) = &self.source {
-            write!(f, "{source}")?;
+            source.fmt(f)?;
         } else if !self.assignments.is_empty() {
-            write!(f, "SET ")?;
-            write!(f, "{}", display_comma_separated(&self.assignments))?;
+            write!(f, "SET")?;
+            indented_list(f, &self.assignments)?;
         } else if let Some(format_clause) = &self.format_clause {
-            write!(f, "{format_clause}")?;
+            format_clause.fmt(f)?;
         } else if self.columns.is_empty() {
             write!(f, "DEFAULT VALUES")?;
         }
@@ -622,7 +626,9 @@ impl Display for Insert {
         }
 
         if let Some(returning) = &self.returning {
-            write!(f, " RETURNING {}", display_comma_separated(returning))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("RETURNING")?;
+            indented_list(f, returning)?;
         }
         Ok(())
     }

src/ast/mod.rs

@@ -4543,7 +4543,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Insert(insert) => write!(f, "{insert}"),
+            Statement::Insert(insert) => insert.fmt(f),
             Statement::Install {
                 extension_name: name,
             } => write!(f, "INSTALL {name}"),

src/ast/query.rs

@@ -2883,14 +2883,26 @@ pub struct Values {
 
 impl fmt::Display for Values {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "VALUES ")?;
-        let prefix = if self.explicit_row { "ROW" } else { "" };
-        let mut delim = "";
-        for row in &self.rows {
-            write!(f, "{delim}")?;
-            delim = ", ";
-            write!(f, "{prefix}({})", display_comma_separated(row))?;
+        struct DisplayRow<'a> {
+            exprs: &'a [Expr],
+            explicit_row: bool,
+        }
+        impl<'a> fmt::Display for DisplayRow<'a> {
+            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+                let comma_separated = display_comma_separated(self.exprs);
+                if self.explicit_row {
+                    write!(f, "ROW({})", comma_separated)
+                } else {
+                    write!(f, "({})", comma_separated)
+                }
+            }
         }
+        let row_display_iter = self.rows.iter().map(|row| DisplayRow {
+            exprs: row,
+            explicit_row: self.explicit_row,
+        });
+        f.write_str("VALUES")?;
+        indented_list(f, row_display_iter)?;
         Ok(())
     }
 }

src/display_utils.rs

@@ -68,12 +68,17 @@ impl Display for SpaceOrNewline {
 
 /// A value that displays a comma-separated list of values.
 /// When pretty-printed (using {:#}), it displays each value on a new line.
-pub(crate) struct DisplayCommaSeparated<'a, T: fmt::Display>(&'a [T]);
+pub(crate) struct DisplayCommaSeparated<I, T: fmt::Display>(I)
+where
+    I: IntoIterator<Item = T> + Clone;
 
-impl<T: fmt::Display> fmt::Display for DisplayCommaSeparated<'_, T> {
+impl<I, T: fmt::Display> fmt::Display for DisplayCommaSeparated<I, T>
+where
+    I: IntoIterator<Item = T> + Clone,
+{
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let mut first = true;
-        for t in self.0 {
+        for t in self.0.clone() {
             if !first {
                 f.write_char(',')?;
                 SpaceOrNewline.fmt(f)?;
@@ -86,9 +91,12 @@ impl<T: fmt::Display> fmt::Display for DisplayCommaSeparated<'_, T> {
 }
 
 /// Displays a whitespace, followed by a comma-separated list that is indented when pretty-printed.
-pub(crate) fn indented_list<T: fmt::Display>(f: &mut fmt::Formatter, slice: &[T]) -> fmt::Result {
+pub(crate) fn indented_list<I, T: fmt::Display>(f: &mut fmt::Formatter, iter: I) -> fmt::Result
+where
+    I: IntoIterator<Item = T> + Clone,
+{
     SpaceOrNewline.fmt(f)?;
-    Indent(DisplayCommaSeparated(slice)).fmt(f)
+    Indent(DisplayCommaSeparated(iter)).fmt(f)
 }
 
 #[cfg(test)]

tests/pretty_print.rs

@@ -171,7 +171,6 @@ string' AS str
 }
 
 #[test]
-#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
 fn test_pretty_print_insert_values() {
     assert_eq!(
         prettify("INSERT INTO my_table (a, b, c) VALUES (1, 2, 3), (4, 5, 6)"),
@@ -186,17 +185,18 @@ VALUES
 }
 
 #[test]
-#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
 fn test_pretty_print_insert_select() {
     assert_eq!(
-        prettify("INSERT INTO my_table (a, b) SELECT x, y FROM source_table"),
+        prettify("INSERT INTO my_table (a, b) SELECT x, y FROM source_table RETURNING a AS id"),
         r#"
 INSERT INTO my_table (a, b)
 SELECT
   x,
   y
 FROM
   source_table
+RETURNING
+  a AS id
 "#
         .trim()
     );