Include UPDATE keyword in statement span

Author: xitep

src/ast/dml.rs

@@ -27,9 +27,10 @@ use sqlparser_derive::{Visit, VisitMut};
 use crate::display_utils::{indented_list, Indent, SpaceOrNewline};
 
 use super::{
-    display_comma_separated, query::InputFormatClause, Assignment, Expr, FromTable, Ident,
-    InsertAliases, MysqlInsertPriority, ObjectName, OnInsert, OrderByExpr, Query, SelectItem,
-    Setting, SqliteOnConflict, TableObject, TableWithJoins, UpdateTableFromKind,
+    display_comma_separated, helpers::attached_token::AttachedToken, query::InputFormatClause,
+    Assignment, Expr, FromTable, Ident, InsertAliases, MysqlInsertPriority, ObjectName, OnInsert,
+    OrderByExpr, Query, SelectItem, Setting, SqliteOnConflict, TableObject, TableWithJoins,
+    UpdateTableFromKind,
 };
 
 /// INSERT statement.
@@ -246,6 +247,8 @@ impl Display for Delete {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct Update {
+    /// Token for the `UPDATE` keyword
+    pub update_token: AttachedToken,
     /// TABLE
     pub table: TableWithJoins,
     /// Column assignments

src/ast/spans.rs

@@ -869,6 +869,7 @@ impl Spanned for Delete {
 impl Spanned for Update {
     fn span(&self) -> Span {
         let Update {
+            update_token,
             table,
             assignments,
             from,
@@ -880,6 +881,7 @@ impl Spanned for Update {
 
         union_spans(
             core::iter::once(table.span())
+                .chain(core::iter::once(update_token.0.span))
                 .chain(assignments.iter().map(|i| i.span()))
                 .chain(from.iter().map(|i| i.span()))
                 .chain(selection.iter().map(|i| i.span()))
@@ -2540,4 +2542,22 @@ ALTER TABLE users
         assert_eq!(stmt_span.start, (2, 13).into());
         assert_eq!(stmt_span.end, (4, 11).into());
     }
+
+    #[test]
+    fn test_update_statement_span() {
+        let sql = r#"-- foo
+      UPDATE foo
+   /* bar */
+   SET bar = 3
+ WHERE quux > 42 ;
+"#;
+
+        let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        let stmt_span = r[0].span();
+
+        assert_eq!(stmt_span.start, (2, 7).into());
+        assert_eq!(stmt_span.end, (5, 17).into());
+    }
 }

src/parser/mod.rs

@@ -590,7 +590,7 @@ impl<'a> Parser<'a> {
                 Keyword::INSERT => self.parse_insert(),
                 Keyword::REPLACE => self.parse_replace(),
                 Keyword::UNCACHE => self.parse_uncache_table(),
-                Keyword::UPDATE => self.parse_update(),
+                Keyword::UPDATE => self.parse_update(next_token),
                 Keyword::ALTER => self.parse_alter(),
                 Keyword::CALL => self.parse_call(),
                 Keyword::COPY => self.parse_copy(),
@@ -12014,7 +12014,7 @@ impl<'a> Parser<'a> {
         } else if self.parse_keyword(Keyword::UPDATE) {
             Ok(Query {
                 with,
-                body: self.parse_update_setexpr_boxed()?,
+                body: self.parse_update_setexpr_boxed(self.get_current_token().clone())?,
                 order_by: None,
                 limit_clause: None,
                 fetch: None,
@@ -15754,11 +15754,14 @@ impl<'a> Parser<'a> {
     /// Parse an UPDATE statement, returning a `Box`ed SetExpr
     ///
     /// This is used to reduce the size of the stack frames in debug builds
-    fn parse_update_setexpr_boxed(&mut self) -> Result<Box<SetExpr>, ParserError> {
-        Ok(Box::new(SetExpr::Update(self.parse_update()?)))
+    fn parse_update_setexpr_boxed(
+        &mut self,
+        update_token: TokenWithSpan,
+    ) -> Result<Box<SetExpr>, ParserError> {
+        Ok(Box::new(SetExpr::Update(self.parse_update(update_token)?)))
     }
 
-    pub fn parse_update(&mut self) -> Result<Statement, ParserError> {
+    pub fn parse_update(&mut self, update_token: TokenWithSpan) -> Result<Statement, ParserError> {
         let or = self.parse_conflict_clause();
         let table = self.parse_table_and_joins()?;
         let from_before_set = if self.parse_keyword(Keyword::FROM) {
@@ -15793,6 +15796,7 @@ impl<'a> Parser<'a> {
             None
         };
         Ok(Update {
+            update_token: update_token.into(),
             table,
             assignments,
             from,

tests/sqlparser_common.rs

@@ -39,8 +39,8 @@ use sqlparser::dialect::{
 };
 use sqlparser::keywords::{Keyword, ALL_KEYWORDS};
 use sqlparser::parser::{Parser, ParserError, ParserOptions};
-use sqlparser::tokenizer::Tokenizer;
-use sqlparser::tokenizer::{Location, Span};
+use sqlparser::tokenizer::{Location, Span, TokenWithSpan};
+use sqlparser::tokenizer::{Token, Tokenizer};
 use test_utils::{
     all_dialects, all_dialects_where, all_dialects_with_options, alter_table_op, assert_eq_vec,
     call, expr_from_projection, join, number, only, table, table_alias, table_from_name,
@@ -456,6 +456,10 @@ fn parse_update_set_from() {
     assert_eq!(
         stmt,
         Statement::Update(Update {
+            update_token: AttachedToken(TokenWithSpan {
+                token: Token::make_keyword("UPDATE"),
+                span: Span::new((1, 1).into(), (1, 7).into()),
+            }),
             table: TableWithJoins {
                 relation: table_from_name(ObjectName::from(vec![Ident::new("t1")])),
                 joins: vec![],
@@ -551,6 +555,7 @@ fn parse_update_with_table_alias() {
             returning,
             or: None,
             limit: None,
+            update_token,
         }) => {
             assert_eq!(
                 TableWithJoins {
@@ -599,6 +604,13 @@ fn parse_update_with_table_alias() {
                 selection
             );
             assert_eq!(None, returning);
+            assert_eq!(
+                AttachedToken(TokenWithSpan {
+                    token: Token::make_keyword("UPDATE"),
+                    span: Span::new((1, 1).into(), (1, 7).into()),
+                }),
+                update_token
+            );
         }
         _ => unreachable!(),
     }

tests/sqlparser_mysql.rs

@@ -26,8 +26,8 @@ use sqlparser::ast::MysqlInsertPriority::{Delayed, HighPriority, LowPriority};
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, MySqlDialect};
 use sqlparser::parser::{ParserError, ParserOptions};
-use sqlparser::tokenizer::Span;
 use sqlparser::tokenizer::Token;
+use sqlparser::tokenizer::{Span, TokenWithSpan};
 use test_utils::*;
 
 #[macro_use]
@@ -2632,6 +2632,7 @@ fn parse_update_with_joins() {
             returning,
             or: None,
             limit: None,
+            update_token,
         }) => {
             assert_eq!(
                 TableWithJoins {
@@ -2706,6 +2707,13 @@ fn parse_update_with_joins() {
                 selection
             );
             assert_eq!(None, returning);
+            assert_eq!(
+                AttachedToken(TokenWithSpan {
+                    token: Token::make_keyword("UPDATE"),
+                    span: Span::new((1, 1).into(), (1, 7).into()),
+                }),
+                update_token
+            );
         }
         _ => unreachable!(),
     }

tests/sqlparser_sqlite.rs

@@ -22,6 +22,7 @@
 #[macro_use]
 mod test_utils;
 
+use sqlparser::ast::helpers::attached_token::AttachedToken;
 use sqlparser::keywords::Keyword;
 use test_utils::*;
 
@@ -30,7 +31,7 @@ use sqlparser::ast::Value::Placeholder;
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, SQLiteDialect};
 use sqlparser::parser::{ParserError, ParserOptions};
-use sqlparser::tokenizer::Token;
+use sqlparser::tokenizer::{Span, Token, TokenWithSpan};
 
 #[test]
 fn pragma_no_value() {
@@ -494,7 +495,11 @@ fn parse_update_tuple_row_values() {
             },
             from: None,
             returning: None,
-            limit: None
+            limit: None,
+            update_token: AttachedToken(TokenWithSpan {
+                token: Token::make_keyword("UPDATE"),
+                span: Span::new((1, 1).into(), (1, 7).into())
+            })
         })
     );
 }