Add PostgreSQL specfic "CREATE TYPE t AS ENUM (...)" support. (apache#1460)

Author: caldwell

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))
+            }
         }
     }
 }

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 },
+    })
+}

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!(),
+    }
+}