support CREATE/DROP STAGE for Snowflake (#833)

Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>

src/ast/helpers/mod.rs

@@ -1 +1,2 @@
 pub mod stmt_create_table;
+pub mod stmt_data_loading;

src/ast/helpers/stmt_data_loading.rs

@@ -0,0 +1,123 @@
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! AST types specific to loading and unloading syntax, like one available in Snowflake which
+//! contains: STAGE ddl operations, PUT upload or COPY INTO
+//! See [this page](https://docs.snowflake.com/en/sql-reference/commands-data-loading) for more details.
+
+#[cfg(not(feature = "std"))]
+use alloc::string::String;
+#[cfg(not(feature = "std"))]
+use alloc::vec::Vec;
+use core::fmt;
+use core::fmt::Formatter;
+
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+#[cfg(feature = "visitor")]
+use sqlparser_derive::{Visit, VisitMut};
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct StageParamsObject {
+    pub url: Option<String>,
+    pub encryption: DataLoadingOptions,
+    pub endpoint: Option<String>,
+    pub storage_integration: Option<String>,
+    pub credentials: DataLoadingOptions,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct DataLoadingOptions {
+    pub options: Vec<DataLoadingOption>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum DataLoadingOptionType {
+    STRING,
+    BOOLEAN,
+    ENUM,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct DataLoadingOption {
+    pub option_name: String,
+    pub option_type: DataLoadingOptionType,
+    pub value: String,
+}
+
+impl fmt::Display for StageParamsObject {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let url = &self.url.as_ref();
+        let storage_integration = &self.storage_integration.as_ref();
+        let endpoint = &self.endpoint.as_ref();
+
+        if url.is_some() {
+            write!(f, " URL='{}'", url.unwrap())?;
+        }
+        if storage_integration.is_some() {
+            write!(f, " STORAGE_INTEGRATION={}", storage_integration.unwrap())?;
+        }
+        if endpoint.is_some() {
+            write!(f, " ENDPOINT='{}'", endpoint.unwrap())?;
+        }
+        if !self.credentials.options.is_empty() {
+            write!(f, " CREDENTIALS=({})", self.credentials)?;
+        }
+        if !self.encryption.options.is_empty() {
+            write!(f, " ENCRYPTION=({})", self.encryption)?;
+        }
+
+        Ok(())
+    }
+}
+
+impl fmt::Display for DataLoadingOptions {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        if !self.options.is_empty() {
+            for option in &self.options {
+                write!(f, "{}", option)?;
+                if !option.eq(self.options.last().unwrap()) {
+                    write!(f, " ")?;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+impl fmt::Display for DataLoadingOption {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self.option_type {
+            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
+                write!(f, "{}={}", self.option_name, self.value)?;
+            }
+        }
+        Ok(())
+    }
+}

src/ast/mod.rs

@@ -44,6 +44,7 @@ pub use self::value::{
     escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value,
 };
 
+use crate::ast::helpers::stmt_data_loading::{DataLoadingOptions, StageParamsObject};
 #[cfg(feature = "visitor")]
 pub use visitor::*;
 
@@ -1524,6 +1525,21 @@ pub enum Statement {
         /// Optional parameters.
         params: CreateFunctionBody,
     },
+    /// ```sql
+    /// CREATE STAGE
+    /// ```
+    /// See <https://docs.snowflake.com/en/sql-reference/sql/create-stage>
+    CreateStage {
+        or_replace: bool,
+        temporary: bool,
+        if_not_exists: bool,
+        name: ObjectName,
+        stage_params: StageParamsObject,
+        directory_table_params: DataLoadingOptions,
+        file_format: DataLoadingOptions,
+        copy_options: DataLoadingOptions,
+        comment: Option<String>,
+    },
     /// `ASSERT <condition> [AS <message>]`
     Assert {
         condition: Expr,
@@ -2746,6 +2762,39 @@ impl fmt::Display for Statement {
                 }
                 write!(f, "")
             }
+            Statement::CreateStage {
+                or_replace,
+                temporary,
+                if_not_exists,
+                name,
+                stage_params,
+                directory_table_params,
+                file_format,
+                copy_options,
+                comment,
+                ..
+            } => {
+                write!(
+                    f,
+                    "CREATE {or_replace}{temp}STAGE {if_not_exists}{name}{stage_params}",
+                    temp = if *temporary { "TEMPORARY " } else { "" },
+                    or_replace = if *or_replace { "OR REPLACE " } else { "" },
+                    if_not_exists = if *if_not_exists { "IF NOT EXISTS " } else { "" },
+                )?;
+                if !directory_table_params.options.is_empty() {
+                    write!(f, " DIRECTORY=({})", directory_table_params)?;
+                }
+                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 comment.is_some() {
+                    write!(f, " COMMENT='{}'", comment.as_ref().unwrap())?;
+                }
+                Ok(())
+            }
         }
     }
 }
@@ -3405,6 +3454,7 @@ pub enum ObjectType {
     Schema,
     Role,
     Sequence,
+    Stage,
 }
 
 impl fmt::Display for ObjectType {
@@ -3416,6 +3466,7 @@ impl fmt::Display for ObjectType {
             ObjectType::Schema => "SCHEMA",
             ObjectType::Role => "ROLE",
             ObjectType::Sequence => "SEQUENCE",
+            ObjectType::Stage => "STAGE",
         })
     }
 }