Parse Snowflake COPY INTO <location> (#1669)

Author: yoavcloud

src/ast/helpers/stmt_data_loading.rs

@@ -58,6 +58,7 @@ pub enum DataLoadingOptionType {
     STRING,
     BOOLEAN,
     ENUM,
+    NUMBER,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -128,12 +129,9 @@ impl fmt::Display for DataLoadingOption {
             DataLoadingOptionType::STRING => {
                 write!(f, "{}='{}'", self.option_name, self.value)?;
             }
-            DataLoadingOptionType::ENUM => {
-                // single quote is omitted
-                write!(f, "{}={}", self.option_name, self.value)?;
-            }
-            DataLoadingOptionType::BOOLEAN => {
-                // single quote is omitted
+            DataLoadingOptionType::ENUM
+            | DataLoadingOptionType::BOOLEAN
+            | DataLoadingOptionType::NUMBER => {
                 write!(f, "{}={}", self.option_name, self.value)?;
             }
         }

src/ast/mod.rs

@@ -2498,24 +2498,30 @@ pub enum Statement {
         values: Vec<Option<String>>,
     },
     /// ```sql
-    /// COPY INTO
+    /// COPY INTO <table> | <location>
     /// ```
-    /// See <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+    /// See:
+    /// <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+    /// <https://docs.snowflake.com/en/sql-reference/sql/copy-into-location>
+    ///
     /// Copy Into syntax available for Snowflake is different than the one implemented in
     /// Postgres. Although they share common prefix, it is reasonable to implement them
     /// in different enums. This can be refactored later once custom dialects
     /// are allowed to have custom Statements.
     CopyIntoSnowflake {
+        kind: CopyIntoSnowflakeKind,
         into: ObjectName,
-        from_stage: ObjectName,
-        from_stage_alias: Option<Ident>,
+        from_obj: Option<ObjectName>,
+        from_obj_alias: Option<Ident>,
         stage_params: StageParamsObject,
         from_transformations: Option<Vec<StageLoadSelectItem>>,
+        from_query: Option<Box<Query>>,
         files: Option<Vec<String>>,
         pattern: Option<String>,
         file_format: DataLoadingOptions,
         copy_options: DataLoadingOptions,
         validation_mode: Option<String>,
+        partition: Option<Box<Expr>>,
     },
     /// ```sql
     /// CLOSE
@@ -5048,60 +5054,69 @@ impl fmt::Display for Statement {
                 Ok(())
             }
             Statement::CopyIntoSnowflake {
+                kind,
                 into,
-                from_stage,
-                from_stage_alias,
+                from_obj,
+                from_obj_alias,
                 stage_params,
                 from_transformations,
+                from_query,
                 files,
                 pattern,
                 file_format,
                 copy_options,
                 validation_mode,
+                partition,
             } => {
                 write!(f, "COPY INTO {}", into)?;
-                if from_transformations.is_none() {
-                    // Standard data load
-                    write!(f, " FROM {}{}", from_stage, stage_params)?;
-                    if from_stage_alias.as_ref().is_some() {
-                        write!(f, " AS {}", from_stage_alias.as_ref().unwrap())?;
-                    }
-                } else {
+                if let Some(from_transformations) = from_transformations {
                     // Data load with transformation
-                    write!(
-                        f,
-                        " FROM (SELECT {} FROM {}{}",
-                        display_separated(from_transformations.as_ref().unwrap(), ", "),
-                        from_stage,
-                        stage_params,
-                    )?;
-                    if from_stage_alias.as_ref().is_some() {
-                        write!(f, " AS {}", from_stage_alias.as_ref().unwrap())?;
+                    if let Some(from_stage) = from_obj {
+                        write!(
+                            f,
+                            " FROM (SELECT {} FROM {}{}",
+                            display_separated(from_transformations, ", "),
+                            from_stage,
+                            stage_params
+                        )?;
+                    }
+                    if let Some(from_obj_alias) = from_obj_alias {
+                        write!(f, " AS {}", from_obj_alias)?;
                     }
                     write!(f, ")")?;
+                } else if let Some(from_obj) = from_obj {
+                    // Standard data load
+                    write!(f, " FROM {}{}", from_obj, stage_params)?;
+                    if let Some(from_obj_alias) = from_obj_alias {
+                        write!(f, " AS {from_obj_alias}")?;
+                    }
+                } else if let Some(from_query) = from_query {
+                    // Data unload from query
+                    write!(f, " FROM ({from_query})")?;
                 }
-                if files.is_some() {
-                    write!(
-                        f,
-                        " FILES = ('{}')",
-                        display_separated(files.as_ref().unwrap(), "', '")
-                    )?;
+
+                if let Some(files) = files {
+                    write!(f, " FILES = ('{}')", display_separated(files, "', '"))?;
+                }
+                if let Some(pattern) = pattern {
+                    write!(f, " PATTERN = '{}'", pattern)?;
                 }
-                if pattern.is_some() {
-                    write!(f, " PATTERN = '{}'", pattern.as_ref().unwrap())?;
+                if let Some(partition) = partition {
+                    write!(f, " PARTITION BY {partition}")?;
                 }
                 if !file_format.options.is_empty() {
                     write!(f, " FILE_FORMAT=({})", file_format)?;
                 }
                 if !copy_options.options.is_empty() {
-                    write!(f, " COPY_OPTIONS=({})", copy_options)?;
+                    match kind {
+                        CopyIntoSnowflakeKind::Table => {
+                            write!(f, " COPY_OPTIONS=({})", copy_options)?
+                        }
+                        CopyIntoSnowflakeKind::Location => write!(f, " {copy_options}")?,
+                    }
                 }
-                if validation_mode.is_some() {
-                    write!(
-                        f,
-                        " VALIDATION_MODE = {}",
-                        validation_mode.as_ref().unwrap()
-                    )?;
+                if let Some(validation_mode) = validation_mode {
+                    write!(f, " VALIDATION_MODE = {}", validation_mode)?;
                 }
                 Ok(())
             }
@@ -8543,6 +8558,19 @@ impl Display for StorageSerializationPolicy {
     }
 }
 
+/// Variants of the Snowflake `COPY INTO` statement
+#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum CopyIntoSnowflakeKind {
+    /// Loads data from files to a table
+    /// See: <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+    Table,
+    /// Unloads data from a table or query to external files
+    /// See: <https://docs.snowflake.com/en/sql-reference/sql/copy-into-location>
+    Location,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/ast/spans.rs

@@ -333,15 +333,18 @@ impl Spanned for Statement {
             } => source.span(),
             Statement::CopyIntoSnowflake {
                 into: _,
-                from_stage: _,
-                from_stage_alias: _,
+                from_obj: _,
+                from_obj_alias: _,
                 stage_params: _,
                 from_transformations: _,
                 files: _,
                 pattern: _,
                 file_format: _,
                 copy_options: _,
                 validation_mode: _,
+                kind: _,
+                from_query: _,
+                partition: _,
             } => Span::empty(),
             Statement::Close { cursor } => match cursor {
                 CloseCursor::All => Span::empty(),