pretty print improvements (#1851)

README.md

@@ -89,10 +89,14 @@ keywords, the following should hold true for all SQL:
 
 ```rust
 // Parse SQL
+let sql = "SELECT 'hello'";
 let ast = Parser::parse_sql(&GenericDialect, sql).unwrap();
 
 // The original SQL text can be generated from the AST
 assert_eq!(ast[0].to_string(), sql);
+
+// The SQL can also be pretty-printed with newlines and indentation
+assert_eq!(format!("{:#}", ast[0]), "SELECT\n  'hello'");
 ```
 
 There are still some cases in this crate where different SQL with seemingly

src/ast/dml.rs

@@ -29,6 +29,8 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
+use crate::display_utils::{indented_list, Indent, SpaceOrNewline};
+
 pub use super::ddl::{ColumnDef, TableConstraint};
 
 use super::{
@@ -579,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, "SET")?;
-            write!(f, "{}", display_comma_separated(&self.assignments))?;
+            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")?;
         }
@@ -620,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(())
     }
@@ -649,32 +657,45 @@ pub struct Delete {
 
 impl Display for Delete {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "DELETE ")?;
+        f.write_str("DELETE")?;
         if !self.tables.is_empty() {
-            write!(f, "{} ", display_comma_separated(&self.tables))?;
+            indented_list(f, &self.tables)?;
         }
         match &self.from {
             FromTable::WithFromKeyword(from) => {
-                write!(f, "FROM {}", display_comma_separated(from))?;
+                f.write_str(" FROM")?;
+                indented_list(f, from)?;
             }
             FromTable::WithoutKeyword(from) => {
-                write!(f, "{}", display_comma_separated(from))?;
+                indented_list(f, from)?;
             }
         }
         if let Some(using) = &self.using {
-            write!(f, " USING {}", display_comma_separated(using))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("USING")?;
+            indented_list(f, using)?;
         }
         if let Some(selection) = &self.selection {
-            write!(f, " WHERE {selection}")?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("WHERE")?;
+            SpaceOrNewline.fmt(f)?;
+            Indent(selection).fmt(f)?;
         }
         if let Some(returning) = &self.returning {
-            write!(f, " RETURNING {}", display_comma_separated(returning))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("RETURNING")?;
+            indented_list(f, returning)?;
         }
         if !self.order_by.is_empty() {
-            write!(f, " ORDER BY {}", display_comma_separated(&self.order_by))?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("ORDER BY")?;
+            indented_list(f, &self.order_by)?;
         }
         if let Some(limit) = &self.limit {
-            write!(f, " LIMIT {limit}")?;
+            SpaceOrNewline.fmt(f)?;
+            f.write_str("LIMIT")?;
+            SpaceOrNewline.fmt(f)?;
+            Indent(limit).fmt(f)?;
         }
         Ok(())
     }

src/ast/mod.rs

@@ -41,7 +41,7 @@ use serde::{Deserialize, Serialize};
 use sqlparser_derive::{Visit, VisitMut};
 
 use crate::{
-    display_utils::SpaceOrNewline,
+    display_utils::{indented_list, SpaceOrNewline},
     tokenizer::{Span, Token},
 };
 use crate::{
@@ -4548,7 +4548,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}"),
@@ -4611,30 +4611,42 @@ impl fmt::Display for Statement {
                 returning,
                 or,
             } => {
-                write!(f, "UPDATE ")?;
+                f.write_str("UPDATE ")?;
                 if let Some(or) = or {
-                    write!(f, "{or} ")?;
+                    or.fmt(f)?;
+                    f.write_str(" ")?;
                 }
-                write!(f, "{table}")?;
+                table.fmt(f)?;
                 if let Some(UpdateTableFromKind::BeforeSet(from)) = from {
-                    write!(f, " FROM {}", display_comma_separated(from))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("FROM")?;
+                    indented_list(f, from)?;
                 }
                 if !assignments.is_empty() {
-                    write!(f, " SET {}", display_comma_separated(assignments))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("SET")?;
+                    indented_list(f, assignments)?;
                 }
                 if let Some(UpdateTableFromKind::AfterSet(from)) = from {
-                    write!(f, " FROM {}", display_comma_separated(from))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("FROM")?;
+                    indented_list(f, from)?;
                 }
                 if let Some(selection) = selection {
-                    write!(f, " WHERE {selection}")?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("WHERE")?;
+                    SpaceOrNewline.fmt(f)?;
+                    Indent(selection).fmt(f)?;
                 }
                 if let Some(returning) = returning {
-                    write!(f, " RETURNING {}", display_comma_separated(returning))?;
+                    SpaceOrNewline.fmt(f)?;
+                    f.write_str("RETURNING")?;
+                    indented_list(f, returning)?;
                 }
                 Ok(())
             }
