Use explicit variants

apache · Jul 8, 2024 · 0ff17d7 · 0ff17d7
1 parent 25d6a40
commit 0ff17d7
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 45 deletions.
diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
@@ -169,14 +169,18 @@ pub enum AlterTableOperation {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum Owner {
     Ident(Ident),
-    Expr(Expr),
+    CurrentRole,
+    CurrentUser,
+    SessionUser,
 }
 
 impl fmt::Display for Owner {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Owner::Ident(ident) => write!(f, "{}", ident),
-            Owner::Expr(expr) => write!(f, "{}", expr),
+            Owner::CurrentRole => write!(f, "CURRENT_ROLE"),
+            Owner::CurrentUser => write!(f, "CURRENT_USER"),
+            Owner::SessionUser => write!(f, "SESSION_USER"),
         }
     }
 }

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -6277,24 +6277,19 @@ impl<'a> Parser<'a> {
         } else if dialect_of!(self is PostgreSqlDialect | GenericDialect)
             && self.parse_keywords(&[Keyword::OWNER, Keyword::TO])
         {
-            let next_token = self.next_token();
-            let new_owner = match next_token.token {
-                Token::DoubleQuotedString(ref s) => Owner::Ident(Ident::new(s.to_string())),
-                Token::Word(ref w) => match w.keyword {
-                    Keyword::CURRENT_USER | Keyword::CURRENT_ROLE | Keyword::SESSION_USER => {
-                        Owner::Expr(Expr::Function(Function {
-                                    name: ObjectName(vec![w.to_ident()]),
-                                    args: FunctionArguments::None,
-                                    null_treatment: None,
-                                    filter: None,
-                                    over: None,
-                                    within_group: vec![],
-                        }))
-                    },
-                    Keyword::NoKeyword => Owner::Ident(w.to_ident()),
-                    _ => self.expected("CURRENT_USER, CURRENT_ROLE, SESSION_USER or identifier after OWNER TO clause", next_token)?,
+            let new_owner = match self.parse_one_of_keywords( &[Keyword::CURRENT_USER, Keyword::CURRENT_ROLE, Keyword::SESSION_USER]) {
+                Some(Keyword::CURRENT_USER) => Owner::CurrentUser,
+                Some(Keyword::CURRENT_ROLE) => Owner::CurrentRole,
+                Some(Keyword::SESSION_USER) => Owner::SessionUser,
+                Some(_) => unreachable!(),
+                None => {
+                    match self.parse_identifier(false) {
+                        Ok(ident) => Owner::Ident(ident),
+                        Err(e) => {
+                            return Err(ParserError::ParserError(format!("Expected CURRENT_USER, CURRENT_ROLE, SESSION_USER or identifier after OWNER TO. {e}")))
+                        }
+                    }
                 },
-                _ => self.expected("Token::Word", next_token)?
             };
 
             AlterTableOperation::OwnerTo { new_owner }

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -725,42 +725,25 @@ fn parse_alter_table_owner_to() {
             sql: "ALTER TABLE tab OWNER TO new_owner",
             expected_owner: Owner::Ident(Ident::new("new_owner".to_string())),
         },
+        TestCase {
+            sql: "ALTER TABLE tab OWNER TO postgres",
+            expected_owner: Owner::Ident(Ident::new("postgres".to_string())),
+        },
         TestCase {
             sql: "ALTER TABLE tab OWNER TO \"new_owner\"",
             expected_owner: Owner::Ident(Ident::with_quote('\"', "new_owner".to_string())),
         },
         TestCase {
             sql: "ALTER TABLE tab OWNER TO CURRENT_USER",
-            expected_owner: Owner::Expr(Expr::Function(Function {
-                name: ObjectName(vec![Ident::new("CURRENT_USER")]),
-                args: FunctionArguments::None,
-                null_treatment: None,
-                filter: None,
-                over: None,
-                within_group: vec![],
-            })),
+            expected_owner: Owner::CurrentUser,
         },
         TestCase {
             sql: "ALTER TABLE tab OWNER TO CURRENT_ROLE",
-            expected_owner: Owner::Expr(Expr::Function(Function {
-                name: ObjectName(vec![Ident::new("CURRENT_ROLE")]),
-                args: FunctionArguments::None,
-                null_treatment: None,
-                filter: None,
-                over: None,
-                within_group: vec![],
-            })),
+            expected_owner: Owner::CurrentRole,
         },
         TestCase {
             sql: "ALTER TABLE tab OWNER TO SESSION_USER",
-            expected_owner: Owner::Expr(Expr::Function(Function {
-                name: ObjectName(vec![Ident::new("SESSION_USER")]),
-                args: FunctionArguments::None,
-                null_treatment: None,
-                filter: None,
-                over: None,
-                within_group: vec![],
-            })),
+            expected_owner: Owner::SessionUser,
         },
     ];
 
@@ -785,9 +768,15 @@ fn parse_alter_table_owner_to() {
         }
     }
 
-    let res = pg().parse_sql_statements("ALTER TABLE tab OWNER TO CREATE");
+    let res = pg().parse_sql_statements("ALTER TABLE tab OWNER TO CREATE FOO");
+    assert_eq!(
+        ParserError::ParserError("Expected end of statement, found: FOO".to_string()),
+        res.unwrap_err()
+    );
+
+    let res = pg().parse_sql_statements("ALTER TABLE tab OWNER TO 4");
     assert_eq!(
-        ParserError::ParserError("Expected CURRENT_USER, CURRENT_ROLE, SESSION_USER or identifier after OWNER TO clause, found: CREATE".to_string()),
+        ParserError::ParserError("Expected CURRENT_USER, CURRENT_ROLE, SESSION_USER or identifier after OWNER TO. sql parser error: Expected identifier, found: 4".to_string()),
         res.unwrap_err()
     );
 }