Fix CASE expression spans (#1874)

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

@@ -1567,18 +1567,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(),
@@ -2464,4 +2470,16 @@ pub mod tests {
 
         assert_eq!(test.get_source(body_span), "SELECT cte.* FROM cte");
     }
+
+    #[test]
+    fn test_case_expr_span() {
+        let dialect = &GenericDialect;
+        let mut test = SpanTest::new(dialect, "CASE 1 WHEN 2 THEN 3 ELSE 4 END");
+        let expr = test.0.parse_expr().unwrap();
+        let expr_span = expr.span();
+        assert_eq!(
+            test.get_source(expr_span),
+            "CASE 1 WHEN 2 THEN 3 ELSE 4 END"
+        );
+    }
 }

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_current_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

@@ -6869,6 +6869,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 {
@@ -6908,6 +6910,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")),
@@ -14650,6 +14654,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 {