-            Statement::Delete(delete) => write!(f, "{delete}"),
+            Statement::Delete(delete) => delete.fmt(f),
-            Statement::Open(open) => write!(f, "{open}"),
+            Statement::Open(open) => open.fmt(f),
             Statement::Close { cursor } => {
                 write!(f, "CLOSE {cursor}")?;
 

src/ast/query.rs

@@ -2888,13 +2888,14 @@ pub struct Values {
 
 impl fmt::Display for Values {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "VALUES ")?;
+        f.write_str("VALUES")?;
         let prefix = if self.explicit_row { "ROW" } else { "" };
         let mut delim = "";
         for row in &self.rows {
-            write!(f, "{delim}")?;
+            f.write_str(delim)?;
-            delim = ", ";
+            delim = ",";
-            write!(f, "{prefix}({})", display_comma_separated(row))?;
+            SpaceOrNewline.fmt(f)?;
+            Indent(format_args!("{prefix}({})", display_comma_separated(row))).fmt(f)?;
         }
         Ok(())
     }

src/display_utils.rs

@@ -31,13 +31,10 @@ where
     T: Write,
 {
     fn write_str(&mut self, s: &str) -> fmt::Result {
-        let mut first = true;
+        self.0.write_str(s)?;
-        for line in s.split('\n') {
+        // Our NewLine and SpaceOrNewline utils always print individual newlines as a single-character string.
-            if !first {
+        if s == "\n" {
-                write!(self.0, "\n{INDENT}")?;
+            self.0.write_str(INDENT)?;
-            }
-            self.0.write_str(line)?;
-            first = false;
         }
         Ok(())
     }
@@ -71,7 +68,7 @@ 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 struct DisplayCommaSeparated<'a, T: fmt::Display>(&'a [T]);
+pub(crate) struct DisplayCommaSeparated<'a, T: fmt::Display>(&'a [T]);
 
 impl<T: fmt::Display> fmt::Display for DisplayCommaSeparated<'_, T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -89,45 +86,33 @@ 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<T: fmt::Display>(f: &mut fmt::Formatter, items: &[T]) -> fmt::Result {
     SpaceOrNewline.fmt(f)?;
-    Indent(DisplayCommaSeparated(slice)).fmt(f)
+    Indent(DisplayCommaSeparated(items)).fmt(f)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    struct DisplayCharByChar<T: Display>(T);
-
-    impl<T: Display> Display for DisplayCharByChar<T> {
-        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-            for c in self.0.to_string().chars() {
-                write!(f, "{}", c)?;
-            }
-            Ok(())
-        }
-    }
-
     #[test]
     fn test_indent() {
-        let original = "line 1\nline 2";
+        struct TwoLines;
-        let indent = Indent(original);
+
+        impl Display for TwoLines {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                f.write_str("line 1")?;
+                SpaceOrNewline.fmt(f)?;
+                f.write_str("line 2")
+            }
+        }
+
+        let indent = Indent(TwoLines);
         assert_eq!(
             indent.to_string(),
-            original,
+            TwoLines.to_string(),
             "Only the alternate form should be indented"
         );
-        let expected = "  line 1\n  line 2";
+        assert_eq!(format!("{:#}", indent), "  line 1\n  line 2");
-        assert_eq!(format!("{:#}", indent), expected);
-        let display_char_by_char = DisplayCharByChar(original);
-        assert_eq!(format!("{:#}", Indent(display_char_by_char)), expected);
-    }
-
-    #[test]
-    fn test_space_or_newline() {
-        let space_or_newline = SpaceOrNewline;
-        assert_eq!(format!("{}", space_or_newline), " ");
-        assert_eq!(format!("{:#}", space_or_newline), "\n");
     }
 }

tests/pretty_print.rs

@@ -155,3 +155,245 @@ FROM
 "#.trim()
     );
 }
