support from style ctes

Author: hntd187

src/ast/query.rs

@@ -251,11 +251,16 @@ impl fmt::Display for With {
 pub struct Cte {
     pub alias: TableAlias,
     pub query: Query,
+    pub from: Option<Ident>,
 }
 
 impl fmt::Display for Cte {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{} AS ({})", self.alias, self.query)
+        write!(f, "{} AS ({})", self.alias, self.query)?;
+        if let Some(ref fr) = self.from {
+            write!(f, " FROM {}", fr)?;
+        }
+        Ok(())
     }
 }
 

src/parser.rs

@@ -2153,24 +2153,36 @@ impl<'a> Parser<'a> {
     fn parse_cte(&mut self) -> Result<Cte, ParserError> {
         let name = self.parse_identifier()?;
 
-        if self.parse_keyword(Keyword::AS) {
+        let mut cte = if self.parse_keyword(Keyword::AS) {
             self.expect_token(&Token::LParen)?;
             let query = self.parse_query()?;
             self.expect_token(&Token::RParen)?;
             let alias = TableAlias {
                 name,
                 columns: vec![],
             };
-            Ok(Cte { alias, query })
+            Cte {
+                alias,
+                query,
+                from: None,
+            }
         } else {
             let columns = self.parse_parenthesized_column_list(Optional)?;
             self.expect_keyword(Keyword::AS)?;
             self.expect_token(&Token::LParen)?;
             let query = self.parse_query()?;
             self.expect_token(&Token::RParen)?;
             let alias = TableAlias { name, columns };
-            Ok(Cte { alias, query })
+            Cte {
+                alias,
+                query,
+                from: None,
+            }
+        };
+        if self.parse_keyword(Keyword::FROM) {
+            cte.from = Some(self.parse_identifier()?);
         }
+        Ok(cte)
     }
 
     /// Parse a "query body", which is an expression with roughly the

tests/sqlparser_common.rs

@@ -2345,7 +2345,7 @@ fn parse_ctes() {
 
     fn assert_ctes_in_select(expected: &[&str], sel: &Query) {
         for (i, exp) in expected.iter().enumerate() {
-            let Cte { alias, query } = &sel.with.as_ref().unwrap().cte_tables[i];
+            let Cte { alias, query, .. } = &sel.with.as_ref().unwrap().cte_tables[i];
             assert_eq!(*exp, query.to_string());
             assert_eq!(
                 if i == 0 {
@@ -2434,6 +2434,7 @@ fn parse_recursive_cte() {
             }],
         },
         query: cte_query,
+        from: None,
     };
     assert_eq!(with.cte_tables.first().unwrap(), &expected);
 }

tests/sqlparser_hive.rs

@@ -198,6 +198,13 @@ fn map_access() {
     hive().verified_stmt(rename);
 }
 
+#[test]
+fn from_cte() {
+    let rename =
+        "WITH cte AS (SELECT * FROM a.b) FROM cte INSERT INTO TABLE a.b PARTITION (a) SELECT *";
+    println!("{}", hive().verified_stmt(rename));
+}
+
 fn hive() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(HiveDialect {})],