Snowflake: ALTER USER and KeyValueOptions Refactoring (apache#2035)

Author: yoavcloud

src/ast/dml.rs

@@ -16,12 +16,7 @@
 // under the License.
 
 #[cfg(not(feature = "std"))]
-use alloc::{
-    boxed::Box,
-    format,
-    string::{String, ToString},
-    vec::Vec,
-};
+use alloc::{boxed::Box, format, string::ToString, vec::Vec};
 
 use core::fmt::{self, Display};
 #[cfg(feature = "serde")]

src/ast/helpers/key_value_options.rs

@@ -19,9 +19,7 @@
 //! 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 alloc::{boxed::Box, string::String, vec::Vec};
 use core::fmt;
 use core::fmt::Formatter;
 
@@ -31,7 +29,7 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
-use crate::ast::display_separated;
+use crate::ast::{display_comma_separated, display_separated, Value};
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -52,20 +50,23 @@ pub enum KeyValueOptionsDelimiter {
 #[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,
+pub struct KeyValueOption {
+    pub option_name: String,
+    pub option_value: KeyValueOptionKind,
 }
 
+/// An option can have a single value, multiple values or a nested list of values.
+///
+/// A value can be numeric, boolean, etc. Enum-style values are represented
+/// as Value::Placeholder. For example: MFA_METHOD=SMS will be represented as
+/// `Value::Placeholder("SMS".to_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 KeyValueOption {
-    pub option_name: String,
-    pub option_type: KeyValueOptionType,
-    pub value: String,
+pub enum KeyValueOptionKind {
+    Single(Value),
+    Multi(Vec<Value>),
+    KeyValueOptions(Box<KeyValueOptions>),
 }
 
 impl fmt::Display for KeyValueOptions {
@@ -80,12 +81,20 @@ impl fmt::Display for KeyValueOptions {
 
 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)?;
+        match &self.option_value {
+            KeyValueOptionKind::Single(value) => {
+                write!(f, "{}={value}", self.option_name)?;
+            }
+            KeyValueOptionKind::Multi(values) => {
+                write!(
+                    f,
+                    "{}=({})",
+                    self.option_name,
+                    display_comma_separated(values)
+                )?;
             }
-            KeyValueOptionType::ENUM | KeyValueOptionType::BOOLEAN | KeyValueOptionType::NUMBER => {
-                write!(f, "{}={}", self.option_name, self.value)?;
+            KeyValueOptionKind::KeyValueOptions(options) => {
+                write!(f, "{}=({options})", self.option_name)?;
             }
         }
         Ok(())

src/ast/helpers/stmt_create_database.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 #[cfg(not(feature = "std"))]
-use alloc::{boxed::Box, format, string::String, vec, vec::Vec};
+use alloc::{format, string::String, vec::Vec};
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};

src/ast/mod.rs

@@ -4310,6 +4310,11 @@ pub enum Statement {
     /// ```
     /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/create-user)
     CreateUser(CreateUser),
+    /// ```sql
+    /// ALTER USER \[ IF EXISTS \] \[ <name> \]
+    /// ```
+    /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/alter-user)
+    AlterUser(AlterUser),
     /// Re-sorts rows and reclaims space in either a specified table or all tables in the current database
     ///
     /// ```sql
@@ -6183,6 +6188,7 @@ impl fmt::Display for Statement {
             Statement::CreateUser(s) => write!(f, "{s}"),
             Statement::AlterSchema(s) => write!(f, "{s}"),
             Statement::Vacuum(s) => write!(f, "{s}"),
+            Statement::AlterUser(s) => write!(f, "{s}"),
         }
     }
 }
@@ -10558,6 +10564,217 @@ impl fmt::Display for CreateUser {
     }
 }
 
