Repository: apache/datafusion-sqlparser-rs (rust)

Issue:
Add `Display` implementation for other `DML` operations
Much like `Select` and `CreateTable`, the ability to programmatically construct queries from AST parts is incredibly useful.

However, `Insert`, `Delete`, and `Alter` do not implement Display, making it difficult to produce these statements as strings.  


Pull request description:
feat: adding Display implementation to DELETE and INSERT
Closes #1420 .

Patch:
diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
index 49f6ae4..21a716d 100644
--- a/src/ast/ddl.rs
+++ b/src/ast/ddl.rs
@@ -1706,6 +1706,8 @@ pub enum UserDefinedTypeRepresentation {
     Composite {
         attributes: Vec<UserDefinedTypeCompositeAttributeDef>,
     },
+    /// Note: this is PostgreSQL-specific. See <https://www.postgresql.org/docs/current/sql-createtype.html>
+    Enum { labels: Vec<Ident> },
 }
 
 impl fmt::Display for UserDefinedTypeRepresentation {
@@ -1714,6 +1716,9 @@ impl fmt::Display for UserDefinedTypeRepresentation {
             UserDefinedTypeRepresentation::Composite { attributes } => {
                 write!(f, "({})", display_comma_separated(attributes))
             }
+            UserDefinedTypeRepresentation::Enum { labels } => {
+                write!(f, "ENUM ({})", display_comma_separated(labels))
+            }
         }
     }
 }
diff --git a/src/ast/dml.rs b/src/ast/dml.rs
index 8121f2c..2932faf 100644
--- a/src/ast/dml.rs
+++ b/src/ast/dml.rs
@@ -16,7 +16,12 @@
 // under the License.
 
 #[cfg(not(feature = "std"))]
-use alloc::{boxed::Box, string::String, vec::Vec};
+use alloc::{
+    boxed::Box,
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 
 use core::fmt::{self, Display};
 #[cfg(feature = "serde")]
@@ -492,6 +497,81 @@ pub struct Insert {
     pub insert_alias: Option<InsertAliases>,
 }
 
+impl Display for Insert {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let table_name = if let Some(alias) = &self.table_alias {
+            format!("{0} AS {alias}", self.table_name)
+        } else {
+            self.table_name.to_string()
+        };
+
+        if let Some(action) = self.or {
+            write!(f, "INSERT OR {action} INTO {table_name} ")?;
+        } else {
+            write!(
+                f,
+                "{start}",
+                start = if self.replace_into {
+                    "REPLACE"
+                } else {
+                    "INSERT"
+                },
+            )?;
+            if let Some(priority) = self.priority {
+                write!(f, " {priority}",)?;
+            }
+
+            write!(
+                f,
+                "{ignore}{over}{int}{tbl} {table_name} ",
+                table_name = table_name,
+                ignore = if self.ignore { " IGNORE" } else { "" },
+                over = if self.overwrite { " OVERWRITE" } else { "" },
+                int = if self.into { " INTO" } else { "" },
+                tbl = if self.table { " TABLE" } else { "" },
+            )?;
+        }
+        if !self.columns.is_empty() {
+            write!(f, "({}) ", display_comma_separated(&self.columns))?;
+        }
+        if let Some(ref parts) = self.partitioned {
+            if !parts.is_empty() {
+                write!(f, "PARTITION ({}) ", display_comma_separated(parts))?;
+            }
+        }
+        if !self.after_columns.is_empty() {
+            write!(f, "({}) ", display_comma_separated(&self.after_columns))?;
+        }
+
+        if let Some(source) = &self.source {
+            write!(f, "{source}")?;
+        }
+
+        if self.source.is_none() && self.columns.is_empty() {
+            write!(f, "DEFAULT VALUES")?;
+        }
+
+        if let Some(insert_alias) = &self.insert_alias {
+            write!(f, " AS {0}", insert_alias.row_alias)?;
+
+            if let Some(col_aliases) = &insert_alias.col_aliases {
+                if !col_aliases.is_empty() {
+                    write!(f, " ({})", display_comma_separated(col_aliases))?;
+                }
+            }
+        }
+
+        if let Some(on) = &self.on {
+            write!(f, "{on}")?;
+        }
+
+        if let Some(returning) = &self.returning {
+            write!(f, " RETURNING {}", display_comma_separated(returning))?;
+        }
+        Ok(())
+    }
+}
+
 /// DELETE statement.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -512,3 +592,36 @@ pub struct Delete {
     /// LIMIT (MySQL)
     pub limit: Option<Expr>,
 }
