Skip to content

Commit 199d1be

Browse files
yoavcloudayman-sigma
yoavcloud
authored and
ayman-sigma
committed
MsSQL SET for session params (apache#1646)
1 parent f661dbc commit 199d1be

File tree

7 files changed

+271
-0
lines changed

7 files changed

+271
-0
lines changed

‎src/ast/mod.rs

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3455,6 +3455,10 @@ pub enum Statement {
34553455
/// Snowflake `REMOVE`
34563456
/// See: <https://docs.snowflake.com/en/sql-reference/sql/remove>
34573457
Remove(FileStagingCommand),
3458+
/// MS-SQL session
3459+
///
3460+
/// See <https://learn.microsoft.com/en-us/sql/t-sql/statements/set-statements-transact-sql>
3461+
SetSessionParam(SetSessionParamKind),
34583462
}
34593463

34603464
impl fmt::Display for Statement {
@@ -5042,6 +5046,7 @@ impl fmt::Display for Statement {
50425046
}
50435047
Statement::List(command) => write!(f, "LIST {command}"),
50445048
Statement::Remove(command) => write!(f, "REMOVE {command}"),
5049+
Statement::SetSessionParam(kind) => write!(f, "SET {kind}"),
50455050
}
50465051
}
50475052
}
@@ -6459,6 +6464,7 @@ pub enum TransactionIsolationLevel {
64596464
ReadCommitted,
64606465
RepeatableRead,
64616466
Serializable,
6467+
Snapshot,
64626468
}
64636469

64646470
impl fmt::Display for TransactionIsolationLevel {
@@ -6469,6 +6475,7 @@ impl fmt::Display for TransactionIsolationLevel {
64696475
ReadCommitted => "READ COMMITTED",
64706476
RepeatableRead => "REPEATABLE READ",
64716477
Serializable => "SERIALIZABLE",
6478+
Snapshot => "SNAPSHOT",
64726479
})
64736480
}
64746481
}
@@ -7955,6 +7962,126 @@ impl fmt::Display for TableObject {
79557962
}
79567963
}
79577964

7965+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
7966+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7967+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
7968+
pub enum SetSessionParamKind {
7969+
Generic(SetSessionParamGeneric),
7970+
IdentityInsert(SetSessionParamIdentityInsert),
7971+
Offsets(SetSessionParamOffsets),
7972+
Statistics(SetSessionParamStatistics),
7973+
}
7974+
7975+
impl fmt::Display for SetSessionParamKind {
7976+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
7977+
match self {
7978+
SetSessionParamKind::Generic(x) => write!(f, "{x}"),
7979+
SetSessionParamKind::IdentityInsert(x) => write!(f, "{x}"),
7980+
SetSessionParamKind::Offsets(x) => write!(f, "{x}"),
7981+
SetSessionParamKind::Statistics(x) => write!(f, "{x}"),
7982+
}
7983+
}
7984+
}
7985+
7986+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
7987+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7988+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
7989+
pub struct SetSessionParamGeneric {
7990+
pub names: Vec<String>,
7991+
pub value: String,
7992+
}
7993+
7994+
impl fmt::Display for SetSessionParamGeneric {
7995+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
7996+
write!(f, "{} {}", display_comma_separated(&self.names), self.value)
7997+
}
7998+
}
7999+
8000+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
8001+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8002+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
8003+
pub struct SetSessionParamIdentityInsert {
8004+
pub obj: ObjectName,
8005+
pub value: SessionParamValue,
8006+
}
8007+
8008+
impl fmt::Display for SetSessionParamIdentityInsert {
8009+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
8010+
write!(f, "IDENTITY_INSERT {} {}", self.obj, self.value)
8011+
}
8012+
}
8013+
8014+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
8015+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8016+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
8017+
pub struct SetSessionParamOffsets {
8018+
pub keywords: Vec<String>,
8019+
pub value: SessionParamValue,
8020+
}
8021+
8022+
impl fmt::Display for SetSessionParamOffsets {
8023+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
8024+
write!(
8025+
f,
8026+
"OFFSETS {} {}",
8027+
display_comma_separated(&self.keywords),
8028+
self.value
8029+
)
8030+
}
8031+
}
8032+
8033+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
8034+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8035+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
8036+
pub struct SetSessionParamStatistics {
8037+
pub topic: SessionParamStatsTopic,
8038+
pub value: SessionParamValue,
8039+
}
8040+
8041+
impl fmt::Display for SetSessionParamStatistics {
8042+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
8043+
write!(f, "STATISTICS {} {}", self.topic, self.value)
8044+
}
8045+
}
8046+
8047+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
8048+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8049+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
8050+
pub enum SessionParamStatsTopic {
8051+
IO,
8052+
Profile,
8053+
Time,
8054+
Xml,
8055+
}
8056+
8057+
impl fmt::Display for SessionParamStatsTopic {
8058+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
8059+
match self {
8060+
SessionParamStatsTopic::IO => write!(f, "IO"),
8061+
SessionParamStatsTopic::Profile => write!(f, "PROFILE"),
8062+
SessionParamStatsTopic::Time => write!(f, "TIME"),
8063+
SessionParamStatsTopic::Xml => write!(f, "XML"),
8064+
}
8065+
}
8066+
}
8067+
8068+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
8069+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8070+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
8071+
pub enum SessionParamValue {
8072+
On,
8073+
Off,
8074+
}
8075+
8076+
impl fmt::Display for SessionParamValue {
8077+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
8078+
match self {
8079+
SessionParamValue::On => write!(f, "ON"),
8080+
SessionParamValue::Off => write!(f, "OFF"),
8081+
}
8082+
}
8083+
}
8084+
79588085
#[cfg(test)]
79598086
mod tests {
79608087
use super::*;

‎src/ast/spans.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,7 @@ impl Spanned for Statement {
496496
Statement::UNLISTEN { .. } => Span::empty(),
497497
Statement::RenameTable { .. } => Span::empty(),
498498
Statement::List(..) | Statement::Remove(..) => Span::empty(),
499+
Statement::SetSessionParam { .. } => Span::empty(),
499500
}
500501
}
501502
}

