Add support for MSSQL IF/ELSE statements. (#1791)

Co-authored-by: Roman Borschel <roman@cluvio.com>

src/ast/mod.rs

@@ -37,7 +37,8 @@ use serde::{Deserialize, Serialize};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
-use crate::tokenizer::Span;
+use crate::keywords::Keyword;
+use crate::tokenizer::{Span, Token};
 
 pub use self::data_type::{
     ArrayElemTypeDef, BinaryLength, CharLengthUnits, CharacterLength, DataType, EnumMember,
@@ -2118,20 +2119,23 @@ pub enum Password {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct CaseStatement {
+    /// The `CASE` token that starts the statement.
+    pub case_token: AttachedToken,
     pub match_expr: Option<Expr>,
-    pub when_blocks: Vec<ConditionalStatements>,
+    pub when_blocks: Vec<ConditionalStatementBlock>,
-    pub else_block: Option<Vec<Statement>>,
+    pub else_block: Option<ConditionalStatementBlock>,
-    /// TRUE if the statement ends with `END CASE` (vs `END`).
+    /// The last token of the statement (`END` or `CASE`).
-    pub has_end_case: bool,
+    pub end_case_token: AttachedToken,
 }
 
 impl fmt::Display for CaseStatement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let CaseStatement {
+            case_token: _,
             match_expr,
             when_blocks,
             else_block,
-            has_end_case,
+            end_case_token: AttachedToken(end),
         } = self;
 
         write!(f, "CASE")?;
@@ -2145,14 +2149,16 @@ impl fmt::Display for CaseStatement {
         }
 
         if let Some(else_block) = else_block {
-            write!(f, " ELSE ")?;
+            write!(f, " {else_block}")?;
-            format_statement_list(f, else_block)?;
         }
 
         write!(f, " END")?;
-        if *has_end_case {
+
+        if let Token::Word(w) = &end.token {
+            if w.keyword == Keyword::CASE {
                 write!(f, " CASE")?;
             }
+        }
 
         Ok(())
     }
@@ -2160,7 +2166,7 @@ impl fmt::Display for CaseStatement {
 
 /// An `IF` statement.
 ///
-/// Examples:
+/// Example (BigQuery or Snowflake):
 /// ```sql
 /// IF TRUE THEN
 ///     SELECT 1;
@@ -2171,16 +2177,22 @@ impl fmt::Display for CaseStatement {
 ///     SELECT 4;
 /// END IF
 /// ```
-///
 /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#if)
 /// [Snowflake](https://docs.snowflake.com/en/sql-reference/snowflake-scripting/if)
+///
+/// Example (MSSQL):
+/// ```sql
+/// IF 1=1 SELECT 1 ELSE SELECT 2
+/// ```
+/// [MSSQL](https://learn.microsoft.com/en-us/sql/t-sql/language-elements/if-else-transact-sql?view=sql-server-ver16)
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct IfStatement {
-    pub if_block: ConditionalStatements,
+    pub if_block: ConditionalStatementBlock,
-    pub elseif_blocks: Vec<ConditionalStatements>,
+    pub elseif_blocks: Vec<ConditionalStatementBlock>,
-    pub else_block: Option<Vec<Statement>>,
+    pub else_block: Option<ConditionalStatementBlock>,
+    pub end_token: Option<AttachedToken>,
 }
 
 impl fmt::Display for IfStatement {
@@ -2189,80 +2201,126 @@ impl fmt::Display for IfStatement {
             if_block,
             elseif_blocks,
             else_block,
+            end_token,
         } = self;
 
         write!(f, "{if_block}")?;
 
-        if !elseif_blocks.is_empty() {
+        for elseif_block in elseif_blocks {
-            write!(f, " {}", display_separated(elseif_blocks, " "))?;
+            write!(f, " {elseif_block}")?;
         }
 
         if let Some(else_block) = else_block {
-            write!(f, " ELSE ")?;
+            write!(f, " {else_block}")?;
-            format_statement_list(f, else_block)?;
         }
 
-        write!(f, " END IF")?;
+        if let Some(AttachedToken(end_token)) = end_token {
+            write!(f, " END {end_token}")?;
+        }
 
         Ok(())
     }
 }
 
-/// Represents a type of [ConditionalStatements]
-#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub enum ConditionalStatementKind {
-    /// `WHEN <condition> THEN <statements>`
-    When,
-    /// `IF <condition> THEN <statements>`
-    If,
-    /// `ELSEIF <condition> THEN <statements>`
-    ElseIf,
-}
-
 /// A block within a [Statement::Case] or [Statement::If]-like statement
 ///
-/// Examples:
+/// Example 1:
 /// ```sql
 /// WHEN EXISTS(SELECT 1) THEN SELECT 1;
+/// ```
 ///
+/// Example 2:
+/// ```sql
 /// IF TRUE THEN SELECT 1; SELECT 2;
 /// ```
+///
+/// Example 3:
+/// ```sql
+/// ELSE SELECT 1; SELECT 2;
+/// ```
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub struct ConditionalStatements {
+pub struct ConditionalStatementBlock {
-    /// The condition expression.
+    pub start_token: AttachedToken,
-    pub condition: Expr,
+    pub condition: Option<Expr>,
-    /// Statement list of the `THEN` clause.
+    pub then_token: Option<AttachedToken>,
-    pub statements: Vec<Statement>,
+    pub conditional_statements: ConditionalStatements,
-    pub kind: ConditionalStatementKind,
+}
+
+impl ConditionalStatementBlock {
+    pub fn statements(&self) -> &Vec<Statement> {
+        self.conditional_statements.statements()
+    }
+}
+
+impl fmt::Display for ConditionalStatementBlock {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let ConditionalStatementBlock {
+            start_token: AttachedToken(start_token),
+            condition,
+            then_token,
+            conditional_statements,
+        } = self;
+
+        write!(f, "{start_token}")?;
+
+        if let Some(condition) = condition {
+            write!(f, " {condition}")?;
+        }
+
+        if then_token.is_some() {
+            write!(f, " THEN")?;
+        }
+
+        if !conditional_statements.statements().is_empty() {
+            write!(f, " {conditional_statements}")?;
+        }
+
+        Ok(())
+    }
+}
+
+/// A list of statements in a [ConditionalStatementBlock].
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum ConditionalStatements {
+    /// SELECT 1; SELECT 2; SELECT 3; ...
+    Sequence { statements: Vec<Statement> },
+    /// BEGIN SELECT 1; SELECT 2; SELECT 3; ... END
+    BeginEnd {
+        begin_token: AttachedToken,
+        statements: Vec<Statement>,
+        end_token: AttachedToken,
+    },
+}
+
+impl ConditionalStatements {
+    pub fn statements(&self) -> &Vec<Statement> {
+        match self {
+            ConditionalStatements::Sequence { statements } => statements,
+            ConditionalStatements::BeginEnd { statements, .. } => statements,
+        }
+    }
 }
 
 impl fmt::Display for ConditionalStatements {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let ConditionalStatements {
+        match self {
-            condition: expr,
+            ConditionalStatements::Sequence { statements } => {
-            statements,
-            kind,
-        } = self;
-
-        let kind = match kind {
-            ConditionalStatementKind::When => "WHEN",
-            ConditionalStatementKind::If => "IF",
-            ConditionalStatementKind::ElseIf => "ELSEIF",
-        };
-
-        write!(f, "{kind} {expr} THEN")?;
-
                 if !statements.is_empty() {
-            write!(f, " ")?;
                     format_statement_list(f, statements)?;
                 }
-
                 Ok(())
             }
+            ConditionalStatements::BeginEnd { statements, .. } => {
+                write!(f, "BEGIN ")?;
+                format_statement_list(f, statements)?;
+                write!(f, " END")
+            }
+        }
+    }
 }
 
 /// A `RAISE` statement.

src/ast/spans.rs

@@ -22,21 +22,22 @@ use crate::tokenizer::Span;
 
 use super::{
     dcl::SecondaryRoles, value::ValueWithSpan, AccessExpr, AlterColumnOperation,
-    AlterIndexOperation, AlterTableOperation, Array, Assignment, AssignmentTarget, CaseStatement,
+    AlterIndexOperation, AlterTableOperation, Array, Assignment, AssignmentTarget, AttachedToken,
-    CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef, ConditionalStatements,
+    CaseStatement, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef,
-    ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
+    ConditionalStatementBlock, ConditionalStatements, ConflictTarget, ConnectBy,
-    CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
+    ConstraintCharacteristics, CopySource, CreateIndex, CreateTable, CreateTableOptions, Cte,
-    ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
+    Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr, ExprWithAlias, Fetch, FromTable,
-    FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
+    Function, FunctionArg, FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList,
-    IfStatement, IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join, JoinConstraint,
+    FunctionArguments, GroupByExpr, HavingBound, IfStatement, IlikeSelectItem, Insert, Interpolate,
-    JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause, MatchRecognizePattern, Measure,
+    InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView,
-    NamedWindowDefinition, ObjectName, ObjectNamePart, Offset, OnConflict, OnConflictAction,
+    LimitClause, MatchRecognizePattern, Measure, NamedWindowDefinition, ObjectName, ObjectNamePart,
-    OnInsert, OrderBy, OrderByExpr, OrderByKind, Partition, PivotValueSource, ProjectionSelect,
+    Offset, OnConflict, OnConflictAction, OnInsert, OrderBy, OrderByExpr, OrderByKind, Partition,
-    Query, RaiseStatement, RaiseStatementValue, ReferentialAction, RenameSelectItem,
+    PivotValueSource, ProjectionSelect, Query, RaiseStatement, RaiseStatementValue,
-    ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption,
+    ReferentialAction, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select,
-    Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint,
+    SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias,
-    TableFactor, TableObject, TableOptionsClustered, TableWithJoins, UpdateTableFromKind, Use,
+    TableAliasColumnDef, TableConstraint, TableFactor, TableObject, TableOptionsClustered,
-    Value, Values, ViewColumnDef, WildcardAdditionalOptions, With, WithFill,
+    TableWithJoins, UpdateTableFromKind, Use, Value, Values, ViewColumnDef,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -739,19 +740,14 @@ impl Spanned for CreateIndex {
 impl Spanned for CaseStatement {
     fn span(&self) -> Span {
         let CaseStatement {
-            match_expr,
+            case_token: AttachedToken(start),
-            when_blocks,
+            match_expr: _,
-            else_block,
+            when_blocks: _,
-            has_end_case: _,
+            else_block: _,
+            end_case_token: AttachedToken(end),
         } = self;
 
-        union_spans(
+        union_spans([start.span, end.span].into_iter())
-            match_expr
-                .iter()
-                .map(|e| e.span())
-                .chain(when_blocks.iter().map(|b| b.span()))
-                .chain(else_block.iter().flat_map(|e| e.iter().map(|s| s.span()))),
-        )
     }
 }
 
@@ -761,25 +757,48 @@ impl Spanned for IfStatement {
             if_block,
             elseif_blocks,
             else_block,
+            end_token,
         } = self;
 
         union_spans(
             iter::once(if_block.span())
                 .chain(elseif_blocks.iter().map(|b| b.span()))
-                .chain(else_block.iter().flat_map(|e| e.iter().map(|s| s.span()))),
+                .chain(else_block.as_ref().map(|b| b.span()))
+                .chain(end_token.as_ref().map(|AttachedToken(t)| t.span)),
         )
     }
 }
 
 impl Spanned for ConditionalStatements {
     fn span(&self) -> Span {
-        let ConditionalStatements {
+        match self {
+            ConditionalStatements::Sequence { statements } => {
+                union_spans(statements.iter().map(|s| s.span()))
+            }
+            ConditionalStatements::BeginEnd {
+                begin_token: AttachedToken(start),
+                statements: _,
+                end_token: AttachedToken(end),
+            } => union_spans([start.span, end.span].into_iter()),
+        }
+    }
+}
+
+impl Spanned for ConditionalStatementBlock {
+    fn span(&self) -> Span {
+        let ConditionalStatementBlock {
+            start_token: AttachedToken(start_token),
             condition,
-            statements,
+            then_token,
-            kind: _,
+            conditional_statements,
         } = self;
 
-        union_spans(iter::once(condition.span()).chain(statements.iter().map(|s| s.span())))
+        union_spans(
+            iter::once(start_token.span)
+                .chain(condition.as_ref().map(|c| c.span()))
+                .chain(then_token.as_ref().map(|AttachedToken(t)| t.span))
+                .chain(iter::once(conditional_statements.span())),
+        )
     }
 }
 

src/dialect/mssql.rs

@@ -15,7 +15,16 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::ast::helpers::attached_token::AttachedToken;
+use crate::ast::{ConditionalStatementBlock, ConditionalStatements, IfStatement, Statement};
 use crate::dialect::Dialect;
+use crate::keywords::{self, Keyword};
+use crate::parser::{Parser, ParserError};
+use crate::tokenizer::Token;
+#[cfg(not(feature = "std"))]
+use alloc::{vec, vec::Vec};
+
+const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[Keyword::IF, Keyword::ELSE];
 
 /// A [`Dialect`] for [Microsoft SQL Server](https://www.microsoft.com/en-us/sql-server/)
 #[derive(Debug)]
@@ -106,4 +115,123 @@ impl Dialect for MsSqlDialect {
     fn supports_object_name_double_dot_notation(&self) -> bool {
         true
     }
+
+    fn is_column_alias(&self, kw: &Keyword, _parser: &mut Parser) -> bool {
+        !keywords::RESERVED_FOR_COLUMN_ALIAS.contains(kw) && !RESERVED_FOR_COLUMN_ALIAS.contains(kw)
+    }
+
+    fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
+        if parser.peek_keyword(Keyword::IF) {
+            Some(self.parse_if_stmt(parser))
+        } else {
+            None
+        }
+    }
+}
+
+impl MsSqlDialect {
+    /// ```sql
+    /// IF boolean_expression
+    ///     { sql_statement | statement_block }
+    /// [ ELSE
+    ///     { sql_statement | statement_block } ]
+    /// ```
+    fn parse_if_stmt(&self, parser: &mut Parser) -> Result<Statement, ParserError> {
+        let if_token = parser.expect_keyword(Keyword::IF)?;
+
+        let condition = parser.parse_expr()?;
+
+        let if_block = if parser.peek_keyword(Keyword::BEGIN) {
+            let begin_token = parser.expect_keyword(Keyword::BEGIN)?;
+            let statements = self.parse_statement_list(parser, Some(Keyword::END))?;
+            let end_token = parser.expect_keyword(Keyword::END)?;
+            ConditionalStatementBlock {
+                start_token: AttachedToken(if_token),
+                condition: Some(condition),
+                then_token: None,
+                conditional_statements: ConditionalStatements::BeginEnd {
+                    begin_token: AttachedToken(begin_token),
+                    statements,
+                    end_token: AttachedToken(end_token),
+                },
+            }
+        } else {
+            let stmt = parser.parse_statement()?;
+            ConditionalStatementBlock {
+                start_token: AttachedToken(if_token),
+                condition: Some(condition),
+                then_token: None,
+                conditional_statements: ConditionalStatements::Sequence {
+                    statements: vec![stmt],
+                },
+            }
+        };
+
+        while let Token::SemiColon = parser.peek_token_ref().token {
+            parser.advance_token();
+        }
+
+        let mut else_block = None;
+        if parser.peek_keyword(Keyword::ELSE) {
+            let else_token = parser.expect_keyword(Keyword::ELSE)?;
+            if parser.peek_keyword(Keyword::BEGIN) {
+                let begin_token = parser.expect_keyword(Keyword::BEGIN)?;
+                let statements = self.parse_statement_list(parser, Some(Keyword::END))?;
+                let end_token = parser.expect_keyword(Keyword::END)?;
+                else_block = Some(ConditionalStatementBlock {
+                    start_token: AttachedToken(else_token),
+                    condition: None,
+                    then_token: None,
+                    conditional_statements: ConditionalStatements::BeginEnd {
+                        begin_token: AttachedToken(begin_token),
+                        statements,
+                        end_token: AttachedToken(end_token),
+                    },
+                });
+            } else {
+                let stmt = parser.parse_statement()?;
+                else_block = Some(ConditionalStatementBlock {
+                    start_token: AttachedToken(else_token),
+                    condition: None,
+                    then_token: None,
+                    conditional_statements: ConditionalStatements::Sequence {
+                        statements: vec![stmt],
+                    },
+                });
+            }
+        }
+
+        Ok(Statement::If(IfStatement {
+            if_block,
+            else_block,
+            elseif_blocks: Vec::new(),
+            end_token: None,
+        }))
+    }
+
+    /// Parse a sequence of statements, optionally separated by semicolon.
+    ///
+    /// Stops parsing when reaching EOF or the given keyword.
+    fn parse_statement_list(
+        &self,
+        parser: &mut Parser,
+        terminal_keyword: Option<Keyword>,
+    ) -> Result<Vec<Statement>, ParserError> {
+        let mut stmts = Vec::new();
+        loop {
+            if let Token::EOF = parser.peek_token_ref().token {
+                break;
+            }
+            if let Some(term) = terminal_keyword {
+                if parser.peek_keyword(term) {
+                    break;
+                }
+            }
+            stmts.push(parser.parse_statement()?);
+            while let Token::SemiColon = parser.peek_token_ref().token {
+                parser.advance_token();
+            }
+        }
+        Ok(stmts)
+    }
 }

src/parser/mod.rs

@@ -631,7 +631,7 @@ impl<'a> Parser<'a> {
     ///
     /// See [Statement::Case]
     pub fn parse_case_stmt(&mut self) -> Result<Statement, ParserError> {
-        self.expect_keyword_is(Keyword::CASE)?;
+        let case_token = self.expect_keyword(Keyword::CASE)?;
 
         let match_expr = if self.peek_keyword(Keyword::WHEN) {
             None
@@ -641,26 +641,26 @@ impl<'a> Parser<'a> {
 
         self.expect_keyword_is(Keyword::WHEN)?;
         let when_blocks = self.parse_keyword_separated(Keyword::WHEN, |parser| {
-            parser.parse_conditional_statements(
+            parser.parse_conditional_statement_block(&[Keyword::WHEN, Keyword::ELSE, Keyword::END])
-                ConditionalStatementKind::When,
-                &[Keyword::WHEN, Keyword::ELSE, Keyword::END],
-            )
         })?;
 
         let else_block = if self.parse_keyword(Keyword::ELSE) {
-            Some(self.parse_statement_list(&[Keyword::END])?)
+            Some(self.parse_conditional_statement_block(&[Keyword::END])?)
         } else {
             None
         };
 
-        self.expect_keyword_is(Keyword::END)?;
+        let mut end_case_token = self.expect_keyword(Keyword::END)?;
-        let has_end_case = self.parse_keyword(Keyword::CASE);
+        if self.peek_keyword(Keyword::CASE) {
+            end_case_token = self.expect_keyword(Keyword::CASE)?;
+        }
 
         Ok(Statement::Case(CaseStatement {
+            case_token: AttachedToken(case_token),
             match_expr,
             when_blocks,
             else_block,
-            has_end_case,
+            end_case_token: AttachedToken(end_case_token),
         }))
     }
 
@@ -669,34 +669,38 @@ impl<'a> Parser<'a> {
     /// See [Statement::If]
     pub fn parse_if_stmt(&mut self) -> Result<Statement, ParserError> {
         self.expect_keyword_is(Keyword::IF)?;
-        let if_block = self.parse_conditional_statements(
+        let if_block = self.parse_conditional_statement_block(&[
-            ConditionalStatementKind::If,
+            Keyword::ELSE,
-            &[Keyword::ELSE, Keyword::ELSEIF, Keyword::END],
+            Keyword::ELSEIF,
-        )?;
+            Keyword::END,
+        ])?;
 
         let elseif_blocks = if self.parse_keyword(Keyword::ELSEIF) {
             self.parse_keyword_separated(Keyword::ELSEIF, |parser| {
-                parser.parse_conditional_statements(
+                parser.parse_conditional_statement_block(&[
-                    ConditionalStatementKind::ElseIf,
+                    Keyword::ELSEIF,
-                    &[Keyword::ELSEIF, Keyword::ELSE, Keyword::END],
+                    Keyword::ELSE,
-                )
+                    Keyword::END,
+                ])
             })?
         } else {
             vec![]
         };
 
         let else_block = if self.parse_keyword(Keyword::ELSE) {
-            Some(self.parse_statement_list(&[Keyword::END])?)
+            Some(self.parse_conditional_statement_block(&[Keyword::END])?)
         } else {
             None
         };
 
-        self.expect_keywords(&[Keyword::END, Keyword::IF])?;
+        self.expect_keyword_is(Keyword::END)?;
+        let end_token = self.expect_keyword(Keyword::IF)?;
 
         Ok(Statement::If(IfStatement {
             if_block,
             elseif_blocks,
             else_block,
+            end_token: Some(AttachedToken(end_token)),
         }))
     }
 
@@ -707,19 +711,29 @@ impl<'a> Parser<'a> {
     /// ```sql
     /// IF condition THEN statement1; statement2;
     /// ```
-    fn parse_conditional_statements(
+    fn parse_conditional_statement_block(
         &mut self,
-        kind: ConditionalStatementKind,
         terminal_keywords: &[Keyword],
-    ) -> Result<ConditionalStatements, ParserError> {
+    ) -> Result<ConditionalStatementBlock, ParserError> {
-        let condition = self.parse_expr()?;
+        let start_token = self.get_current_token().clone(); // self.expect_keyword(keyword)?;
-        self.expect_keyword_is(Keyword::THEN)?;
+        let mut then_token = None;
+
+        let condition = match &start_token.token {
+            Token::Word(w) if w.keyword == Keyword::ELSE => None,
+            _ => {
+                let expr = self.parse_expr()?;
+                then_token = Some(AttachedToken(self.expect_keyword(Keyword::THEN)?));
+                Some(expr)
+            }
+        };
+
         let statements = self.parse_statement_list(terminal_keywords)?;
 
-        Ok(ConditionalStatements {
+        Ok(ConditionalStatementBlock {
+            start_token: AttachedToken(start_token),
             condition,
-            statements,
+            then_token,
-            kind,
+            conditional_statements: ConditionalStatements::Sequence { statements },
         })
     }
 

tests/sqlparser_common.rs

@@ -14229,9 +14229,12 @@ fn parse_case_statement() {
     };
 
     assert_eq!(Some(Expr::value(number("1"))), stmt.match_expr);
-    assert_eq!(Expr::value(number("2")), stmt.when_blocks[0].condition);
+    assert_eq!(
-    assert_eq!(2, stmt.when_blocks[0].statements.len());
+        Some(Expr::value(number("2"))),
-    assert_eq!(1, stmt.else_block.unwrap().len());
+        stmt.when_blocks[0].condition
+    );
+    assert_eq!(2, stmt.when_blocks[0].statements().len());
+    assert_eq!(1, stmt.else_block.unwrap().statements().len());
 
     verified_stmt(concat!(
         "CASE 1",
@@ -14274,17 +14277,35 @@ fn parse_case_statement() {
     );
 }
 
+#[test]
+fn test_case_statement_span() {
+    let sql = "CASE 1 WHEN 2 THEN SELECT 1; SELECT 2; ELSE SELECT 3; END CASE";
+    let mut parser = Parser::new(&GenericDialect {}).try_with_sql(sql).unwrap();
+    assert_eq!(
+        parser.parse_statement().unwrap().span(),
+        Span::new(Location::new(1, 1), Location::new(1, sql.len() as u64 + 1))
+    );
+}
+
 #[test]
 fn parse_if_statement() {
+    let dialects = all_dialects_except(|d| d.is::<MsSqlDialect>());
+
     let sql = "IF 1 THEN SELECT 1; ELSEIF 2 THEN SELECT 2; ELSE SELECT 3; END IF";
-    let Statement::If(stmt) = verified_stmt(sql) else {
+    let Statement::If(IfStatement {
+        if_block,
+        elseif_blocks,
+        else_block,
+        ..
+    }) = dialects.verified_stmt(sql)
+    else {
         unreachable!()
     };
-    assert_eq!(Expr::value(number("1")), stmt.if_block.condition);
+    assert_eq!(Some(Expr::value(number("1"))), if_block.condition);
-    assert_eq!(Expr::value(number("2")), stmt.elseif_blocks[0].condition);
+    assert_eq!(Some(Expr::value(number("2"))), elseif_blocks[0].condition);
-    assert_eq!(1, stmt.else_block.unwrap().len());
+    assert_eq!(1, else_block.unwrap().statements().len());
 
-    verified_stmt(concat!(
+    dialects.verified_stmt(concat!(
         "IF 1 THEN",
         " SELECT 1;",
         " SELECT 2;",
@@ -14300,7 +14321,7 @@ fn parse_if_statement() {
         " SELECT 9;",
         " END IF"
     ));
-    verified_stmt(concat!(
+    dialects.verified_stmt(concat!(
         "IF 1 THEN",
         " SELECT 1;",
         " SELECT 2;",
@@ -14309,7 +14330,7 @@ fn parse_if_statement() {
         " SELECT 4;",
         " END IF"
     ));
-    verified_stmt(concat!(
+    dialects.verified_stmt(concat!(
         "IF 1 THEN",
         " SELECT 1;",
         " SELECT 2;",
@@ -14319,22 +14340,79 @@ fn parse_if_statement() {
         " SELECT 4;",
         " END IF"
     ));
-    verified_stmt(concat!("IF 1 THEN", " SELECT 1;", " SELECT 2;", " END IF"));
+    dialects.verified_stmt(concat!("IF 1 THEN", " SELECT 1;", " SELECT 2;", " END IF"));
-    verified_stmt(concat!(
+    dialects.verified_stmt(concat!(
         "IF (1) THEN",
         " SELECT 1;",
         " SELECT 2;",
         " END IF"
     ));
-    verified_stmt("IF 1 THEN END IF");
+    dialects.verified_stmt("IF 1 THEN END IF");
-    verified_stmt("IF 1 THEN SELECT 1; ELSEIF 1 THEN END IF");
+    dialects.verified_stmt("IF 1 THEN SELECT 1; ELSEIF 1 THEN END IF");
 
     assert_eq!(
         ParserError::ParserError("Expected: IF, found: EOF".to_string()),
-        parse_sql_statements("IF 1 THEN SELECT 1; ELSEIF 1 THEN SELECT 2; END").unwrap_err()
+        dialects
+            .parse_sql_statements("IF 1 THEN SELECT 1; ELSEIF 1 THEN SELECT 2; END")
+            .unwrap_err()
     );
 }
 
+#[test]
+fn test_if_statement_span() {
+    let sql = "IF 1=1 THEN SELECT 1; ELSEIF 1=2 THEN SELECT 2; ELSE SELECT 3; END IF";
+    let mut parser = Parser::new(&GenericDialect {}).try_with_sql(sql).unwrap();
+    assert_eq!(
+        parser.parse_statement().unwrap().span(),
+        Span::new(Location::new(1, 1), Location::new(1, sql.len() as u64 + 1))
+    );
+}
+
+#[test]
+fn test_if_statement_multiline_span() {
+    let sql_line1 = "IF 1 = 1 THEN SELECT 1;";
+    let sql_line2 = "ELSEIF 1 = 2 THEN SELECT 2;";
+    let sql_line3 = "ELSE SELECT 3;";
+    let sql_line4 = "END IF";
+    let sql = [sql_line1, sql_line2, sql_line3, sql_line4].join("\n");
+    let mut parser = Parser::new(&GenericDialect {}).try_with_sql(&sql).unwrap();
+    assert_eq!(
+        parser.parse_statement().unwrap().span(),
+        Span::new(
+            Location::new(1, 1),
+            Location::new(4, sql_line4.len() as u64 + 1)
+        )
+    );
+}
+
+#[test]
+fn test_conditional_statement_span() {
+    let sql = "IF 1=1 THEN SELECT 1; ELSEIF 1=2 THEN SELECT 2; ELSE SELECT 3; END IF";
+    let mut parser = Parser::new(&GenericDialect {}).try_with_sql(sql).unwrap();
+    match parser.parse_statement().unwrap() {
+        Statement::If(IfStatement {
+            if_block,
+            elseif_blocks,
+            else_block,
+            ..
+        }) => {
+            assert_eq!(
+                Span::new(Location::new(1, 1), Location::new(1, 21)),
+                if_block.span()
+            );
+            assert_eq!(
+                Span::new(Location::new(1, 23), Location::new(1, 47)),
+                elseif_blocks[0].span()
+            );
+            assert_eq!(
+                Span::new(Location::new(1, 49), Location::new(1, 62)),
+                else_block.unwrap().span()
+            );
+        }
+        stmt => panic!("Unexpected statement: {:?}", stmt),
+    }
+}
+
 #[test]
 fn parse_raise_statement() {
     let sql = "RAISE USING MESSAGE = 42";

tests/sqlparser_mssql.rs

@@ -23,7 +23,7 @@
 mod test_utils;
 
 use helpers::attached_token::AttachedToken;
-use sqlparser::tokenizer::Span;
+use sqlparser::tokenizer::{Location, Span};
 use test_utils::*;
 
 use sqlparser::ast::DataType::{Int, Text, Varbinary};
@@ -31,7 +31,7 @@ use sqlparser::ast::DeclareAssignment::MsSqlAssignment;
 use sqlparser::ast::Value::SingleQuotedString;
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, MsSqlDialect};
-use sqlparser::parser::ParserError;
+use sqlparser::parser::{Parser, ParserError};
 
 #[test]
 fn parse_mssql_identifiers() {
@@ -1857,6 +1857,104 @@ fn parse_mssql_set_session_value() {
     ms().verified_stmt("SET ANSI_NULLS, ANSI_PADDING ON");
 }
 
+#[test]
+fn parse_mssql_if_else() {
+    // Simple statements and blocks
+    ms().verified_stmt("IF 1 = 1 SELECT '1'; ELSE SELECT '2';");
+    ms().verified_stmt("IF 1 = 1 BEGIN SET @A = 1; END ELSE SET @A = 2;");
+    ms().verified_stmt(
+        "IF DATENAME(weekday, GETDATE()) IN (N'Saturday', N'Sunday') SELECT 'Weekend'; ELSE SELECT 'Weekday';"
+    );
+    ms().verified_stmt(
+        "IF (SELECT COUNT(*) FROM a.b WHERE c LIKE 'x%') > 1 SELECT 'yes'; ELSE SELECT 'No';",
+    );
+
+    // Multiple statements
+    let stmts = ms()
+        .parse_sql_statements("DECLARE @A INT; IF 1=1 BEGIN SET @A = 1 END ELSE SET @A = 2")
+        .unwrap();
+    match &stmts[..] {
+        [Statement::Declare { .. }, Statement::If(stmt)] => {
+            assert_eq!(
+                stmt.to_string(),
+                "IF 1 = 1 BEGIN SET @A = 1; END ELSE SET @A = 2;"
+            );
+        }
+        _ => panic!("Unexpected statements: {:?}", stmts),
+    }
+}
+
+#[test]
+fn test_mssql_if_else_span() {
+    let sql = "IF 1 = 1 SELECT '1' ELSE SELECT '2'";
+    let mut parser = Parser::new(&MsSqlDialect {}).try_with_sql(sql).unwrap();
+    assert_eq!(
+        parser.parse_statement().unwrap().span(),
+        Span::new(Location::new(1, 1), Location::new(1, sql.len() as u64 + 1))
+    );
+}
+
+#[test]
+fn test_mssql_if_else_multiline_span() {
+    let sql_line1 = "IF 1 = 1";
+    let sql_line2 = "SELECT '1'";
+    let sql_line3 = "ELSE SELECT '2'";
+    let sql = [sql_line1, sql_line2, sql_line3].join("\n");
+    let mut parser = Parser::new(&MsSqlDialect {}).try_with_sql(&sql).unwrap();
+    assert_eq!(
+        parser.parse_statement().unwrap().span(),
+        Span::new(
+            Location::new(1, 1),
+            Location::new(3, sql_line3.len() as u64 + 1)
+        )
+    );
+}
+
+#[test]
+fn test_mssql_if_statements_span() {
+    // Simple statements
+    let mut sql = "IF 1 = 1 SELECT '1' ELSE SELECT '2'";
+    let mut parser = Parser::new(&MsSqlDialect {}).try_with_sql(sql).unwrap();
+    match parser.parse_statement().unwrap() {
+        Statement::If(IfStatement {
+            if_block,
+            else_block: Some(else_block),
+            ..
+        }) => {
+            assert_eq!(
+                if_block.span(),
+                Span::new(Location::new(1, 1), Location::new(1, 20))
+            );
+            assert_eq!(
+                else_block.span(),
+                Span::new(Location::new(1, 21), Location::new(1, 36))
+            );
+        }
+        stmt => panic!("Unexpected statement: {:?}", stmt),
+    }
+
+    // Blocks
+    sql = "IF 1 = 1 BEGIN SET @A = 1; END ELSE BEGIN SET @A = 2 END";
+    parser = Parser::new(&MsSqlDialect {}).try_with_sql(sql).unwrap();
+    match parser.parse_statement().unwrap() {
+        Statement::If(IfStatement {
+            if_block,
+            else_block: Some(else_block),
+            ..
+        }) => {
+            assert_eq!(
+                if_block.span(),
+                Span::new(Location::new(1, 1), Location::new(1, 31))
+            );
+            assert_eq!(
+                else_block.span(),
+                Span::new(Location::new(1, 32), Location::new(1, 57))
+            );
+        }
+        stmt => panic!("Unexpected statement: {:?}", stmt),
+    }
+}
+
 #[test]
 fn parse_mssql_varbinary_max_length() {
     let sql = "CREATE TABLE example (var_binary_col VARBINARY(MAX))";
@@ -1918,6 +2016,7 @@ fn parse_mssql_table_identifier_with_default_schema() {
 fn ms() -> TestedDialects {
     TestedDialects::new(vec![Box::new(MsSqlDialect {})])
 }
+
 fn ms_and_generic() -> TestedDialects {
     TestedDialects::new(vec![Box::new(MsSqlDialect {}), Box::new(GenericDialect {})])
 }