support COPY INTO in snowflake

Signed-off-by: Pawel Leszczynski <[email protected]>

Author: pawel-big-lebowski

src/ast/helpers/stmt_data_loading.rs

@@ -24,6 +24,7 @@ use core::fmt::Formatter;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
+use crate::ast::Ident;
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
@@ -63,6 +64,16 @@ pub struct DataLoadingOption {
     pub value: String,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct StageLoadSelectItem {
+    pub alias: Option<Ident>,
+    pub file_col_num: i32,
+    pub element: Option<Ident>,
+    pub item_as: Option<Ident>,
+}
+
 impl fmt::Display for StageParamsObject {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let url = &self.url.as_ref();
@@ -121,3 +132,19 @@ impl fmt::Display for DataLoadingOption {
         Ok(())
     }
 }
+
+impl fmt::Display for StageLoadSelectItem {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if self.alias.is_some() {
+            write!(f, "{}.", self.alias.as_ref().unwrap())?;
+        }
+        write!(f, "${}", self.file_col_num)?;
+        if self.element.is_some() {
+            write!(f, ":{}", self.element.as_ref().unwrap())?;
+        }
+        if self.item_as.is_some() {
+            write!(f, " AS {}", self.item_as.as_ref().unwrap())?;
+        }
+        Ok(())
+    }
+}

src/ast/mod.rs

@@ -44,7 +44,9 @@ pub use self::value::{
     escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value,
 };
 
-use crate::ast::helpers::stmt_data_loading::{DataLoadingOptions, StageParamsObject};
+use crate::ast::helpers::stmt_data_loading::{
+    DataLoadingOptions, StageLoadSelectItem, StageParamsObject,
+};
 #[cfg(feature = "visitor")]
 pub use visitor::*;
 
@@ -1187,6 +1189,26 @@ pub enum Statement {
         /// VALUES a vector of values to be copied
         values: Vec<Option<String>>,
     },
+    /// ```sql
+    /// COPY INTO
+    /// ```
+    /// See <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+    /// 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 {
+        into: ObjectName,
+        from_stage: ObjectName,
+        from_stage_alias: Option<Ident>,
+        stage_params: StageParamsObject,
+        from_transformations: Option<Vec<StageLoadSelectItem>>,
+        files: Option<Vec<String>>,
+        pattern: Option<String>,
+        file_format: DataLoadingOptions,
+        copy_options: DataLoadingOptions,
+        validation_mode: Option<String>,
+    },
     /// Close - closes the portal underlying an open cursor.
     Close {
         /// Cursor name
@@ -2795,6 +2817,64 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::CopyIntoSnowflake {
+                into,
+                from_stage,
+                from_stage_alias,
+                stage_params,
+                from_transformations,
+                files,
+                pattern,
+                file_format,
+                copy_options,
+                validation_mode,
+            } => {
+                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 {
+                    // 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())?;
+                    }
+                    write!(f, ")")?;
+                }
+                if files.is_some() {
+                    write!(
+                        f,
+                        " FILES = ('{}')",
+                        display_separated(files.as_ref().unwrap(), "', '")
+                    )?;
+                }
+                if pattern.is_some() {
+                    write!(f, " PATTERN = '{}'", pattern.as_ref().unwrap())?;
+                }
+                if !file_format.options.is_empty() {
+                    write!(f, " FILE_FORMAT=({})", file_format)?;
+                }
+                if !copy_options.options.is_empty() {
+                    write!(f, " COPY_OPTIONS=({})", copy_options)?;
+                }
+                if validation_mode.is_some() {
+                    write!(
+                        f,
+                        " VALIDATION_MODE = {}",
+                        validation_mode.as_ref().unwrap()
+                    )?;
+                }
+                Ok(())
+            }
         }
     }
 }