‎src/dialect/mod.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -802,6 +802,12 @@ pub trait Dialect: Debug + Any {
802802
fn supports_insert_format(&self) -> bool {
803803
false
804804
}
805+
806+
/// Returns true if this dialect supports `SET` statements without an explicit
807+
/// assignment operator such as `=`. For example: `SET SHOWPLAN_XML ON`.
808+
fn supports_set_stmt_without_operator(&self) -> bool {
809+
false
810+
}
805811
}
806812

807813
/// This represents the operators for which precedence must be defined

‎src/dialect/mssql.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,9 @@ impl Dialect for MsSqlDialect {
8585
fn supports_end_transaction_modifier(&self) -> bool {
8686
true
8787
}
88+
89+
/// See: <https://learn.microsoft.com/en-us/sql/t-sql/statements/set-statements-transact-sql>
90+
fn supports_set_stmt_without_operator(&self) -> bool {
91+
true
92+
}
8893
}

‎src/keywords.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,7 @@ define_keywords!(
388388
HOURS,
389389
ID,
390390
IDENTITY,
391+
IDENTITY_INSERT,
391392
IF,
392393
IGNORE,
393394
ILIKE,
@@ -426,6 +427,7 @@ define_keywords!(
426427
INTERVAL,
427428
INTO,
428429
INVOKER,
430+
IO,
429431
IS,
430432
ISODOW,
431433
ISOLATION,
@@ -557,7 +559,9 @@ define_keywords!(
557559
OCTETS,
558560
OCTET_LENGTH,
559561
OF,
562+
OFF,
560563
OFFSET,
564+
OFFSETS,
561565
OLD,
562566
OMIT,
563567
ON,
@@ -623,6 +627,7 @@ define_keywords!(
623627
PRIOR,
624628
PRIVILEGES,
625629
PROCEDURE,
630+
PROFILE,
626631
PROGRAM,
627632
PROJECTION,
628633
PUBLIC,

‎src/parser/mod.rs

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10438,11 +10438,75 @@ impl<'a> Parser<'a> {
1043810438
snapshot: None,
1043910439
session: false,
1044010440
})
10441+
} else if self.dialect.supports_set_stmt_without_operator() {
10442+
self.prev_token();
10443+
self.parse_set_session_params()
1044110444
} else {
1044210445
self.expected("equals sign or TO", self.peek_token())
1044310446
}
1044410447
}
1044510448

10449+
pub fn parse_set_session_params(&mut self) -> Result<Statement, ParserError> {
10450+
if self.parse_keyword(Keyword::STATISTICS) {
10451+
let topic = match self.parse_one_of_keywords(&[
10452+
Keyword::IO,
10453+
Keyword::PROFILE,
10454+
Keyword::TIME,
10455+
Keyword::XML,
10456+
]) {
10457+
Some(Keyword::IO) => SessionParamStatsTopic::IO,
10458+
Some(Keyword::PROFILE) => SessionParamStatsTopic::Profile,
10459+
Some(Keyword::TIME) => SessionParamStatsTopic::Time,
10460+
Some(Keyword::XML) => SessionParamStatsTopic::Xml,
10461+
_ => return self.expected("IO, PROFILE, TIME or XML", self.peek_token()),
10462+
};
10463+
let value = self.parse_session_param_value()?;
10464+
Ok(Statement::SetSessionParam(SetSessionParamKind::Statistics(
10465+
SetSessionParamStatistics { topic, value },
10466+
)))
10467+
} else if self.parse_keyword(Keyword::IDENTITY_INSERT) {
10468+
let obj = self.parse_object_name(false)?;
10469+
let value = self.parse_session_param_value()?;
10470+
Ok(Statement::SetSessionParam(
10471+
SetSessionParamKind::IdentityInsert(SetSessionParamIdentityInsert { obj, value }),
10472+
))
10473+
} else if self.parse_keyword(Keyword::OFFSETS) {
10474+
let keywords = self.parse_comma_separated(|parser| {
10475+
let next_token = parser.next_token();
10476+
match &next_token.token {
10477+
Token::Word(w) => Ok(w.to_string()),
10478+
_ => parser.expected("SQL keyword", next_token),
10479+
}
10480+
})?;
10481+
let value = self.parse_session_param_value()?;
10482+
Ok(Statement::SetSessionParam(SetSessionParamKind::Offsets(
10483+
SetSessionParamOffsets { keywords, value },
10484+
)))
10485+
} else {
10486+
let names = self.parse_comma_separated(|parser| {
10487+
let next_token = parser.next_token();
10488+
match next_token.token {
10489+
Token::Word(w) => Ok(w.to_string()),
10490+
_ => parser.expected("Session param name", next_token),
10491+
}
10492+
})?;
10493+
let value = self.parse_expr()?.to_string();
10494+
Ok(Statement::SetSessionParam(SetSessionParamKind::Generic(
10495+
SetSessionParamGeneric { names, value },
10496+
)))
10497+
}
10498+
}
10499+
10500+
fn parse_session_param_value(&mut self) -> Result<SessionParamValue, ParserError> {
10501+
if self.parse_keyword(Keyword::ON) {
10502+
Ok(SessionParamValue::On)
10503+
} else if self.parse_keyword(Keyword::OFF) {
10504+
Ok(SessionParamValue::Off)
10505+
} else {
10506+
self.expected("ON or OFF", self.peek_token())
10507+
}
10508+
}
10509+
1044610510
pub fn parse_show(&mut self) -> Result<Statement, ParserError> {
1044710511
let terse = self.parse_keyword(Keyword::TERSE);
1044810512
let extended = self.parse_keyword(Keyword::EXTENDED);
@@ -13014,6 +13078,8 @@ impl<'a> Parser<'a> {
1301413078
TransactionIsolationLevel::RepeatableRead
1301513079
} else if self.parse_keyword(Keyword::SERIALIZABLE) {
1301613080
TransactionIsolationLevel::Serializable
13081+
} else if self.parse_keyword(Keyword::SNAPSHOT) {
13082+
TransactionIsolationLevel::Snapshot
1301713083
} else {
1301813084
self.expected("isolation level", self.peek_token())?
1301913085
};

‎tests/sqlparser_mssql.rs

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1679,6 +1679,67 @@ fn parse_true_false_as_identifiers() {
16791679
);
16801680
}
16811681

1682+
#[test]
1683+
fn parse_mssql_set_session_value() {
1684+
ms().verified_stmt(
1685+
"SET OFFSETS SELECT, FROM, ORDER, TABLE, PROCEDURE, STATEMENT, PARAM, EXECUTE ON",
1686+
);
1687+
ms().verified_stmt("SET IDENTITY_INSERT dbo.Tool ON");
1688+
ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED");
1689+
ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL READ COMMITTED");
1690+
ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ");
1691+
ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL SNAPSHOT");
1692+
ms().verified_stmt("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE");
1693+
ms().verified_stmt("SET STATISTICS IO ON");
1694+
ms().verified_stmt("SET STATISTICS XML ON");
1695+
ms().verified_stmt("SET STATISTICS PROFILE ON");
1696+
ms().verified_stmt("SET STATISTICS TIME ON");
1697+
ms().verified_stmt("SET DATEFIRST 7");
1698+
ms().verified_stmt("SET DATEFIRST @xxx");
1699+
ms().verified_stmt("SET DATEFIRST @@xxx");
1700+
ms().verified_stmt("SET DATEFORMAT dmy");
1701+
ms().verified_stmt("SET DATEFORMAT @datevar");
1702+
ms().verified_stmt("SET DATEFORMAT @@datevar");
1703+
ms().verified_stmt("SET DEADLOCK_PRIORITY 'LOW'");
1704+
ms().verified_stmt("SET DEADLOCK_PRIORITY LOW");
1705+
ms().verified_stmt("SET DEADLOCK_PRIORITY 8");
1706+
ms().verified_stmt("SET DEADLOCK_PRIORITY -8");
1707+
ms().verified_stmt("SET DEADLOCK_PRIORITY @xxx");
1708+
ms().verified_stmt("SET DEADLOCK_PRIORITY @@xxx");
1709+
ms().verified_stmt("SET LOCK_TIMEOUT 1800");
1710+
ms().verified_stmt("SET CONCAT_NULL_YIELDS_NULL ON");
1711+
ms().verified_stmt("SET CURSOR_CLOSE_ON_COMMIT ON");
1712+
ms().verified_stmt("SET FIPS_FLAGGER 'level'");
1713+
ms().verified_stmt("SET FIPS_FLAGGER OFF");
1714+
ms().verified_stmt("SET LANGUAGE Italian");
1715+
ms().verified_stmt("SET QUOTED_IDENTIFIER ON");
1716+
ms().verified_stmt("SET ARITHABORT ON");
1717+
ms().verified_stmt("SET ARITHIGNORE OFF");
1718+
ms().verified_stmt("SET FMTONLY ON");
1719+
ms().verified_stmt("SET NOCOUNT OFF");
1720+
ms().verified_stmt("SET NOEXEC ON");
1721+
ms().verified_stmt("SET NUMERIC_ROUNDABORT ON");
1722+
ms().verified_stmt("SET QUERY_GOVERNOR_COST_LIMIT 11");
1723+
ms().verified_stmt("SET ROWCOUNT 4");
1724+
ms().verified_stmt("SET ROWCOUNT @xxx");
1725+
ms().verified_stmt("SET ROWCOUNT @@xxx");
1726+
ms().verified_stmt("SET TEXTSIZE 11");
1727+
ms().verified_stmt("SET ANSI_DEFAULTS ON");
1728+
ms().verified_stmt("SET ANSI_NULL_DFLT_OFF ON");
1729+
ms().verified_stmt("SET ANSI_NULL_DFLT_ON ON");
1730+
ms().verified_stmt("SET ANSI_NULLS ON");
1731+
ms().verified_stmt("SET ANSI_PADDING ON");
1732+
ms().verified_stmt("SET ANSI_WARNINGS ON");
1733+
ms().verified_stmt("SET FORCEPLAN ON");
1734+
ms().verified_stmt("SET SHOWPLAN_ALL ON");
1735+
ms().verified_stmt("SET SHOWPLAN_TEXT ON");
1736+
ms().verified_stmt("SET SHOWPLAN_XML ON");
1737+
ms().verified_stmt("SET IMPLICIT_TRANSACTIONS ON");
1738+
ms().verified_stmt("SET REMOTE_PROC_TRANSACTIONS ON");
1739+
ms().verified_stmt("SET XACT_ABORT ON");
1740+
ms().verified_stmt("SET ANSI_NULLS, ANSI_PADDING ON");
1741+
}
1742+
16821743
fn ms() -> TestedDialects {
16831744
TestedDialects::new(vec![Box::new(MsSqlDialect {})])
16841745
}

0 commit comments

Comments
 (0)