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

Issue:
Snowflake COMMENT ON support
Support for https://docs.snowflake.com/en/sql-reference/sql/comment, for example:
```sql
COMMENT ON my_schema.my_table IS 'hello world'
```

Currently, parsing fails:

```
Expected: an SQL statement, found: COMMENT
```

Pull request description:
Add support of COMMENT ON syntax for Snowflake
The COMMENT ON syntax has been supported in PostgreSQL. This PR only moves it into the common and adds `supports_comment_on` to enable this feature. Also, this PR extends `SCHEMA`, `DATABASE`, `USER`, `ROLE` as its allowed object type.

This closes #1511.

Patch:
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 848e6bd..505386f 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -1884,6 +1884,10 @@ pub enum CommentObject {
     Column,
     Table,
     Extension,
+    Schema,
+    Database,
+    User,
+    Role,
 }
 
 impl fmt::Display for CommentObject {
@@ -1892,6 +1896,10 @@ impl fmt::Display for CommentObject {
             CommentObject::Column => f.write_str("COLUMN"),
             CommentObject::Table => f.write_str("TABLE"),
             CommentObject::Extension => f.write_str("EXTENSION"),
+            CommentObject::Schema => f.write_str("SCHEMA"),
+            CommentObject::Database => f.write_str("DATABASE"),
+            CommentObject::User => f.write_str("USER"),
+            CommentObject::Role => f.write_str("ROLE"),
         }
     }
 }
diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs
index 0a5464c..8cfac21 100644
--- a/src/dialect/generic.rs
+++ b/src/dialect/generic.rs
@@ -111,4 +111,8 @@ impl Dialect for GenericDialect {
     fn supports_try_convert(&self) -> bool {
         true
     }
+
+    fn supports_comment_on(&self) -> bool {
+        true
+    }
 }
diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs
index f37c0d8..d95d7c7 100644
--- a/src/dialect/mod.rs
+++ b/src/dialect/mod.rs
@@ -611,7 +611,7 @@ pub trait Dialect: Debug + Any {
         false
     }
 
