Snowflake: parse top-level UNSET (single / parenthesized session variables)

Author: sabir-akhadov-localstack

src/ast/mod.rs

@@ -3797,6 +3797,16 @@ pub enum Statement {
     Analyze(Analyze),
     /// `SET` statements (session, transaction, timezone, etc.).
     Set(Set),
+    /// Snowflake `UNSET` of one or more session variables.
+    ///
+    /// ```sql
+    /// UNSET v
+    /// UNSET (v1, v2)
+    /// ```
+    Unset {
+        /// Session variable names to remove.
+        names: Vec<ObjectName>,
+    },
     /// ```sql
     /// TRUNCATE
     /// ```
@@ -6377,6 +6387,13 @@ impl fmt::Display for Statement {
                 Ok(())
             }
             Self::Set(set) => write!(f, "{set}"),
+            Statement::Unset { names } => {
+                if names.len() == 1 {
+                    write!(f, "UNSET {}", names[0])
+                } else {
+                    write!(f, "UNSET ({})", display_comma_separated(names))
+                }
+            }
             Statement::ShowVariable { variable } => {
                 write!(f, "SHOW")?;
                 if !variable.is_empty() {

src/ast/spans.rs

@@ -439,6 +439,7 @@ impl Spanned for Statement {
             Statement::Flush { .. } => Span::empty(),
             Statement::Discard { .. } => Span::empty(),
             Statement::Set(_) => Span::empty(),
+            Statement::Unset { .. } => Span::empty(),
             Statement::ShowFunctions { .. } => Span::empty(),
             Statement::ShowVariable { .. } => Span::empty(),
             Statement::ShowStatus { .. } => Span::empty(),

src/parser/mod.rs

@@ -667,6 +667,7 @@ impl<'a> Parser<'a> {
                 }
                 Keyword::CLOSE => self.parse_close(),
                 Keyword::SET => self.parse_set(),
+                Keyword::UNSET if dialect_of!(self is SnowflakeDialect) => self.parse_unset(),
                 Keyword::SHOW => self.parse_show(),
                 Keyword::USE => self.parse_use(),
                 Keyword::GRANT => self.parse_grant().map(Into::into),
@@ -15740,6 +15741,20 @@ impl<'a> Parser<'a> {
         Ok(SetAssignment { scope, name, value })
     }
 
+    /// Parse a Snowflake top-level `UNSET v` / `UNSET (v1, v2)` statement.
+    fn parse_unset(&mut self) -> Result<Statement, ParserError> {
+        let names = if self.consume_token(&Token::LParen) {
+            let names = self.parse_comma_separated(|parser: &mut Parser<'a>| {
+                Ok(ObjectName::from(vec![parser.parse_identifier()?]))
+            })?;
+            self.expect_token(&Token::RParen)?;
+            names
+        } else {
+            vec![self.parse_object_name(false)?]
+        };
+        Ok(Statement::Unset { names })
+    }
+
     fn parse_set(&mut self) -> Result<Statement, ParserError> {
         let hivevar = self.parse_keyword(Keyword::HIVEVAR);