Skip to content

Commit 3ace97c

Browse files
osipovartemosipovartem
authored
Implement SnowFlake ALTER SESSION (apache#1712)
1 parent 28736da commit 3ace97c

File tree

8 files changed

+335
-197
lines changed

8 files changed

+335
-197
lines changed

‎src/ast/helpers/key_value_options.rs‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Key-value options for SQL statements.
19+
//! See [this page](https://docs.snowflake.com/en/sql-reference/commands-data-loading) for more details.
20+
21+
#[cfg(not(feature = "std"))]
22+
use alloc::string::String;
23+
#[cfg(not(feature = "std"))]
24+
use alloc::vec::Vec;
25+
use core::fmt;
26+
use core::fmt::Formatter;
27+
28+
#[cfg(feature = "serde")]
29+
use serde::{Deserialize, Serialize};
30+
31+
#[cfg(feature = "visitor")]
32+
use sqlparser_derive::{Visit, VisitMut};
33+
34+
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
35+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
36+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
37+
pub struct KeyValueOptions {
38+
pub options: Vec<KeyValueOption>,
39+
}
40+
41+
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
42+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
43+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
44+
pub enum KeyValueOptionType {
45+
STRING,
46+
BOOLEAN,
47+
ENUM,
48+
NUMBER,
49+
}
50+
51+
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
52+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
53+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
54+
pub struct KeyValueOption {
55+
pub option_name: String,
56+
pub option_type: KeyValueOptionType,
57+
pub value: String,
58+
}
59+
60+
impl fmt::Display for KeyValueOptions {
61+
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
62+
if !self.options.is_empty() {
63+
let mut first = false;
64+
for option in &self.options {
65+
if !first {
66+
first = true;
67+
} else {
68+
f.write_str(" ")?;
69+
}
70+
write!(f, "{}", option)?;
71+
}
72+
}
73+
Ok(())
74+
}
75+
}
76+
77+
impl fmt::Display for KeyValueOption {
78+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
79+
match self.option_type {
80+
KeyValueOptionType::STRING => {
81+
write!(f, "{}='{}'", self.option_name, self.value)?;
82+
}
83+
KeyValueOptionType::ENUM | KeyValueOptionType::BOOLEAN | KeyValueOptionType::NUMBER => {
84+
write!(f, "{}={}", self.option_name, self.value)?;
85+
}
86+
}
87+
Ok(())
88+
}
89+
}

‎src/ast/helpers/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,6 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717
pub mod attached_token;
18+
pub mod key_value_options;
1819
pub mod stmt_create_table;
1920
pub mod stmt_data_loading;

‎src/ast/helpers/stmt_data_loading.rs‎

Lines changed: 3 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,11 @@ use alloc::string::String;
2424
#[cfg(not(feature = "std"))]
2525
use alloc::vec::Vec;
2626
use core::fmt;
27-
use core::fmt::Formatter;
2827

2928
#[cfg(feature = "serde")]
3029
use serde::{Deserialize, Serialize};
3130

31+
use crate::ast::helpers::key_value_options::KeyValueOptions;
3232
use crate::ast::{Ident, ObjectName};
3333
#[cfg(feature = "visitor")]
3434
use sqlparser_derive::{Visit, VisitMut};
@@ -38,36 +38,10 @@ use sqlparser_derive::{Visit, VisitMut};
3838
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
3939
pub struct StageParamsObject {
4040
pub url: Option<String>,
41-
pub encryption: DataLoadingOptions,
41+
pub encryption: KeyValueOptions,
4242
pub endpoint: Option<String>,
4343
pub storage_integration: Option<String>,
44-
pub credentials: DataLoadingOptions,
45-
}
46-
47-
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
48-
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
49-
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
50-
pub struct DataLoadingOptions {
51-
pub options: Vec<DataLoadingOption>,
52-
}
53-
54-
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
55-
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
56-
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
57-
pub enum DataLoadingOptionType {
58-
STRING,
59-
BOOLEAN,
60-
ENUM,
61-
NUMBER,
62-
}
63-
64-
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
65-
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
66-
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
67-
pub struct DataLoadingOption {
68-
pub option_name: String,
69-
pub option_type: DataLoadingOptionType,
70-
pub value: String,
44+
pub credentials: KeyValueOptions,
7145
}
7246

7347
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -106,39 +80,6 @@ impl fmt::Display for StageParamsObject {
10680
}
10781
}
10882

