MsSQL SET for session params (#1646)

src/ast/mod.rs

@@ -3437,6 +3437,10 @@ pub enum Statement {
     /// Snowflake `REMOVE`
     /// See: <https://docs.snowflake.com/en/sql-reference/sql/remove>
     Remove(FileStagingCommand),
+    /// MS-SQL session
+    ///
+    /// See <https://learn.microsoft.com/en-us/sql/t-sql/statements/set-statements-transact-sql>
+    SetSessionParam(SetSessionParamKind),
 }
 
 impl fmt::Display for Statement {
@@ -5024,6 +5028,7 @@ impl fmt::Display for Statement {
             }
             Statement::List(command) => write!(f, "LIST {command}"),
             Statement::Remove(command) => write!(f, "REMOVE {command}"),
+            Statement::SetSessionParam(kind) => write!(f, "SET {kind}"),
         }
     }
 }
@@ -6441,6 +6446,7 @@ pub enum TransactionIsolationLevel {
     ReadCommitted,
     RepeatableRead,
     Serializable,
+    Snapshot,
 }
 
 impl fmt::Display for TransactionIsolationLevel {
@@ -6451,6 +6457,7 @@ impl fmt::Display for TransactionIsolationLevel {
             ReadCommitted => "READ COMMITTED",
             RepeatableRead => "REPEATABLE READ",
             Serializable => "SERIALIZABLE",
+            Snapshot => "SNAPSHOT",
         })
     }
 }