+
+impl Display for Delete {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "DELETE ")?;
+        if !self.tables.is_empty() {
+            write!(f, "{} ", display_comma_separated(&self.tables))?;
+        }
+        match &self.from {
+            FromTable::WithFromKeyword(from) => {
+                write!(f, "FROM {}", display_comma_separated(from))?;
+            }
+            FromTable::WithoutKeyword(from) => {
+                write!(f, "{}", display_comma_separated(from))?;
+            }
+        }
+        if let Some(using) = &self.using {
+            write!(f, " USING {}", display_comma_separated(using))?;
+        }
+        if let Some(selection) = &self.selection {
+            write!(f, " WHERE {selection}")?;
+        }
+        if let Some(returning) = &self.returning {
+            write!(f, " RETURNING {}", display_comma_separated(returning))?;
+        }
+        if !self.order_by.is_empty() {
+            write!(f, " ORDER BY {}", display_comma_separated(&self.order_by))?;
+        }
+        if let Some(limit) = &self.limit {
+            write!(f, " LIMIT {limit}")?;
+        }
+        Ok(())
+    }
+}
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 4900307..790a39b 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -2176,6 +2176,18 @@ pub enum FromTable {
     /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#delete_statement>
     WithoutKeyword(Vec<TableWithJoins>),
 }
+impl Display for FromTable {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            FromTable::WithFromKeyword(tables) => {
+                write!(f, "FROM {}", display_comma_separated(tables))
+            }
+            FromTable::WithoutKeyword(tables) => {
+                write!(f, "{}", display_comma_separated(tables))
+            }
+        }
+    }
+}
 
 /// Policy type for a `CREATE POLICY` statement.
 /// ```sql
@@ -3533,93 +3545,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Insert(insert) => {
-                let Insert {
-                    or,
-                    ignore,
-                    into,
-                    table_name,
-                    table_alias,
-                    overwrite,
-                    partitioned,
-                    columns,
-                    after_columns,
-                    source,
-                    table,
-                    on,
-                    returning,
-                    replace_into,
-                    priority,
-                    insert_alias,
-                } = insert;
-                let table_name = if let Some(alias) = table_alias {
-                    format!("{table_name} AS {alias}")
-                } else {
-                    table_name.to_string()
-                };
-
-                if let Some(action) = or {
-                    write!(f, "INSERT OR {action} INTO {table_name} ")?;
-                } else {
-                    write!(
-                        f,
-                        "{start}",
-                        start = if *replace_into { "REPLACE" } else { "INSERT" },
-                    )?;
-                    if let Some(priority) = priority {
-                        write!(f, " {priority}",)?;
-                    }
-
-                    write!(
-                        f,
-                        "{ignore}{over}{int}{tbl} {table_name} ",
-                        table_name = table_name,
-                        ignore = if *ignore { " IGNORE" } else { "" },
-                        over = if *overwrite { " OVERWRITE" } else { "" },
-                        int = if *into { " INTO" } else { "" },
-                        tbl = if *table { " TABLE" } else { "" },
-                    )?;
-                }
-                if !columns.is_empty() {
-                    write!(f, "({}) ", display_comma_separated(columns))?;
-                }
-                if let Some(ref parts) = partitioned {
-                    if !parts.is_empty() {
-                        write!(f, "PARTITION ({}) ", display_comma_separated(parts))?;
-                    }
-                }
-                if !after_columns.is_empty() {
-                    write!(f, "({}) ", display_comma_separated(after_columns))?;
-                }
-
-                if let Some(source) = source {
-                    write!(f, "{source}")?;
-                }
-
-                if source.is_none() && columns.is_empty() {
-                    write!(f, "DEFAULT VALUES")?;
-                }
-
-                if let Some(insert_alias) = insert_alias {
-                    write!(f, " AS {0}", insert_alias.row_alias)?;
-
-                    if let Some(col_aliases) = &insert_alias.col_aliases {
-                        if !col_aliases.is_empty() {
-                            write!(f, " ({})", display_comma_separated(col_aliases))?;
-                        }
-                    }
-                }
-
-                if let Some(on) = on {
-                    write!(f, "{on}")?;
-                }
-
-                if let Some(returning) = returning {
-                    write!(f, " RETURNING {}", display_comma_separated(returning))?;
-                }
-
-                Ok(())
-            }
+            Statement::Insert(insert) => write!(f, "{insert}"),
             Statement::Install {
                 extension_name: name,
             } => write!(f, "INSTALL {name}"),
