moved partition check into partition statement

Author: jonathanlehto

src/parser/mod.rs

@@ -8410,16 +8410,9 @@ impl<'a> Parser<'a> {
             } else {
                 let columns = self.parse_parenthesized_column_list(Optional, is_mysql)?;
 
-                let partitioned = self.parse_insert_partition()?;
-
-                // make sure we are not about to consume a query
-                let mut after_columns = vec![];
-                match self.peek_nth_token(1).token {
-                    Token::Word(w) if w.keyword != Keyword::SELECT => {
-                        // Hive allows you to specify columns after partitions as well if you want.
-                        after_columns = self.parse_parenthesized_column_list(Optional, false)?;
-                    }
-                    _ => {}
+                let (partitioned, after_columns) = match self.parse_insert_partition()? {
+                    Some((partitioned, after_columns)) => (Some(partitioned), after_columns),
+                    None => (None, vec![])
                 };
 
                 let source = Some(Box::new(self.parse_query()?));
@@ -8500,12 +8493,16 @@ impl<'a> Parser<'a> {
         }
     }
 
-    pub fn parse_insert_partition(&mut self) -> Result<Option<Vec<Expr>>, ParserError> {
+    pub fn parse_insert_partition(&mut self) -> Result<Option<(Vec<Expr>, Vec<Ident>)>, ParserError> {
         if self.parse_keyword(Keyword::PARTITION) {
             self.expect_token(&Token::LParen)?;
-            let partition_cols = Some(self.parse_comma_separated(Parser::parse_expr)?);
+            let partition_cols = self.parse_comma_separated(Parser::parse_expr)?;
             self.expect_token(&Token::RParen)?;
-            Ok(partition_cols)
+
+            // Hive allows you to specify columns after partitions as well if you want.
+            let after_columns = self.parse_parenthesized_column_list(Optional, false)?;
+
+            Ok(Some((partition_cols, after_columns)))
         } else {
             Ok(None)
         }