Implement SnowFlake ALTER SESSION (apache#1712)

Author: osipovartem

src/ast/helpers/key_value_options.rs

@@ -0,0 +1,89 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+//! Key-value options for SQL statements.
+//! 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 KeyValueOptions {
+    pub options: Vec<KeyValueOption>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum KeyValueOptionType {
+    STRING,
+    BOOLEAN,
+    ENUM,
+    NUMBER,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct KeyValueOption {
+    pub option_name: String,
+    pub option_type: KeyValueOptionType,
+    pub value: String,
+}
+
+impl fmt::Display for KeyValueOptions {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        if !self.options.is_empty() {
+            let mut first = false;
+            for option in &self.options {
+                if !first {
+                    first = true;
+                } else {
+                    f.write_str(" ")?;
+                }
+                write!(f, "{}", option)?;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl fmt::Display for KeyValueOption {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self.option_type {
+            KeyValueOptionType::STRING => {
+                write!(f, "{}='{}'", self.option_name, self.value)?;
+            }
+            KeyValueOptionType::ENUM | KeyValueOptionType::BOOLEAN | KeyValueOptionType::NUMBER => {
+                write!(f, "{}={}", self.option_name, self.value)?;
+            }
+        }
+        Ok(())
+    }
+}

src/ast/helpers/mod.rs

@@ -15,5 +15,6 @@
 // specific language governing permissions and limitations
 // under the License.
 pub mod attached_token;
+pub mod key_value_options;
 pub mod stmt_create_table;
 pub mod stmt_data_loading;

src/ast/helpers/stmt_data_loading.rs

@@ -24,11 +24,11 @@ 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};
 
+use crate::ast::helpers::key_value_options::KeyValueOptions;
 use crate::ast::{Ident, ObjectName};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
@@ -38,36 +38,10 @@ use sqlparser_derive::{Visit, VisitMut};
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct StageParamsObject {
     pub url: Option<String>,
-    pub encryption: DataLoadingOptions,
+    pub encryption: KeyValueOptions,
     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,
-    NUMBER,
-}
-
-#[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,
+    pub credentials: KeyValueOptions,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -106,39 +80,6 @@ impl fmt::Display for StageParamsObject {
     }
 }
 
-impl fmt::Display for DataLoadingOptions {
-    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        if !self.options.is_empty() {
-            let mut first = false;
-            for option in &self.options {
-                if !first {
-                    first = true;
-                } else {
-                    f.write_str(" ")?;
-                }
-                write!(f, "{}", option)?;
-            }
-        }
-        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
-            | DataLoadingOptionType::BOOLEAN
-            | DataLoadingOptionType::NUMBER => {
-                write!(f, "{}={}", self.option_name, self.value)?;
-            }
-        }
-        Ok(())
-    }
-}
-
 impl fmt::Display for StageLoadSelectItem {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         if self.alias.is_some() {

src/ast/mod.rs

@@ -89,9 +89,8 @@ pub use self::value::{
     NormalizationForm, TrimWhereField, Value,
 };
 
-use crate::ast::helpers::stmt_data_loading::{
-    DataLoadingOptions, StageLoadSelectItem, StageParamsObject,
-};
+use crate::ast::helpers::key_value_options::KeyValueOptions;
+use crate::ast::helpers::stmt_data_loading::{StageLoadSelectItem, StageParamsObject};
 #[cfg(feature = "visitor")]
 pub use visitor::*;
 
@@ -2504,8 +2503,8 @@ pub enum Statement {
         from_query: Option<Box<Query>>,
         files: Option<Vec<String>>,
         pattern: Option<String>,
-        file_format: DataLoadingOptions,
-        copy_options: DataLoadingOptions,
+        file_format: KeyValueOptions,
+        copy_options: KeyValueOptions,
         validation_mode: Option<String>,
         partition: Option<Box<Expr>>,
     },
@@ -2713,6 +2712,17 @@ pub enum Statement {
         owner: Option<ddl::AlterConnectorOwner>,
     },
     /// ```sql
+    /// ALTER SESSION SET sessionParam
+    /// ALTER SESSION UNSET <param_name> [ , <param_name> , ... ]
+    /// ```
+    /// See <https://docs.snowflake.com/en/sql-reference/sql/alter-session>
+    AlterSession {
+        /// true is to set for the session parameters, false is to unset
+        set: bool,
+        /// The session parameters to set or unset
+        session_params: KeyValueOptions,
+    },
+    /// ```sql
     /// ATTACH DATABASE 'path/to/file' AS alias
     /// ```
     /// (SQLite-specific)
@@ -3240,9 +3250,9 @@ pub enum Statement {
         if_not_exists: bool,
         name: ObjectName,
         stage_params: StageParamsObject,
-        directory_table_params: DataLoadingOptions,
-        file_format: DataLoadingOptions,
-        copy_options: DataLoadingOptions,
+        directory_table_params: KeyValueOptions,
+        file_format: KeyValueOptions,
+        copy_options: KeyValueOptions,
         comment: Option<String>,
     },
     /// ```sql
@@ -4467,6 +4477,29 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::AlterSession {
+                set,
+                session_params,
+            } => {
+                write!(
+                    f,
+                    "ALTER SESSION {set}",
+                    set = if *set { "SET" } else { "UNSET" }
+                )?;
+                if !session_params.options.is_empty() {
+                    if *set {
+                        write!(f, " {}", session_params)?;
+                    } else {
+                        let options = session_params
+                            .options
+                            .iter()
+                            .map(|p| p.option_name.clone())
+                            .collect::<Vec<_>>();
+                        write!(f, " {}", display_separated(&options, ", "))?;
+                    }
+                }
+                Ok(())
+            }
             Statement::Drop {
                 object_type,
                 if_exists,

src/ast/spans.rs

@@ -430,6 +430,7 @@ impl Spanned for Statement {
             // These statements need to be implemented
             Statement::AlterType { .. } => Span::empty(),
             Statement::AlterRole { .. } => Span::empty(),
+            Statement::AlterSession { .. } => Span::empty(),
             Statement::AttachDatabase { .. } => Span::empty(),
             Statement::AttachDuckDBDatabase { .. } => Span::empty(),
             Statement::DetachDuckDBDatabase { .. } => Span::empty(),