Snowflake: Add support for PUT (apache#2341)

Author: ramnes

src/ast/mod.rs

@@ -4852,6 +4852,20 @@ pub enum Statement {
     /// Snowflake `LIST`
     /// See: <https://docs.snowflake.com/en/sql-reference/sql/list>
     List(FileStagingCommand),
+    /// Snowflake `PUT`
+    /// ```sql
+    /// PUT 'file://<path>' <internalStage> [ <option> = <value> ... ]
+    /// ```
+    /// Options include `PARALLEL`, `AUTO_COMPRESS`, `SOURCE_COMPRESSION`, `OVERWRITE`.
+    /// See: <https://docs.snowflake.com/en/sql-reference/sql/put>
+    Put {
+        /// Local source URI as written in the statement, e.g. `file:///tmp/data.csv`.
+        source: String,
+        /// Target internal stage (e.g. `@mystage`, `@~`, `@%table`).
+        stage: ObjectName,
+        /// Trailing options (`PARALLEL=4`, `AUTO_COMPRESS=TRUE`, ...).
+        options: KeyValueOptions,
+    },
     /// Snowflake `REMOVE`
     /// See: <https://docs.snowflake.com/en/sql-reference/sql/remove>
     Remove(FileStagingCommand),
@@ -6372,6 +6386,17 @@ impl fmt::Display for Statement {
             Statement::WaitFor(s) => write!(f, "{s}"),
             Statement::Return(r) => write!(f, "{r}"),
             Statement::List(command) => write!(f, "LIST {command}"),
+            Statement::Put {
+                source,
+                stage,
+                options,
+            } => {
+                write!(f, "PUT '{source}' {stage}")?;
+                if !options.options.is_empty() {
+                    write!(f, " {options}")?;
+                }
+                Ok(())
+            }
             Statement::Remove(command) => write!(f, "REMOVE {command}"),
             Statement::ExportData(e) => write!(f, "{e}"),
             Statement::CreateUser(s) => write!(f, "{s}"),

src/ast/spans.rs

@@ -501,7 +501,7 @@ impl Spanned for Statement {
             Statement::Print { .. } => Span::empty(),
             Statement::WaitFor(_) => Span::empty(),
             Statement::Return { .. } => Span::empty(),
-            Statement::List(..) | Statement::Remove(..) => Span::empty(),
+            Statement::List(..) | Statement::Put { .. } | Statement::Remove(..) => Span::empty(),
             Statement::ExportData(ExportData {
                 options,
                 query,

src/dialect/snowflake.rs

@@ -354,6 +354,10 @@ impl Dialect for SnowflakeDialect {
             return Some(parse_file_staging_command(kw, parser));
         }
 
+        if parser.parse_keyword(Keyword::PUT) {
+            return Some(parse_put(parser));
+        }
+
         if parser.parse_keyword(Keyword::SHOW) {
             let terse = parser.parse_keyword(Keyword::TERSE);
             if parser.parse_keyword(Keyword::OBJECTS) {
@@ -696,6 +700,21 @@ fn peek_for_limit_options(parser: &Parser) -> bool {
     }
 }
 
+/// Parse a Snowflake `PUT <source> <stage> [ options ]` statement. The caller
+/// is expected to have already consumed `PUT`.
+///
+/// See <https://docs.snowflake.com/en/sql-reference/sql/put>.
+fn parse_put(parser: &mut Parser) -> Result<Statement, ParserError> {
+    let source = parser.parse_literal_string()?;
+    let stage = parse_snowflake_stage_name(parser)?;
+    let options = parser.parse_key_value_options(false, &[])?;
+    Ok(Statement::Put {
+        source,
+        stage,
+        options,
+    })
+}
+
 fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {
     let stage = parse_snowflake_stage_name(parser)?;
     let pattern = if parser.parse_keyword(Keyword::PATTERN) {

src/keywords.rs

@@ -822,6 +822,7 @@ define_keywords!(
     PUBLIC,
     PURCHASE,
     PURGE,
+    PUT,
     QUALIFY,
     QUARTER,
     QUERIES,

tests/sqlparser_snowflake.rs

@@ -3861,6 +3861,71 @@ fn parse_ls_and_rm() {
         .unwrap();
 }
 
+#[test]
+fn test_put() {
+    let sql = "PUT 'file:///tmp/data.csv' @my_stage";
+    match snowflake().verified_stmt(sql) {
+        Statement::Put {
+            source,
+            stage,
+            options,
+        } => {
+            assert_eq!("file:///tmp/data.csv", source);
+            assert_eq!(ObjectName::from(vec!["@my_stage".into()]), stage);
+            assert!(options.options.is_empty());
+        }
+        _ => unreachable!(),
+    };
+    assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
+}
+
+#[test]
+fn test_put_with_quoted_stage() {
+    // Stage names can be quoted (e.g. Snowflake driver `write_pandas`)
+    let sql = r#"PUT 'file:///tmp/data.csv' @"my stage" PARALLEL=4"#;
+    match snowflake().verified_stmt(sql) {
+        Statement::Put { stage, .. } => {
+            assert_eq!(ObjectName::from(vec![r#"@"my stage""#.into()]), stage);
+        }
+        _ => unreachable!(),
+    };
+    assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
+}
+
+#[test]
+fn test_put_with_options() {
+    let sql = concat!(
+        "PUT 'file:///tmp/data.csv' @my_stage ",
+        "PARALLEL=8 AUTO_COMPRESS=true SOURCE_COMPRESSION=GZIP OVERWRITE=false"
+    );
+    match snowflake().verified_stmt(sql) {
+        Statement::Put { options, .. } => {
+            assert!(options.options.contains(&KeyValueOption {
+                option_name: "PARALLEL".to_string(),
+                option_value: KeyValueOptionKind::Single(
+                    Value::Number("8".parse().unwrap(), false).with_empty_span()
+                ),
+            }));
+            assert!(options.options.contains(&KeyValueOption {
+                option_name: "AUTO_COMPRESS".to_string(),
+                option_value: KeyValueOptionKind::Single(Value::Boolean(true).with_empty_span()),
+            }));
+            assert!(options.options.contains(&KeyValueOption {
+                option_name: "SOURCE_COMPRESSION".to_string(),
+                option_value: KeyValueOptionKind::Single(
+                    Value::Placeholder("GZIP".to_string()).with_empty_span()
+                ),
+            }));
+            assert!(options.options.contains(&KeyValueOption {
+                option_name: "OVERWRITE".to_string(),
+                option_value: KeyValueOptionKind::Single(Value::Boolean(false).with_empty_span()),
+            }));
+        }
+        _ => unreachable!(),
+    };
+    assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
+}
+
 #[test]
 fn test_sql_keywords_as_select_item_ident() {
     // Some keywords that should be parsed as an alias