Include INSERT keyword in statement span

Author: xitep

src/ast/dml.rs

@@ -38,6 +38,8 @@ use super::{
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct Insert {
+    /// Token for the `INSERT` keyword (or its substitutes)
+    pub insert_token: AttachedToken,
     /// Only for Sqlite
     pub or: Option<SqliteOnConflict>,
     /// Only for mysql

src/ast/spans.rs

@@ -1221,6 +1221,7 @@ impl Spanned for AlterIndexOperation {
 impl Spanned for Insert {
     fn span(&self) -> Span {
         let Insert {
+            insert_token,
             or: _,     // enum, sqlite specific
             ignore: _, // bool
             into: _,   // bool
@@ -1244,6 +1245,7 @@ impl Spanned for Insert {
 
         union_spans(
             core::iter::once(table.span())
+                .chain(core::iter::once(insert_token.0.span))
                 .chain(table_alias.as_ref().map(|i| i.span))
                 .chain(columns.iter().map(|i| i.span))
                 .chain(source.as_ref().map(|q| q.span()))

src/dialect/sqlite.rs

@@ -68,7 +68,7 @@ impl Dialect for SQLiteDialect {
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
         if parser.parse_keyword(Keyword::REPLACE) {
             parser.prev_token();
-            Some(parser.parse_insert())
+            Some(parser.parse_insert(parser.get_current_token().clone()))
         } else {
             None
         }

src/parser/mod.rs

@@ -587,8 +587,8 @@ impl<'a> Parser<'a> {
                 Keyword::DECLARE => self.parse_declare(),
                 Keyword::FETCH => self.parse_fetch_statement(),
                 Keyword::DELETE => self.parse_delete(next_token),
-                Keyword::INSERT => self.parse_insert(),
-                Keyword::REPLACE => self.parse_replace(),
+                Keyword::INSERT => self.parse_insert(next_token),
+                Keyword::REPLACE => self.parse_replace(next_token),
                 Keyword::UNCACHE => self.parse_uncache_table(),
                 Keyword::UPDATE => self.parse_update(next_token),
                 Keyword::ALTER => self.parse_alter(),
@@ -12001,7 +12001,7 @@ impl<'a> Parser<'a> {
         if self.parse_keyword(Keyword::INSERT) {
             Ok(Query {
                 with,
-                body: self.parse_insert_setexpr_boxed()?,
+                body: self.parse_insert_setexpr_boxed(self.get_current_token().clone())?,
                 order_by: None,
                 limit_clause: None,
                 fetch: None,
@@ -15471,15 +15471,15 @@ impl<'a> Parser<'a> {
     }
 
     /// Parse an REPLACE statement
-    pub fn parse_replace(&mut self) -> Result<Statement, ParserError> {
+    pub fn parse_replace(&mut self, replace_token: TokenWithSpan) -> Result<Statement, ParserError> {
         if !dialect_of!(self is MySqlDialect | GenericDialect) {
             return parser_err!(
                 "Unsupported statement REPLACE",
                 self.peek_token().span.start
             );
         }
 
-        let mut insert = self.parse_insert()?;
+        let mut insert = self.parse_insert(replace_token)?;
         if let Statement::Insert(Insert { replace_into, .. }) = &mut insert {
             *replace_into = true;
         }
@@ -15490,12 +15490,12 @@ impl<'a> Parser<'a> {
     /// Parse an INSERT statement, returning a `Box`ed SetExpr
     ///
     /// This is used to reduce the size of the stack frames in debug builds
-    fn parse_insert_setexpr_boxed(&mut self) -> Result<Box<SetExpr>, ParserError> {
-        Ok(Box::new(SetExpr::Insert(self.parse_insert()?)))
+    fn parse_insert_setexpr_boxed(&mut self, insert_token: TokenWithSpan) -> Result<Box<SetExpr>, ParserError> {
+        Ok(Box::new(SetExpr::Insert(self.parse_insert(insert_token)?)))
     }
 
     /// Parse an INSERT statement
-    pub fn parse_insert(&mut self) -> Result<Statement, ParserError> {
+    pub fn parse_insert(&mut self, insert_token: TokenWithSpan) -> Result<Statement, ParserError> {
         let or = self.parse_conflict_clause();
         let priority = if !dialect_of!(self is MySqlDialect | GenericDialect) {
             None
@@ -15664,6 +15664,7 @@ impl<'a> Parser<'a> {
             };
 
             Ok(Statement::Insert(Insert {
+                insert_token: insert_token.into(),
                 or,
                 table: table_object,
                 table_alias,

tests/sqlparser_postgres.rs

@@ -23,7 +23,7 @@
 mod test_utils;
 
 use helpers::attached_token::AttachedToken;
-use sqlparser::tokenizer::Span;
+use sqlparser::tokenizer::{Span, Token, TokenWithSpan};
 use test_utils::*;
 
 use sqlparser::ast::*;
@@ -5140,6 +5140,10 @@ fn test_simple_postgres_insert_with_alias() {
     assert_eq!(
         statement,
         Statement::Insert(Insert {
+            insert_token: AttachedToken(TokenWithSpan {
+                token: Token::make_keyword("INSERT"),
+                span: Span::new((1, 1).into(), (1, 7).into()),
+            }),
             or: None,
             ignore: false,
             into: true,
@@ -5210,6 +5214,10 @@ fn test_simple_postgres_insert_with_alias() {
     assert_eq!(
         statement,
         Statement::Insert(Insert {
+            insert_token: AttachedToken(TokenWithSpan {
+                token: Token::make_keyword("INSERT"),
+                span: Span::new((1, 1).into(), (1, 7).into()),
+            }),
             or: None,
             ignore: false,
             into: true,
@@ -5282,6 +5290,10 @@ fn test_simple_insert_with_quoted_alias() {
     assert_eq!(
         statement,
         Statement::Insert(Insert {
+            insert_token: AttachedToken(TokenWithSpan {
+                token: Token::make_keyword("INSERT"),
+                span: Span::new((1, 1).into(), (1, 7).into()),
+            }),
             or: None,
             ignore: false,
             into: true,