+
+#[test]
+fn test_pretty_print_multiline_string() {
+    assert_eq!(
+        prettify("SELECT 'multiline\nstring' AS str"),
+        r#"
+SELECT
+  'multiline
+string' AS str
+"#
+        .trim(),
+        "A  literal string with a newline should be kept as is. The contents of the string should not be indented."
+    );
+}
+
+#[test]
+fn test_pretty_print_insert_values() {
+    assert_eq!(
+        prettify("INSERT INTO my_table (a, b, c) VALUES (1, 2, 3), (4, 5, 6)"),
+        r#"
+INSERT INTO my_table (a, b, c)
+VALUES
+  (1, 2, 3),
+  (4, 5, 6)
+"#
+        .trim()
+    );
+}
+
+#[test]
+fn test_pretty_print_insert_select() {
+    assert_eq!(
+        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()
+    );
+}
+
+#[test]
+fn test_pretty_print_update() {
+    assert_eq!(
+        prettify("UPDATE my_table SET a = 1, b = 2 WHERE x > 0 RETURNING id, name"),
+        r#"
+UPDATE my_table
+SET
+  a = 1,
+  b = 2
+WHERE
+  x > 0
+RETURNING
+  id,
+  name
+"#
+        .trim()
+    );
+}
+
+#[test]
+fn test_pretty_print_delete() {
+    assert_eq!(
+        prettify("DELETE FROM my_table WHERE x > 0 RETURNING id, name"),
+        r#"
+DELETE FROM
+  my_table
+WHERE
+  x > 0
+RETURNING
+  id,
+  name
+"#
+        .trim()
+    );
+
+    assert_eq!(
+        prettify("DELETE table1, table2"),
+        r#"
+DELETE
+  table1,
+  table2
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_create_table() {
+    assert_eq!(
+        prettify("CREATE TABLE my_table (id INT PRIMARY KEY, name VARCHAR(255) NOT NULL, CONSTRAINT fk_other FOREIGN KEY (id) REFERENCES other_table(id))"),
+        r#"
+CREATE TABLE my_table (
+  id INT PRIMARY KEY,
+  name VARCHAR(255) NOT NULL,
+  CONSTRAINT fk_other FOREIGN KEY (id) REFERENCES other_table(id)
+)
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_create_view() {
+    assert_eq!(
+        prettify("CREATE VIEW my_view AS SELECT a, b FROM my_table WHERE x > 0"),
+        r#"
+CREATE VIEW my_view AS
+SELECT
+  a,
+  b
+FROM
+  my_table
+WHERE
+  x > 0
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_create_function() {
+    assert_eq!(
+        prettify("CREATE FUNCTION my_func() RETURNS INT BEGIN SELECT COUNT(*) INTO @count FROM my_table; RETURN @count; END"),
+        r#"
+CREATE FUNCTION my_func() RETURNS INT
+BEGIN
+  SELECT COUNT(*) INTO @count FROM my_table;
+  RETURN @count;
+END
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_json_table() {
+    assert_eq!(
+        prettify("SELECT * FROM JSON_TABLE(@json, '$[*]' COLUMNS (id INT PATH '$.id', name VARCHAR(255) PATH '$.name')) AS jt"),
+        r#"
+SELECT
+  *
+FROM
+  JSON_TABLE(
+    @json,
+    '$[*]' COLUMNS (
+      id INT PATH '$.id',
+      name VARCHAR(255) PATH '$.name'
+    )
+  ) AS jt
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_transaction_blocks() {
+    assert_eq!(
+        prettify("BEGIN; UPDATE my_table SET x = 1; COMMIT;"),
+        r#"
+BEGIN;
+UPDATE my_table SET x = 1;
+COMMIT;
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_control_flow() {
+    assert_eq!(
+        prettify("IF x > 0 THEN SELECT 'positive'; ELSE SELECT 'negative'; END IF;"),
+        r#"
+IF x > 0 THEN
+  SELECT 'positive';
+ELSE
+  SELECT 'negative';
+END IF;
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_merge() {
+    assert_eq!(
+        prettify("MERGE INTO target_table t USING source_table s ON t.id = s.id WHEN MATCHED THEN UPDATE SET t.value = s.value WHEN NOT MATCHED THEN INSERT (id, value) VALUES (s.id, s.value)"),
+        r#"
+MERGE INTO target_table t
+USING source_table s ON t.id = s.id
+WHEN MATCHED THEN
+  UPDATE SET t.value = s.value
+WHEN NOT MATCHED THEN
+  INSERT (id, value) VALUES (s.id, s.value)
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_create_index() {
+    assert_eq!(
+        prettify("CREATE INDEX idx_name ON my_table (column1, column2)"),
+        r#"
+CREATE INDEX idx_name
+ON my_table (column1, column2)
+"#
+        .trim()
+    );
+}
+
+#[test]
+#[ignore = "https://github.com/apache/datafusion-sqlparser-rs/issues/1850"]
+fn test_pretty_print_explain() {
+    assert_eq!(
+        prettify("EXPLAIN ANALYZE SELECT * FROM my_table WHERE x > 0"),
+        r#"
+EXPLAIN ANALYZE
+SELECT
+  *
+FROM
+  my_table
+WHERE
+  x > 0
+"#
+        .trim()
+    );
+}