Skip CREATE STAGE by DFParser (pass to dialect parser)

datafusion/sql/src/parser.rs

@@ -557,7 +557,11 @@ impl<'a> DFParser<'a> {
             self.parser.expect_keyword(Keyword::EXTERNAL)?;
             self.parse_create_external_table(true)
         } else {
-            Ok(Statement::Statement(Box::from(self.parser.parse_create()?)))
+            // Push back CREATE
+            self.parser.prev_token();
+            Ok(Statement::Statement(Box::from(
+                self.parser.parse_statement()?,
+            )))
         }
     }
 
@@ -885,6 +889,7 @@ mod tests {
     use super::*;
     use sqlparser::ast::Expr::Identifier;
     use sqlparser::ast::{BinaryOperator, DataType, Expr, Ident};
+    use sqlparser::dialect::SnowflakeDialect;
     use sqlparser::tokenizer::Span;
 
     fn expect_parse_ok(sql: &str, expected: Statement) -> Result<(), ParserError> {
@@ -929,6 +934,26 @@ mod tests {
         }
     }
 
+    #[test]
+    fn skip_create_stage_snowflake() -> Result<(), ParserError> {
+        let sql =
+            "CREATE OR REPLACE STAGE stage URL='s3://data.csv' FILE_FORMAT=(TYPE=csv)";
+        let dialect = Box::new(SnowflakeDialect);
+        let statements = DFParser::parse_sql_with_dialect(sql, dialect.as_ref())?;
+
+        assert_eq!(
+            statements.len(),
+            1,
+            "Expected to parse exactly one statement"
+        );
+        match &statements[0] {
+            Statement::Statement(stmt) => {
+                assert_eq!(stmt.to_string(), sql);
+            }
+            _ => panic!("Expected statement type"),
+        }
+        Ok(())
+    }
     #[test]
     fn create_external_table() -> Result<(), ParserError> {
         // positive case
@@ -1411,6 +1436,26 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn skip_copy_into_snowflake() -> Result<(), ParserError> {
+        let sql = "COPY INTO foo FROM @~/staged FILE_FORMAT=(FORMAT_NAME='mycsv')";
+        let dialect = Box::new(SnowflakeDialect);
+        let statements = DFParser::parse_sql_with_dialect(sql, dialect.as_ref())?;
+
+        assert_eq!(
+            statements.len(),
+            1,
+            "Expected to parse exactly one statement"
+        );
+        match &statements[0] {
+            Statement::Statement(stmt) => {
+                assert_eq!(stmt.to_string(), sql);
+            }
+            _ => panic!("Expected statement type"),
+        }
+        Ok(())
+    }
+
     #[test]
     fn explain_copy_to_table_to_table() -> Result<(), ParserError> {
         let cases = vec![