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

@@ -16,32 +16,32 @@
 // under the License.
 
 use crate::ast::{
-    ddl::AlterSchema, query::SelectItemQualifiedWildcardKind, AlterSchemaOperation, AlterTable,
-    ColumnOptions, CreateView, ExportData, Owner, TypedString,
+    AlterSchemaOperation, AlterTable, ColumnOptions, CreateView, ExportData, Owner, TypedString,
+    ddl::AlterSchema, query::SelectItemQualifiedWildcardKind,
 };
 use core::iter;
 
 use crate::tokenizer::Span;
 
 use super::{
-    dcl::SecondaryRoles, value::ValueWithSpan, AccessExpr, AlterColumnOperation,
-    AlterIndexOperation, AlterTableOperation, Analyze, Array, Assignment, AssignmentTarget,
-    AttachedToken, BeginEndStatements, CaseStatement, CloseCursor, ClusteredIndex, ColumnDef,
-    ColumnOption, ColumnOptionDef, ConditionalStatementBlock, ConditionalStatements,
-    ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
-    CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
-    ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
-    FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
-    IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr, Join,
-    JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause,
+    AccessExpr, AlterColumnOperation, AlterIndexOperation, AlterTableOperation, Analyze, Array,
+    Assignment, AssignmentTarget, AttachedToken, BeginEndStatements, CaseStatement, CloseCursor,
+    ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef, ConditionalStatementBlock,
+    ConditionalStatements, ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource,
+    CreateIndex, CreateTable, CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem,
+    ExcludeSelectItem, Expr, ExprWithAlias, Fetch, FromTable, Function, FunctionArg,
+    FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr,
+    HavingBound, IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr,
+    Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause,
     MatchRecognizePattern, Measure, NamedParenthesizedList, NamedWindowDefinition, ObjectName,
     ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement, OrderBy,
     OrderByExpr, OrderByKind, Partition, PivotValueSource, ProjectionSelect, Query, RaiseStatement,
     RaiseStatementValue, ReferentialAction, RenameSelectItem, ReplaceSelectElement,
     ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript,
     SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint, TableFactor, TableObject,
     TableOptionsClustered, TableWithJoins, Update, UpdateTableFromKind, Use, Value, Values,
-    ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With, WithFill,
+    ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With, WithFill, dcl::SecondaryRoles,
+    value::ValueWithSpan,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -867,6 +867,7 @@ impl Spanned for Delete {
 impl Spanned for Update {
     fn span(&self) -> Span {
         let Update {
+            update_token,
             table,
             assignments,
             from,
@@ -878,6 +879,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()))
@@ -2438,25 +2440,37 @@ pub mod tests {
     #[test]
     pub fn test_cte() {
         let dialect = &GenericDialect;
-        let mut test = SpanTest::new(dialect, "WITH cte_outer AS (SELECT a FROM postgres.public.source), cte_ignored AS (SELECT a FROM cte_outer), cte_inner AS (SELECT a FROM cte_outer) SELECT a FROM cte_inner");
+        let mut test = SpanTest::new(
+            dialect,
+            "WITH cte_outer AS (SELECT a FROM postgres.public.source), cte_ignored AS (SELECT a FROM cte_outer), cte_inner AS (SELECT a FROM cte_outer) SELECT a FROM cte_inner",
+        );
 
         let query = test.0.parse_query().unwrap();
 
         let select_span = query.span();
 
-        assert_eq!(test.get_source(select_span), "WITH cte_outer AS (SELECT a FROM postgres.public.source), cte_ignored AS (SELECT a FROM cte_outer), cte_inner AS (SELECT a FROM cte_outer) SELECT a FROM cte_inner");
+        assert_eq!(
+            test.get_source(select_span),
+            "WITH cte_outer AS (SELECT a FROM postgres.public.source), cte_ignored AS (SELECT a FROM cte_outer), cte_inner AS (SELECT a FROM cte_outer) SELECT a FROM cte_inner"
+        );
     }
 
     #[test]
     pub fn test_snowflake_lateral_flatten() {
         let dialect = &SnowflakeDialect;
-        let mut test = SpanTest::new(dialect, "SELECT FLATTENED.VALUE:field::TEXT AS FIELD FROM SNOWFLAKE.SCHEMA.SOURCE AS S, LATERAL FLATTEN(INPUT => S.JSON_ARRAY) AS FLATTENED");
+        let mut test = SpanTest::new(
+            dialect,
+            "SELECT FLATTENED.VALUE:field::TEXT AS FIELD FROM SNOWFLAKE.SCHEMA.SOURCE AS S, LATERAL FLATTEN(INPUT => S.JSON_ARRAY) AS FLATTENED",
+        );
 
         let query = test.0.parse_select().unwrap();
 
         let select_span = query.span();
 
-        assert_eq!(test.get_source(select_span), "SELECT FLATTENED.VALUE:field::TEXT AS FIELD FROM SNOWFLAKE.SCHEMA.SOURCE AS S, LATERAL FLATTEN(INPUT => S.JSON_ARRAY) AS FLATTENED");
+        assert_eq!(
+            test.get_source(select_span),
+            "SELECT FLATTENED.VALUE:field::TEXT AS FIELD FROM SNOWFLAKE.SCHEMA.SOURCE AS S, LATERAL FLATTEN(INPUT => S.JSON_ARRAY) AS FLATTENED"
+        );
     }
 
     #[test]
@@ -2535,4 +2549,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(),
@@ -12009,7 +12009,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,
@@ -15734,11 +15734,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) {
@@ -15773,6 +15776,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,
@@ -440,6 +440,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![],
@@ -535,6 +539,7 @@ fn parse_update_with_table_alias() {
             returning,
             or: None,
             limit: None,
+            update_token,
         }) => {
             assert_eq!(
                 TableWithJoins {
@@ -583,6 +588,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]
@@ -2623,6 +2623,7 @@ fn parse_update_with_joins() {
             returning,
             or: None,
             limit: None,
+            update_token,
         }) => {
             assert_eq!(
                 TableWithJoins {
@@ -2697,6 +2698,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())
+            })
         })
     );
 }