fix: case expression spans to include leading and trailing keywords

Author: eliaperantoni

src/ast/mod.rs

@@ -967,6 +967,8 @@ pub enum Expr {
     /// not `< 0` nor `1, 2, 3` as allowed in a `<simple when clause>` per
     /// <https://jakewheat.github.io/sql-overview/sql-2011-foundation-grammar.html#simple-when-clause>
     Case {
+        case_token: AttachedToken,
+        end_token: AttachedToken,
         operand: Option<Box<Expr>>,
         conditions: Vec<CaseWhen>,
         else_result: Option<Box<Expr>>,
@@ -1675,6 +1677,8 @@ impl fmt::Display for Expr {
             }
             Expr::Function(fun) => fun.fmt(f),
             Expr::Case {
+                case_token: _,
+                end_token: _,
                 operand,
                 conditions,
                 else_result,

src/ast/spans.rs

@@ -1566,18 +1566,24 @@ impl Spanned for Expr {
             ),
             Expr::Prefixed { value, .. } => value.span(),
             Expr::Case {
+                case_token,
+                end_token,
                 operand,
                 conditions,
                 else_result,
             } => union_spans(
-                operand
-                    .as_ref()
-                    .map(|i| i.span())
-                    .into_iter()
-                    .chain(conditions.iter().flat_map(|case_when| {
-                        [case_when.condition.span(), case_when.result.span()]
-                    }))
-                    .chain(else_result.as_ref().map(|i| i.span())),
+                iter::once(case_token.0.span)
+                    .chain(
+                        operand
+                            .as_ref()
+                            .map(|i| i.span())
+                            .into_iter()
+                            .chain(conditions.iter().flat_map(|case_when| {
+                                [case_when.condition.span(), case_when.result.span()]
+                            }))
+                            .chain(else_result.as_ref().map(|i| i.span())),
+                    )
+                    .chain(iter::once(end_token.0.span)),
             ),
             Expr::Exists { subquery, .. } => subquery.span(),
             Expr::Subquery(query) => query.span(),

src/parser/mod.rs

@@ -2274,6 +2274,7 @@ impl<'a> Parser<'a> {
     }
 
     pub fn parse_case_expr(&mut self) -> Result<Expr, ParserError> {
+        let case_token = AttachedToken(self.get_previous_token().clone());
         let mut operand = None;
         if !self.parse_keyword(Keyword::WHEN) {
             operand = Some(Box::new(self.parse_expr()?));
@@ -2294,8 +2295,10 @@ impl<'a> Parser<'a> {
         } else {
             None
         };
-        self.expect_keyword_is(Keyword::END)?;
+        let end_token = AttachedToken(self.expect_keyword(Keyword::END)?);
         Ok(Expr::Case {
+            case_token,
+            end_token,
             operand,
             conditions,
             else_result,

tests/sqlparser_common.rs

@@ -6861,6 +6861,8 @@ fn parse_searched_case_expr() {
     let select = verified_only_select(sql);
     assert_eq!(
         &Case {
+            case_token: AttachedToken::empty(),
+            end_token: AttachedToken::empty(),
             operand: None,
             conditions: vec![
                 CaseWhen {
@@ -6900,6 +6902,8 @@ fn parse_simple_case_expr() {
     use self::Expr::{Case, Identifier};
     assert_eq!(
         &Case {
+            case_token: AttachedToken::empty(),
+            end_token: AttachedToken::empty(),
             operand: Some(Box::new(Identifier(Ident::new("foo")))),
             conditions: vec![CaseWhen {
                 condition: Expr::value(number("1")),
@@ -14464,6 +14468,16 @@ fn test_case_statement_span() {
     );
 }
 
+#[test]
+fn test_case_expr_span() {
+    let sql = "CASE 1 WHEN 2 THEN 3 ELSE 4 END";
+    let mut parser = Parser::new(&GenericDialect {}).try_with_sql(sql).unwrap();
+    assert_eq!(
+        parser.parse_expr().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>());
@@ -14642,6 +14656,8 @@ fn test_lambdas() {
                 Expr::Lambda(LambdaFunction {
                     params: OneOrManyWithParens::Many(vec![Ident::new("p1"), Ident::new("p2")]),
                     body: Box::new(Expr::Case {
+                        case_token: AttachedToken::empty(),
+                        end_token: AttachedToken::empty(),
                         operand: None,
                         conditions: vec![
                             CaseWhen {

tests/sqlparser_databricks.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use sqlparser::ast::helpers::attached_token::AttachedToken;
 use sqlparser::ast::*;
 use sqlparser::dialect::{DatabricksDialect, GenericDialect};
 use sqlparser::parser::ParserError;
@@ -108,6 +109,8 @@ fn test_databricks_lambdas() {
                 Expr::Lambda(LambdaFunction {
                     params: OneOrManyWithParens::Many(vec![Ident::new("p1"), Ident::new("p2")]),
                     body: Box::new(Expr::Case {
+                        case_token: AttachedToken::empty(),
+                        end_token: AttachedToken::empty(),
                         operand: None,
                         conditions: vec![
                             CaseWhen {