@@ -7937,6 +7944,126 @@ impl fmt::Display for TableObject {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum SetSessionParamKind {
+    Generic(SetSessionParamGeneric),
+    IdentityInsert(SetSessionParamIdentityInsert),
+    Offsets(SetSessionParamOffsets),
+    Statistics(SetSessionParamStatistics),
+}
+
+impl fmt::Display for SetSessionParamKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SetSessionParamKind::Generic(x) => write!(f, "{x}"),
+            SetSessionParamKind::IdentityInsert(x) => write!(f, "{x}"),
+            SetSessionParamKind::Offsets(x) => write!(f, "{x}"),
+            SetSessionParamKind::Statistics(x) => write!(f, "{x}"),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SetSessionParamGeneric {
+    pub names: Vec<String>,
+    pub value: String,
+}
+
+impl fmt::Display for SetSessionParamGeneric {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{} {}", display_comma_separated(&self.names), self.value)
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SetSessionParamIdentityInsert {
+    pub obj: ObjectName,
+    pub value: SessionParamValue,
+}
+
+impl fmt::Display for SetSessionParamIdentityInsert {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "IDENTITY_INSERT {} {}", self.obj, self.value)
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SetSessionParamOffsets {
+    pub keywords: Vec<String>,
+    pub value: SessionParamValue,
+}
+
+impl fmt::Display for SetSessionParamOffsets {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "OFFSETS {} {}",
+            display_comma_separated(&self.keywords),
+            self.value
+        )
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct SetSessionParamStatistics {
+    pub topic: SessionParamStatsTopic,
+    pub value: SessionParamValue,
+}
+
+impl fmt::Display for SetSessionParamStatistics {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "STATISTICS {} {}", self.topic, self.value)
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum SessionParamStatsTopic {
+    IO,
+    Profile,
+    Time,
+    Xml,
+}
+
+impl fmt::Display for SessionParamStatsTopic {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SessionParamStatsTopic::IO => write!(f, "IO"),
+            SessionParamStatsTopic::Profile => write!(f, "PROFILE"),
+            SessionParamStatsTopic::Time => write!(f, "TIME"),
+            SessionParamStatsTopic::Xml => write!(f, "XML"),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum SessionParamValue {
+    On,
+    Off,
+}
+
+impl fmt::Display for SessionParamValue {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SessionParamValue::On => write!(f, "ON"),
+            SessionParamValue::Off => write!(f, "OFF"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/ast/spans.rs

@@ -496,6 +496,7 @@ impl Spanned for Statement {
             Statement::UNLISTEN { .. } => Span::empty(),
             Statement::RenameTable { .. } => Span::empty(),
             Statement::List(..) | Statement::Remove(..) => Span::empty(),
+            Statement::SetSessionParam { .. } => Span::empty(),
         }
     }
 }

src/dialect/mod.rs

@@ -802,6 +802,12 @@ pub trait Dialect: Debug + Any {
     fn supports_insert_format(&self) -> bool {
         false
     }
+
+    /// Returns true if this dialect supports `SET` statements without an explicit
+    /// assignment operator such as `=`. For example: `SET SHOWPLAN_XML ON`.
+    fn supports_set_stmt_without_operator(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/dialect/mssql.rs

@@ -85,4 +85,9 @@ impl Dialect for MsSqlDialect {
     fn supports_end_transaction_modifier(&self) -> bool {
         true
     }
+
+    /// See: <https://learn.microsoft.com/en-us/sql/t-sql/statements/set-statements-transact-sql>
+    fn supports_set_stmt_without_operator(&self) -> bool {
+        true
+    }
 }

src/keywords.rs

@@ -388,6 +388,7 @@ define_keywords!(
     HOURS,
     ID,
     IDENTITY,
+    IDENTITY_INSERT,
     IF,
     IGNORE,
     ILIKE,
@@ -426,6 +427,7 @@ define_keywords!(
     INTERVAL,
     INTO,
     INVOKER,
+    IO,
     IS,
     ISODOW,
     ISOLATION,
@@ -557,7 +559,9 @@ define_keywords!(
     OCTETS,
     OCTET_LENGTH,
     OF,
+    OFF,
     OFFSET,
+    OFFSETS,
     OLD,
     OMIT,
     ON,
@@ -623,6 +627,7 @@ define_keywords!(
     PRIOR,
     PRIVILEGES,
     PROCEDURE,
+    PROFILE,
     PROGRAM,
     PROJECTION,
     PUBLIC,

src/parser/mod.rs

@@ -10428,11 +10428,75 @@ impl<'a> Parser<'a> {
                 snapshot: None,
                 session: false,
             })
+        } else if self.dialect.supports_set_stmt_without_operator() {
+            self.prev_token();
+            self.parse_set_session_params()
         } else {
             self.expected("equals sign or TO", self.peek_token())
         }
     }
 
+    pub fn parse_set_session_params(&mut self) -> Result<Statement, ParserError> {
+        if self.parse_keyword(Keyword::STATISTICS) {
+            let topic = match self.parse_one_of_keywords(&[
+                Keyword::IO,
+                Keyword::PROFILE,
+                Keyword::TIME,
+                Keyword::XML,
+            ]) {
+                Some(Keyword::IO) => SessionParamStatsTopic::IO,
+                Some(Keyword::PROFILE) => SessionParamStatsTopic::Profile,
+                Some(Keyword::TIME) => SessionParamStatsTopic::Time,
+                Some(Keyword::XML) => SessionParamStatsTopic::Xml,
+                _ => return self.expected("IO, PROFILE, TIME or XML", self.peek_token()),
+            };
+            let value = self.parse_session_param_value()?;
+            Ok(Statement::SetSessionParam(SetSessionParamKind::Statistics(
+                SetSessionParamStatistics { topic, value },
+            )))
+        } else if self.parse_keyword(Keyword::IDENTITY_INSERT) {
+            let obj = self.parse_object_name(false)?;
+            let value = self.parse_session_param_value()?;
+            Ok(Statement::SetSessionParam(
+                SetSessionParamKind::IdentityInsert(SetSessionParamIdentityInsert { obj, value }),
+            ))
+        } else if self.parse_keyword(Keyword::OFFSETS) {
+            let keywords = self.parse_comma_separated(|parser| {
+                let next_token = parser.next_token();
+                match &next_token.token {
+                    Token::Word(w) => Ok(w.to_string()),
+                    _ => parser.expected("SQL keyword", next_token),
+                }
+            })?;
+            let value = self.parse_session_param_value()?;
+            Ok(Statement::SetSessionParam(SetSessionParamKind::Offsets(
+                SetSessionParamOffsets { keywords, value },
+            )))
+        } else {
+            let names = self.parse_comma_separated(|parser| {
+                let next_token = parser.next_token();
+                match next_token.token {
+                    Token::Word(w) => Ok(w.to_string()),
+                    _ => parser.expected("Session param name", next_token),
+                }
+            })?;
+            let value = self.parse_expr()?.to_string();
+            Ok(Statement::SetSessionParam(SetSessionParamKind::Generic(
+                SetSessionParamGeneric { names, value },
+            )))
+        }
+    }
+
+    fn parse_session_param_value(&mut self) -> Result<SessionParamValue, ParserError> {
+        if self.parse_keyword(Keyword::ON) {
+            Ok(SessionParamValue::On)
+        } else if self.parse_keyword(Keyword::OFF) {
+            Ok(SessionParamValue::Off)
+        } else {
+            self.expected("ON or OFF", self.peek_token())
+        }
+    }
+
     pub fn parse_show(&mut self) -> Result<Statement, ParserError> {
         let terse = self.parse_keyword(Keyword::TERSE);
         let extended = self.parse_keyword(Keyword::EXTENDED);
@@ -13004,6 +13068,8 @@ impl<'a> Parser<'a> {
                     TransactionIsolationLevel::RepeatableRead
                 } else if self.parse_keyword(Keyword::SERIALIZABLE) {
                     TransactionIsolationLevel::Serializable
+                } else if self.parse_keyword(Keyword::SNAPSHOT) {
+                    TransactionIsolationLevel::Snapshot
                 } else {
                     self.expected("isolation level", self.peek_token())?
                 };

tests/sqlparser_mssql.rs

@@ -1679,6 +1679,67 @@ fn parse_true_false_as_identifiers() {
     );
 }
 
+#[test]
+fn parse_mssql_set_session_value() {
+    ms().verified_stmt(
+        "SET OFFSETS SELECT, FROM, ORDER, TABLE, PROCEDURE, STATEMENT, PARAM, EXECUTE ON",
+    );
+    ms().verified_stmt("SET IDENTITY_INSERT dbo.Tool ON");
+    ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED");
+    ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL READ COMMITTED");
+    ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ");
+    ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL SNAPSHOT");
+    ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE");
+    ms().verified_stmt("SET STATISTICS IO ON");
+    ms().verified_stmt("SET STATISTICS XML ON");
+    ms().verified_stmt("SET STATISTICS PROFILE ON");
+    ms().verified_stmt("SET STATISTICS TIME ON");
+    ms().verified_stmt("SET DATEFIRST 7");
+    ms().verified_stmt("SET DATEFIRST @xxx");
+    ms().verified_stmt("SET DATEFIRST @@xxx");
+    ms().verified_stmt("SET DATEFORMAT dmy");
+    ms().verified_stmt("SET DATEFORMAT @datevar");
+    ms().verified_stmt("SET DATEFORMAT @@datevar");
+    ms().verified_stmt("SET DEADLOCK_PRIORITY 'LOW'");
+    ms().verified_stmt("SET DEADLOCK_PRIORITY LOW");
+    ms().verified_stmt("SET DEADLOCK_PRIORITY 8");
+    ms().verified_stmt("SET DEADLOCK_PRIORITY -8");
+    ms().verified_stmt("SET DEADLOCK_PRIORITY @xxx");
+    ms().verified_stmt("SET DEADLOCK_PRIORITY @@xxx");
+    ms().verified_stmt("SET LOCK_TIMEOUT 1800");
+    ms().verified_stmt("SET CONCAT_NULL_YIELDS_NULL ON");
+    ms().verified_stmt("SET CURSOR_CLOSE_ON_COMMIT ON");
+    ms().verified_stmt("SET FIPS_FLAGGER 'level'");
+    ms().verified_stmt("SET FIPS_FLAGGER OFF");
+    ms().verified_stmt("SET LANGUAGE Italian");
+    ms().verified_stmt("SET QUOTED_IDENTIFIER ON");
+    ms().verified_stmt("SET ARITHABORT ON");
+    ms().verified_stmt("SET ARITHIGNORE OFF");
+    ms().verified_stmt("SET FMTONLY ON");
+    ms().verified_stmt("SET NOCOUNT OFF");
+    ms().verified_stmt("SET NOEXEC ON");
+    ms().verified_stmt("SET NUMERIC_ROUNDABORT ON");
+    ms().verified_stmt("SET QUERY_GOVERNOR_COST_LIMIT 11");
+    ms().verified_stmt("SET ROWCOUNT 4");
+    ms().verified_stmt("SET ROWCOUNT @xxx");
+    ms().verified_stmt("SET ROWCOUNT @@xxx");
+    ms().verified_stmt("SET TEXTSIZE 11");
+    ms().verified_stmt("SET ANSI_DEFAULTS ON");
+    ms().verified_stmt("SET ANSI_NULL_DFLT_OFF ON");
+    ms().verified_stmt("SET ANSI_NULL_DFLT_ON ON");
+    ms().verified_stmt("SET ANSI_NULLS ON");
+    ms().verified_stmt("SET ANSI_PADDING ON");
+    ms().verified_stmt("SET ANSI_WARNINGS ON");
+    ms().verified_stmt("SET FORCEPLAN ON");
+    ms().verified_stmt("SET SHOWPLAN_ALL ON");
+    ms().verified_stmt("SET SHOWPLAN_TEXT ON");
+    ms().verified_stmt("SET SHOWPLAN_XML ON");
+    ms().verified_stmt("SET IMPLICIT_TRANSACTIONS ON");
+    ms().verified_stmt("SET REMOTE_PROC_TRANSACTIONS ON");
+    ms().verified_stmt("SET XACT_ABORT ON");
+    ms().verified_stmt("SET ANSI_NULLS, ANSI_PADDING ON");
+}
+
 fn ms() -> TestedDialects {
     TestedDialects::new(vec![Box::new(MsSqlDialect {})])
 }