109-
impl fmt::Display for DataLoadingOptions {
110-
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
111-
if !self.options.is_empty() {
112-
let mut first = false;
113-
for option in &self.options {
114-
if !first {
115-
first = true;
116-
} else {
117-
f.write_str(" ")?;
118-
}
119-
write!(f, "{}", option)?;
120-
}
121-
}
122-
Ok(())
123-
}
124-
}
125-
126-
impl fmt::Display for DataLoadingOption {
127-
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
128-
match self.option_type {
129-
DataLoadingOptionType::STRING => {
130-
write!(f, "{}='{}'", self.option_name, self.value)?;
131-
}
132-
DataLoadingOptionType::ENUM
133-
| DataLoadingOptionType::BOOLEAN
134-
| DataLoadingOptionType::NUMBER => {
135-
write!(f, "{}={}", self.option_name, self.value)?;
136-
}
137-
}
138-
Ok(())
139-
}
140-
}
141-
14283
impl fmt::Display for StageLoadSelectItem {
14384
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
14485
if self.alias.is_some() {

‎src/ast/mod.rs‎

Lines changed: 41 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -89,9 +89,8 @@ pub use self::value::{
8989
NormalizationForm, TrimWhereField, Value,
9090
};
9191

92-
use crate::ast::helpers::stmt_data_loading::{
93-
DataLoadingOptions, StageLoadSelectItem, StageParamsObject,
94-
};
92+
use crate::ast::helpers::key_value_options::KeyValueOptions;
93+
use crate::ast::helpers::stmt_data_loading::{StageLoadSelectItem, StageParamsObject};
9594
#[cfg(feature = "visitor")]
9695
pub use visitor::*;
9796

@@ -2504,8 +2503,8 @@ pub enum Statement {
25042503
from_query: Option<Box<Query>>,
25052504
files: Option<Vec<String>>,
25062505
pattern: Option<String>,
2507-
file_format: DataLoadingOptions,
2508-
copy_options: DataLoadingOptions,
2506+
file_format: KeyValueOptions,
2507+
copy_options: KeyValueOptions,
25092508
validation_mode: Option<String>,
25102509
partition: Option<Box<Expr>>,
25112510
},
@@ -2713,6 +2712,17 @@ pub enum Statement {
27132712
owner: Option<ddl::AlterConnectorOwner>,
27142713
},
27152714
/// ```sql
2715+
/// ALTER SESSION SET sessionParam
2716+
/// ALTER SESSION UNSET <param_name> [ , <param_name> , ... ]
2717+
/// ```
2718+
/// See <https://docs.snowflake.com/en/sql-reference/sql/alter-session>
2719+
AlterSession {
2720+
/// true is to set for the session parameters, false is to unset
2721+
set: bool,
2722+
/// The session parameters to set or unset
2723+
session_params: KeyValueOptions,
2724+
},
2725+
/// ```sql
27162726
/// ATTACH DATABASE 'path/to/file' AS alias
27172727
/// ```
27182728
/// (SQLite-specific)
@@ -3240,9 +3250,9 @@ pub enum Statement {
32403250
if_not_exists: bool,
32413251
name: ObjectName,
32423252
stage_params: StageParamsObject,
3243-
directory_table_params: DataLoadingOptions,
3244-
file_format: DataLoadingOptions,
3245-
copy_options: DataLoadingOptions,
3253+
directory_table_params: KeyValueOptions,
3254+
file_format: KeyValueOptions,
3255+
copy_options: KeyValueOptions,
32463256
comment: Option<String>,
32473257
},
32483258
/// ```sql
@@ -4467,6 +4477,29 @@ impl fmt::Display for Statement {
44674477
}
44684478
Ok(())
44694479
}
4480+
Statement::AlterSession {
4481+
set,
4482+
session_params,
4483+
} => {
4484+
write!(
4485+
f,
4486+
"ALTER SESSION {set}",
4487+
set = if *set { "SET" } else { "UNSET" }
4488+
)?;
4489+
if !session_params.options.is_empty() {
4490+
if *set {
4491+
write!(f, " {}", session_params)?;
4492+
} else {
4493+
let options = session_params
4494+
.options
4495+
.iter()
4496+
.map(|p| p.option_name.clone())
4497+
.collect::<Vec<_>>();
4498+
write!(f, " {}", display_separated(&options, ", "))?;
4499+
}
4500+
}
4501+
Ok(())
4502+
}
44704503
Statement::Drop {
44714504
object_type,
44724505
if_exists,

‎src/ast/spans.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -430,6 +430,7 @@ impl Spanned for Statement {
430430
// These statements need to be implemented
431431
Statement::AlterType { .. } => Span::empty(),
432432
Statement::AlterRole { .. } => Span::empty(),
433+
Statement::AlterSession { .. } => Span::empty(),
433434
Statement::AttachDatabase { .. } => Span::empty(),
434435
Statement::AttachDuckDBDatabase { .. } => Span::empty(),
435436
Statement::DetachDuckDBDatabase { .. } => Span::empty(),

0 commit comments

Comments
 (0)