+/// Modifies the properties of a user
+///
+/// Syntax:
+/// ```sql
+/// ALTER USER [ IF EXISTS ] [ <name> ] [ OPTIONS ]
+/// ```
+///
+/// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/alter-user)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterUser {
+    pub if_exists: bool,
+    pub name: Ident,
+    /// The following fields are Snowflake-specific: <https://docs.snowflake.com/en/sql-reference/sql/alter-user#syntax>
+    pub rename_to: Option<Ident>,
+    pub reset_password: bool,
+    pub abort_all_queries: bool,
+    pub add_role_delegation: Option<AlterUserAddRoleDelegation>,
+    pub remove_role_delegation: Option<AlterUserRemoveRoleDelegation>,
+    pub enroll_mfa: bool,
+    pub set_default_mfa_method: Option<MfaMethodKind>,
+    pub remove_mfa_method: Option<MfaMethodKind>,
+    pub modify_mfa_method: Option<AlterUserModifyMfaMethod>,
+    pub add_mfa_method_otp: Option<AlterUserAddMfaMethodOtp>,
+    pub set_policy: Option<AlterUserSetPolicy>,
+    pub unset_policy: Option<UserPolicyKind>,
+    pub set_tag: KeyValueOptions,
+    pub unset_tag: Vec<String>,
+    pub set_props: KeyValueOptions,
+    pub unset_props: Vec<String>,
+}
+
+/// ```sql
+/// ALTER USER [ IF EXISTS ] [ <name> ] ADD DELEGATED AUTHORIZATION OF ROLE <role_name> TO SECURITY INTEGRATION <integration_name>
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterUserAddRoleDelegation {
+    pub role: Ident,
+    pub integration: Ident,
+}
+
+/// ```sql
+/// ALTER USER [ IF EXISTS ] [ <name> ] REMOVE DELEGATED { AUTHORIZATION OF ROLE <role_name> | AUTHORIZATIONS } FROM SECURITY INTEGRATION <integration_name>
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterUserRemoveRoleDelegation {
+    pub role: Option<Ident>,
+    pub integration: Ident,
+}
+
+/// ```sql
+/// ADD MFA METHOD OTP [ COUNT = number ]
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterUserAddMfaMethodOtp {
+    pub count: Option<Value>,
+}
+
+/// ```sql
+/// ALTER USER [ IF EXISTS ] [ <name> ] MODIFY MFA METHOD <mfa_method> SET COMMENT = '<string>'
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterUserModifyMfaMethod {
+    pub method: MfaMethodKind,
+    pub comment: String,
+}
+
+/// Types of MFA methods
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum MfaMethodKind {
+    PassKey,
+    Totp,
+    Duo,
+}
+
+impl fmt::Display for MfaMethodKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            MfaMethodKind::PassKey => write!(f, "PASSKEY"),
+            MfaMethodKind::Totp => write!(f, "TOTP"),
+            MfaMethodKind::Duo => write!(f, "DUO"),
+        }
+    }
+}
+
+/// ```sql
+/// ALTER USER [ IF EXISTS ] [ <name> ] SET { AUTHENTICATION | PASSWORD | SESSION } POLICY <policy_name>
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterUserSetPolicy {
+    pub policy_kind: UserPolicyKind,
+    pub policy: Ident,
+}
+
+/// Types of user-based policies
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum UserPolicyKind {
+    Authentication,
+    Password,
+    Session,
+}
+
+impl fmt::Display for UserPolicyKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            UserPolicyKind::Authentication => write!(f, "AUTHENTICATION"),
+            UserPolicyKind::Password => write!(f, "PASSWORD"),
+            UserPolicyKind::Session => write!(f, "SESSION"),
+        }
+    }
+}
+
+impl fmt::Display for AlterUser {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "ALTER")?;
+        write!(f, " USER")?;
+        if self.if_exists {
+            write!(f, " IF EXISTS")?;
+        }
+        write!(f, " {}", self.name)?;
+        if let Some(new_name) = &self.rename_to {
+            write!(f, " RENAME TO {new_name}")?;
+        }
+        if self.reset_password {
+            write!(f, " RESET PASSWORD")?;
+        }
+        if self.abort_all_queries {
+            write!(f, " ABORT ALL QUERIES")?;
+        }
+        if let Some(role_delegation) = &self.add_role_delegation {
+            let role = &role_delegation.role;
+            let integration = &role_delegation.integration;
+            write!(
+                f,
+                " ADD DELEGATED AUTHORIZATION OF ROLE {role} TO SECURITY INTEGRATION {integration}"
+            )?;
+        }
+        if let Some(role_delegation) = &self.remove_role_delegation {
+            write!(f, " REMOVE DELEGATED")?;
+            match &role_delegation.role {
+                Some(role) => write!(f, " AUTHORIZATION OF ROLE {role}")?,
+                None => write!(f, " AUTHORIZATIONS")?,
+            }
+            let integration = &role_delegation.integration;
+            write!(f, " FROM SECURITY INTEGRATION {integration}")?;
+        }
+        if self.enroll_mfa {
+            write!(f, " ENROLL MFA")?;
+        }
+        if let Some(method) = &self.set_default_mfa_method {
+            write!(f, " SET DEFAULT_MFA_METHOD {method}")?
+        }
+        if let Some(method) = &self.remove_mfa_method {
+            write!(f, " REMOVE MFA METHOD {method}")?;
+        }
+        if let Some(modify) = &self.modify_mfa_method {
+            let method = &modify.method;
+            let comment = &modify.comment;
+            write!(
+                f,
+                " MODIFY MFA METHOD {method} SET COMMENT '{}'",
+                value::escape_single_quote_string(comment)
+            )?;
+        }
+        if let Some(add_mfa_method_otp) = &self.add_mfa_method_otp {
+            write!(f, " ADD MFA METHOD OTP")?;
+            if let Some(count) = &add_mfa_method_otp.count {
+                write!(f, " COUNT = {count}")?;
+            }
+        }
+        if let Some(policy) = &self.set_policy {
+            let policy_kind = &policy.policy_kind;
+            let name = &policy.policy;
+            write!(f, " SET {policy_kind} POLICY {name}")?;
+        }
+        if let Some(policy_kind) = &self.unset_policy {
+            write!(f, " UNSET {policy_kind} POLICY")?;
+        }
+        if !self.set_tag.options.is_empty() {
+            write!(f, " SET TAG {}", self.set_tag)?;
+        }
+        if !self.unset_tag.is_empty() {
+            write!(f, " UNSET TAG {}", display_comma_separated(&self.unset_tag))?;
+        }
+        let has_props = !self.set_props.options.is_empty();
+        if has_props {
+            write!(f, " SET")?;
+            write!(f, " {}", &self.set_props)?;
+        }
+        if !self.unset_props.is_empty() {
+            write!(f, " UNSET {}", display_comma_separated(&self.unset_props))?;
+        }
+        Ok(())
+    }
+}
+
 /// Specifies how to create a new table based on an existing table's schema.
 /// '''sql
 /// CREATE TABLE new LIKE old ...

src/ast/spans.rs

@@ -555,6 +555,7 @@ impl Spanned for Statement {
             Statement::CreateUser(..) => Span::empty(),
             Statement::AlterSchema(s) => s.span(),
             Statement::Vacuum(..) => Span::empty(),
+            Statement::AlterUser(..) => Span::empty(),
         }
     }
 }