-    /// Returns true if this dialect expects the the `TOP` option
+    /// Returns true if this dialect expects the `TOP` option
     /// before the `ALL`/`DISTINCT` options in a `SELECT` statement.
     fn supports_top_before_distinct(&self) -> bool {
         false
@@ -628,6 +628,11 @@ pub trait Dialect: Debug + Any {
     fn supports_show_like_before_in(&self) -> bool {
         false
     }
+
+    /// Returns true if this dialect supports the `COMMENT` statement
+    fn supports_comment_on(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined
diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs
index 72841c6..5af1ab8 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, ObjectName, Statement, UserDefinedTypeRepresentation};
+use crate::ast::{ObjectName, Statement, UserDefinedTypeRepresentation};
 use crate::dialect::{Dialect, Precedence};
 use crate::keywords::Keyword;
 use crate::parser::{Parser, ParserError};
@@ -136,9 +136,7 @@ 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) {
+        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 {
@@ -206,42 +204,11 @@ impl Dialect for PostgreSqlDialect {
     fn supports_factorial_operator(&self) -> bool {
         true
     }
-}
-
-pub fn parse_comment(parser: &mut Parser) -> Result<Statement, ParserError> {
-    let if_exists = parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
-
-    parser.expect_keyword(Keyword::ON)?;
-    let token = parser.next_token();
-
-    let (object_type, object_name) = match token.token {
-        Token::Word(w) if w.keyword == Keyword::COLUMN => {
-            let object_name = parser.parse_object_name(false)?;
-            (CommentObject::Column, object_name)
-        }
-        Token::Word(w) if w.keyword == Keyword::TABLE => {
-            let object_name = parser.parse_object_name(false)?;
-            (CommentObject::Table, object_name)
-        }
-        Token::Word(w) if w.keyword == Keyword::EXTENSION => {
-            let object_name = parser.parse_object_name(false)?;
-            (CommentObject::Extension, object_name)
-        }
-        _ => parser.expected("comment object_type", token)?,
-    };
 
-    parser.expect_keyword(Keyword::IS)?;
-    let comment = if parser.parse_keyword(Keyword::NULL) {
-        None
-    } else {
-        Some(parser.parse_literal_string()?)
-    };
-    Ok(Statement::Comment {
-        object_type,
-        object_name,
-        comment,
-        if_exists,
-    })
+    /// see <https://www.postgresql.org/docs/current/sql-comment.html>
+    fn supports_comment_on(&self) -> bool {
+        true
+    }
 }
 
 pub fn parse_create(parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs
index 98e8f5e..b584ed9 100644
--- a/src/dialect/snowflake.rs
+++ b/src/dialect/snowflake.rs
@@ -96,6 +96,11 @@ impl Dialect for SnowflakeDialect {
         true
     }
 
+    /// See [doc](https://docs.snowflake.com/en/sql-reference/sql/comment)
+    fn supports_comment_on(&self) -> bool {
+        true
+    }
+
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
         if parser.parse_keyword(Keyword::CREATE) {
             // possibly CREATE STAGE
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 7f6961a..756f4d6 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -551,6 +551,8 @@ impl<'a> Parser<'a> {
                 Keyword::OPTIMIZE if dialect_of!(self is ClickHouseDialect | GenericDialect) => {
                     self.parse_optimize_table()
                 }
+                // `COMMENT` is snowflake specific https://docs.snowflake.com/en/sql-reference/sql/comment
+                Keyword::COMMENT if self.dialect.supports_comment_on() => self.parse_comment(),
                 _ => self.expected("an SQL statement", next_token),
             },
             Token::LParen => {
@@ -561,6 +563,51 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_comment(&mut self) -> Result<Statement, ParserError> {
+        let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+
+        self.expect_keyword(Keyword::ON)?;
+        let token = self.next_token();
+
+        let (object_type, object_name) = match token.token {
+            Token::Word(w) if w.keyword == Keyword::COLUMN => {
+                (CommentObject::Column, self.parse_object_name(false)?)
+            }
+            Token::Word(w) if w.keyword == Keyword::TABLE => {
+                (CommentObject::Table, self.parse_object_name(false)?)
+            }
+            Token::Word(w) if w.keyword == Keyword::EXTENSION => {
+                (CommentObject::Extension, self.parse_object_name(false)?)
+            }
+            Token::Word(w) if w.keyword == Keyword::SCHEMA => {
+                (CommentObject::Schema, self.parse_object_name(false)?)
+            }
+            Token::Word(w) if w.keyword == Keyword::DATABASE => {
+                (CommentObject::Database, self.parse_object_name(false)?)
+            }
+            Token::Word(w) if w.keyword == Keyword::USER => {
+                (CommentObject::User, self.parse_object_name(false)?)
+            }
+            Token::Word(w) if w.keyword == Keyword::ROLE => {
+                (CommentObject::Role, self.parse_object_name(false)?)
+            }
+            _ => self.expected("comment object_type", token)?,
+        };
+
+        self.expect_keyword(Keyword::IS)?;
+        let comment = if self.parse_keyword(Keyword::NULL) {
+            None
+        } else {
+            Some(self.parse_literal_string()?)
+        };
+        Ok(Statement::Comment {
+            object_type,
+            object_name,
+            comment,
+            if_exists,
+        })
+    }
+
     pub fn parse_flush(&mut self) -> Result<Statement, ParserError> {
         let mut channel = None;
         let mut tables: Vec<ObjectName> = vec![];


Test patch:
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index d08e19d..25bf306 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -11629,3 +11629,91 @@ fn parse_factorial_operator() {
         );
     }
 }
+
+#[test]
+fn parse_comments() {
+    match all_dialects_where(|d| d.supports_comment_on())
+        .verified_stmt("COMMENT ON COLUMN tab.name IS 'comment'")
+    {
+        Statement::Comment {
+            object_type,
+            object_name,
+            comment: Some(comment),
+            if_exists,
+        } => {
+            assert_eq!("comment", comment);
+            assert_eq!("tab.name", object_name.to_string());
+            assert_eq!(CommentObject::Column, object_type);
+            assert!(!if_exists);
+        }
+        _ => unreachable!(),
+    }
+
+    let object_types = [
+        ("COLUMN", CommentObject::Column),
+        ("EXTENSION", CommentObject::Extension),
+        ("TABLE", CommentObject::Table),
+        ("SCHEMA", CommentObject::Schema),
+        ("DATABASE", CommentObject::Database),
+        ("USER", CommentObject::User),
+        ("ROLE", CommentObject::Role),
+    ];
+    for (keyword, expected_object_type) in object_types.iter() {
+        match all_dialects_where(|d| d.supports_comment_on())
+            .verified_stmt(format!("COMMENT IF EXISTS ON {keyword} db.t0 IS 'comment'").as_str())
+        {
+            Statement::Comment {
+                object_type,
+                object_name,
+                comment: Some(comment),
+                if_exists,
+            } => {
+                assert_eq!("comment", comment);
+                assert_eq!("db.t0", object_name.to_string());
+                assert_eq!(*expected_object_type, object_type);
+                assert!(if_exists);
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    match all_dialects_where(|d| d.supports_comment_on())
+        .verified_stmt("COMMENT IF EXISTS ON TABLE public.tab IS NULL")
+    {
+        Statement::Comment {
+            object_type,
+            object_name,
+            comment: None,
+            if_exists,
+        } => {
+            assert_eq!("public.tab", object_name.to_string());
+            assert_eq!(CommentObject::Table, object_type);
+            assert!(if_exists);
+        }
+        _ => unreachable!(),
+    }
+
+    // missing IS statement
+    assert_eq!(
+        all_dialects_where(|d| d.supports_comment_on())
+            .parse_sql_statements("COMMENT ON TABLE t0")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: IS, found: EOF".to_string())
+    );
+
+    // missing comment literal
+    assert_eq!(
+        all_dialects_where(|d| d.supports_comment_on())
+            .parse_sql_statements("COMMENT ON TABLE t0 IS")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: literal string, found: EOF".to_string())
+    );
+
+    // unknown object type
+    assert_eq!(
+        all_dialects_where(|d| d.supports_comment_on())
+            .parse_sql_statements("COMMENT ON UNKNOWN t0 IS 'comment'")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: comment object_type, found: UNKNOWN".to_string())
+    );
+}
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index 100c8ee..a6c480c 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -2891,68 +2891,6 @@ fn test_composite_value() {
     );
 }
 
-#[test]
-fn parse_comments() {
-    match pg().verified_stmt("COMMENT ON COLUMN tab.name IS 'comment'") {
-        Statement::Comment {
-            object_type,
-            object_name,
-            comment: Some(comment),
-            if_exists,
-        } => {
-            assert_eq!("comment", comment);
-            assert_eq!("tab.name", object_name.to_string());
-            assert_eq!(CommentObject::Column, object_type);
-            assert!(!if_exists);
-        }
-        _ => unreachable!(),
-    }
-
-    match pg().verified_stmt("COMMENT ON EXTENSION plpgsql IS 'comment'") {
-        Statement::Comment {
-            object_type,
-            object_name,
-            comment: Some(comment),
-            if_exists,
-        } => {
-            assert_eq!("comment", comment);
-            assert_eq!("plpgsql", object_name.to_string());
-            assert_eq!(CommentObject::Extension, object_type);
-            assert!(!if_exists);
-        }
-        _ => unreachable!(),
-    }
-
-    match pg().verified_stmt("COMMENT ON TABLE public.tab IS 'comment'") {
-        Statement::Comment {
-            object_type,
-            object_name,
-            comment: Some(comment),
-            if_exists,
-        } => {
-            assert_eq!("comment", comment);
-            assert_eq!("public.tab", object_name.to_string());
-            assert_eq!(CommentObject::Table, object_type);
-            assert!(!if_exists);
-        }
-        _ => unreachable!(),
-    }
-
-    match pg().verified_stmt("COMMENT IF EXISTS ON TABLE public.tab IS NULL") {
-        Statement::Comment {
-            object_type,
-            object_name,
-            comment: None,
-            if_exists,
-        } => {
-            assert_eq!("public.tab", object_name.to_string());
-            assert_eq!(CommentObject::Table, object_type);
-            assert!(if_exists);
-        }
-        _ => unreachable!(),
-    }
-}
-
 #[test]
 fn parse_quoted_identifier() {
     pg_and_generic().verified_stmt(r#"SELECT "quoted "" ident""#);