@@ -3696,45 +3622,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Delete(delete) => {
-                let Delete {
-                    tables,
-                    from,
-                    using,
-                    selection,
-                    returning,
-                    order_by,
-                    limit,
-                } = delete;
-                write!(f, "DELETE ")?;
-                if !tables.is_empty() {
-                    write!(f, "{} ", display_comma_separated(tables))?;
-                }
-                match from {
-                    FromTable::WithFromKeyword(from) => {
-                        write!(f, "FROM {}", display_comma_separated(from))?;
-                    }
-                    FromTable::WithoutKeyword(from) => {
-                        write!(f, "{}", display_comma_separated(from))?;
-                    }
-                }
-                if let Some(using) = using {
-                    write!(f, " USING {}", display_comma_separated(using))?;
-                }
-                if let Some(selection) = selection {
-                    write!(f, " WHERE {selection}")?;
-                }
-                if let Some(returning) = returning {
-                    write!(f, " RETURNING {}", display_comma_separated(returning))?;
-                }
-                if !order_by.is_empty() {
-                    write!(f, " ORDER BY {}", display_comma_separated(order_by))?;
-                }
-                if let Some(limit) = limit {
-                    write!(f, " LIMIT {limit}")?;
-                }
-                Ok(())
-            }
+            Statement::Delete(delete) => write!(f, "{delete}"),
             Statement::Close { cursor } => {
                 write!(f, "CLOSE {cursor}")?;
 
diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs
index 2a66705..51dc498 100644
--- a/src/dialect/postgresql.rs
+++ b/src/dialect/postgresql.rs
@@ -28,7 +28,7 @@
 // limitations under the License.
 use log::debug;
 
-use crate::ast::{CommentObject, Statement};
+use crate::ast::{CommentObject, ObjectName, Statement, UserDefinedTypeRepresentation};
 use crate::dialect::{Dialect, Precedence};
 use crate::keywords::Keyword;
 use crate::parser::{Parser, ParserError};
@@ -138,6 +138,9 @@ impl Dialect for PostgreSqlDialect {
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
         if parser.parse_keyword(Keyword::COMMENT) {
             Some(parse_comment(parser))
+        } else if parser.parse_keyword(Keyword::CREATE) {
+            parser.prev_token(); // unconsume the CREATE in case we don't end up parsing anything
+            parse_create(parser)
         } else {
             None
         }
@@ -225,3 +228,33 @@ pub fn parse_comment(parser: &mut Parser) -> Result<Statement, ParserError> {
         if_exists,
     })
 }
+
+pub fn parse_create(parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
+    let name = parser.maybe_parse(|parser| -> Result<ObjectName, ParserError> {
+        parser.expect_keyword(Keyword::CREATE)?;
+        parser.expect_keyword(Keyword::TYPE)?;
+        let name = parser.parse_object_name(false)?;
+        parser.expect_keyword(Keyword::AS)?;
+        parser.expect_keyword(Keyword::ENUM)?;
+        Ok(name)
+    });
+    name.map(|name| parse_create_type_as_enum(parser, name))
+}
+
+// https://www.postgresql.org/docs/current/sql-createtype.html
+pub fn parse_create_type_as_enum(
+    parser: &mut Parser,
+    name: ObjectName,
+) -> Result<Statement, ParserError> {
+    if !parser.consume_token(&Token::LParen) {
+        return parser.expected("'(' after CREATE TYPE AS ENUM", parser.peek_token());
+    }
+
+    let labels = parser.parse_comma_separated0(|p| p.parse_identifier(false), Token::RParen)?;
+    parser.expect_token(&Token::RParen)?;
+
+    Ok(Statement::CreateType {
+        name,
+        representation: UserDefinedTypeRepresentation::Enum { labels },
+    })
+}


Test patch:
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 735b87b..bd37214 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -5128,3 +5128,32 @@ fn arrow_cast_precedence() {
         }
     )
 }
+
+#[test]
+fn parse_create_type_as_enum() {
+    let statement = pg().one_statement_parses_to(
+        r#"CREATE TYPE public.my_type AS ENUM (
+            'label1',
+            'label2',
+            'label3',
+            'label4'
+        );"#,
+        "CREATE TYPE public.my_type AS ENUM ('label1', 'label2', 'label3', 'label4')",
+    );
+    match statement {
+        Statement::CreateType {
+            name,
+            representation: UserDefinedTypeRepresentation::Enum { labels },
+        } => {
+            assert_eq!("public.my_type", name.to_string());
+            assert_eq!(
+                vec!["label1", "label2", "label3", "label4"]
+                    .into_iter()
+                    .map(|l| Ident::with_quote('\'', l))
+                    .collect::<Vec<Ident>>(),
+                labels
+            );
+        }
+        _ => unreachable!(),
+    }
+}