Skip to content

Commit

Permalink
MsSQL SET for session params (#1646)
Browse files Browse the repository at this point in the history
  • Loading branch information
@yoavcloud
yoavcloud authored Jan 11, 2025
1 parent 0c3b6c0 commit 3b4dc0f
Show file tree
Hide file tree
Showing 7 changed files with 271 additions and 0 deletions.
127 changes: 127 additions & 0 deletions src/ast/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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}"),
}
}
}
Expand Down Expand Up @@ -6441,6 +6446,7 @@ pub enum TransactionIsolationLevel {
ReadCommitted,
RepeatableRead,
Serializable,
Snapshot,
}

impl fmt::Display for TransactionIsolationLevel {
Expand All @@ -6451,6 +6457,7 @@ impl fmt::Display for TransactionIsolationLevel {
ReadCommitted => "READ COMMITTED",
RepeatableRead => "REPEATABLE READ",
Serializable => "SERIALIZABLE",
Snapshot => "SNAPSHOT",
})
}
}
Expand Down Expand Up @@ -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::*;
Expand Down
1 change: 1 addition & 0 deletions src/ast/spans.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
}
Expand Down
6 changes: 6 additions & 0 deletions src/dialect/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions src/dialect/mssql.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
5 changes: 5 additions & 0 deletions src/keywords.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -388,6 +388,7 @@ define_keywords!(
HOURS,
ID,
IDENTITY,
IDENTITY_INSERT,
IF,
IGNORE,
ILIKE,
Expand Down Expand Up @@ -426,6 +427,7 @@ define_keywords!(
INTERVAL,
INTO,
INVOKER,
IO,
IS,
ISODOW,
ISOLATION,
Expand Down Expand Up @@ -557,7 +559,9 @@ define_keywords!(
OCTETS,
OCTET_LENGTH,
OF,
OFF,
OFFSET,
OFFSETS,
OLD,
OMIT,
ON,
Expand Down Expand Up @@ -623,6 +627,7 @@ define_keywords!(
PRIOR,
PRIVILEGES,
PROCEDURE,
PROFILE,
PROGRAM,
PROJECTION,
PUBLIC,
Expand Down
66 changes: 66 additions & 0 deletions src/parser/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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())?
};
Expand Down
61 changes: 61 additions & 0 deletions tests/sqlparser_mssql.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {})])
}
Expand Down

0 comments on commit 3b4dc0f

Please sign in to comment.