Skip to content

Commit 4f79997

Browse files
xitepxitep
authored
Include DML keyword in statement span (#2090)
1 parent 78be8b1 commit 4f79997

File tree

8 files changed

+150
-37
lines changed

8 files changed

+150
-37
lines changed

‎src/ast/dml.rs‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,19 @@ use sqlparser_derive::{Visit, VisitMut};
2727
use crate::display_utils::{indented_list, Indent, SpaceOrNewline};
2828

2929
use super::{
30-
display_comma_separated, query::InputFormatClause, Assignment, Expr, FromTable, Ident,
31-
InsertAliases, MysqlInsertPriority, ObjectName, OnInsert, OrderByExpr, Query, SelectItem,
32-
Setting, SqliteOnConflict, TableObject, TableWithJoins, UpdateTableFromKind,
30+
display_comma_separated, helpers::attached_token::AttachedToken, query::InputFormatClause,
31+
Assignment, Expr, FromTable, Ident, InsertAliases, MysqlInsertPriority, ObjectName, OnInsert,
32+
OrderByExpr, Query, SelectItem, Setting, SqliteOnConflict, TableObject, TableWithJoins,
33+
UpdateTableFromKind,
3334
};
3435

3536
/// INSERT statement.
3637
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
3738
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
3839
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
3940
pub struct Insert {
41+
/// Token for the `INSERT` keyword (or its substitutes)
42+
pub insert_token: AttachedToken,
4043
/// Only for Sqlite
4144
pub or: Option<SqliteOnConflict>,
4245
/// Only for mysql
@@ -179,6 +182,8 @@ impl Display for Insert {
179182
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
180183
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
181184
pub struct Delete {
185+
/// Token for the `DELETE` keyword
186+
pub delete_token: AttachedToken,
182187
/// Multi tables delete are supported in mysql
183188
pub tables: Vec<ObjectName>,
184189
/// FROM
@@ -246,6 +251,8 @@ impl Display for Delete {
246251
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
247252
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
248253
pub struct Update {
254+
/// Token for the `UPDATE` keyword
255+
pub update_token: AttachedToken,
249256
/// TABLE
250257
pub table: TableWithJoins,
251258
/// Column assignments

‎src/ast/spans.rs‎

Lines changed: 97 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -839,6 +839,7 @@ impl Spanned for CopySource {
839839
impl Spanned for Delete {
840840
fn span(&self) -> Span {
841841
let Delete {
842+
delete_token,
842843
tables,
843844
from,
844845
using,
@@ -849,26 +850,29 @@ impl Spanned for Delete {
849850
} = self;
850851

851852
union_spans(
852-
tables
853-
.iter()
854-
.map(|i| i.span())
855-
.chain(core::iter::once(from.span()))
856-
.chain(
857-
using
858-
.iter()
859-
.map(|u| union_spans(u.iter().map(|i| i.span()))),
860-
)
861-
.chain(selection.iter().map(|i| i.span()))
862-
.chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
863-
.chain(order_by.iter().map(|i| i.span()))
864-
.chain(limit.iter().map(|i| i.span())),
853+
core::iter::once(delete_token.0.span).chain(
854+
tables
855+
.iter()
856+
.map(|i| i.span())
857+
.chain(core::iter::once(from.span()))
858+
.chain(
859+
using
860+
.iter()
861+
.map(|u| union_spans(u.iter().map(|i| i.span()))),
862+
)
863+
.chain(selection.iter().map(|i| i.span()))
864+
.chain(returning.iter().flat_map(|i| i.iter().map(|k| k.span())))
865+
.chain(order_by.iter().map(|i| i.span()))
866+
.chain(limit.iter().map(|i| i.span())),
867+
),
865868
)
866869
}
867870
}
868871

869872
impl Spanned for Update {
870873
fn span(&self) -> Span {
871874
let Update {
875+
update_token,
872876
table,
873877
assignments,
874878
from,
@@ -880,6 +884,7 @@ impl Spanned for Update {
880884

881885
union_spans(
882886
core::iter::once(table.span())
887+
.chain(core::iter::once(update_token.0.span))
883888
.chain(assignments.iter().map(|i| i.span()))
884889
.chain(from.iter().map(|i| i.span()))
885890
.chain(selection.iter().map(|i| i.span()))
@@ -1217,6 +1222,7 @@ impl Spanned for AlterIndexOperation {
12171222
impl Spanned for Insert {
12181223
fn span(&self) -> Span {
12191224
let Insert {
1225+
insert_token,
12201226
or: _, // enum, sqlite specific
12211227
ignore: _, // bool
12221228
into: _, // bool
@@ -1239,7 +1245,8 @@ impl Spanned for Insert {
12391245
} = self;
12401246

12411247
union_spans(
1242-
core::iter::once(table.span())
1248+
core::iter::once(insert_token.0.span)
1249+
.chain(core::iter::once(table.span()))
12431250
.chain(table_alias.as_ref().map(|i| i.span))
12441251
.chain(columns.iter().map(|i| i.span))
12451252
.chain(source.as_ref().map(|q| q.span()))
@@ -2540,4 +2547,80 @@ ALTER TABLE users
25402547
assert_eq!(stmt_span.start, (2, 13).into());
25412548
assert_eq!(stmt_span.end, (4, 11).into());
25422549
}
2550+
2551+
#[test]
2552+
fn test_update_statement_span() {
2553+
let sql = r#"-- foo
2554+
UPDATE foo
2555+
/* bar */
2556+
SET bar = 3
2557+
WHERE quux > 42 ;
2558+
"#;
2559+
2560+
let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
2561+
assert_eq!(1, r.len());
2562+
2563+
let stmt_span = r[0].span();
2564+
2565+
assert_eq!(stmt_span.start, (2, 7).into());
2566+
assert_eq!(stmt_span.end, (5, 17).into());
2567+
}
2568+
2569+
#[test]
2570+
fn test_insert_statement_span() {
2571+
let sql = r#"
2572+
/* foo */ INSERT INTO FOO (X, Y, Z)
2573+
SELECT 1, 2, 3
2574+
FROM DUAL
2575+
;"#;
2576+
2577+
let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
2578+
assert_eq!(1, r.len());
2579+
2580+
let stmt_span = r[0].span();
2581+
2582+
assert_eq!(stmt_span.start, (2, 11).into());
2583+
assert_eq!(stmt_span.end, (4, 12).into());
2584+
}
2585+
2586+
#[test]
2587+
fn test_replace_statement_span() {
2588+
let sql = r#"
2589+
/* foo */ REPLACE INTO
2590+
cities(name,population)
2591+
SELECT
2592+
name,
2593+
population
2594+
FROM
2595+
cities
2596+
WHERE id = 1
2597+
;"#;
2598+
2599+
let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
2600+
assert_eq!(1, r.len());
2601+
2602+
dbg!(&r[0]);
2603+
2604+
let stmt_span = r[0].span();
2605+
2606+
assert_eq!(stmt_span.start, (2, 11).into());
2607+
assert_eq!(stmt_span.end, (9, 13).into());
2608+
}
2609+
2610+
#[test]
2611+
fn test_delete_statement_span() {
2612+
let sql = r#"-- foo
2613+
DELETE /* quux */
2614+
FROM foo
2615+
WHERE foo.x = 42
2616+
;"#;
2617+
2618+
let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
2619+
assert_eq!(1, r.len());
2620+
2621+
let stmt_span = r[0].span();
2622+
2623+
assert_eq!(stmt_span.start, (2, 7).into());
2624+
assert_eq!(stmt_span.end, (4, 24).into());
2625+
}
25432626
}

‎src/dialect/sqlite.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ impl Dialect for SQLiteDialect {
6868
fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
6969
if parser.parse_keyword(Keyword::REPLACE) {
7070
parser.prev_token();
71-
Some(parser.parse_insert())
71+
Some(parser.parse_insert(parser.get_current_token().clone()))
7272
} else {
7373
None
7474
}

‎src/parser/mod.rs‎

Lines changed: 33 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -586,11 +586,11 @@ impl<'a> Parser<'a> {
586586
Keyword::DISCARD => self.parse_discard(),
587587
Keyword::DECLARE => self.parse_declare(),
588588
Keyword::FETCH => self.parse_fetch_statement(),
589-
Keyword::DELETE => self.parse_delete(),
590-
Keyword::INSERT => self.parse_insert(),
591-
Keyword::REPLACE => self.parse_replace(),
589+
Keyword::DELETE => self.parse_delete(next_token),
590+
Keyword::INSERT => self.parse_insert(next_token),
591+
Keyword::REPLACE => self.parse_replace(next_token),
592592
Keyword::UNCACHE => self.parse_uncache_table(),
593-
Keyword::UPDATE => self.parse_update(),
593+
Keyword::UPDATE => self.parse_update(next_token),
594594
Keyword::ALTER => self.parse_alter(),
595595
Keyword::CALL => self.parse_call(),
596596
Keyword::COPY => self.parse_copy(),
@@ -11817,8 +11817,11 @@ impl<'a> Parser<'a> {
1181711817
/// Parse a DELETE statement, returning a `Box`ed SetExpr
1181811818
///
1181911819
/// This is used to reduce the size of the stack frames in debug builds
11820-
fn parse_delete_setexpr_boxed(&mut self) -> Result<Box<SetExpr>, ParserError> {
11821-
Ok(Box::new(SetExpr::Delete(self.parse_delete()?)))
11820+
fn parse_delete_setexpr_boxed(
11821+
&mut self,
11822+
delete_token: TokenWithSpan,
11823+
) -> Result<Box<SetExpr>, ParserError> {
11824+
Ok(Box::new(SetExpr::Delete(self.parse_delete(delete_token)?)))
1182211825
}
1182311826

1182411827
/// Parse a MERGE statement, returning a `Box`ed SetExpr
@@ -11828,7 +11831,7 @@ impl<'a> Parser<'a> {
1182811831
Ok(Box::new(SetExpr::Merge(self.parse_merge()?)))
1182911832
}
1183011833

11831-
pub fn parse_delete(&mut self) -> Result<Statement, ParserError> {
11834+
pub fn parse_delete(&mut self, delete_token: TokenWithSpan) -> Result<Statement, ParserError> {
1183211835
let (tables, with_from_keyword) = if !self.parse_keyword(Keyword::FROM) {
1183311836
// `FROM` keyword is optional in BigQuery SQL.
1183411837
// https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#delete_statement
@@ -11871,6 +11874,7 @@ impl<'a> Parser<'a> {
1187111874
};
1187211875

1187311876
Ok(Statement::Delete(Delete {
11877+
delete_token: delete_token.into(),
1187411878
tables,
1187511879
from: if with_from_keyword {
1187611880
FromTable::WithFromKeyword(from)
@@ -12000,7 +12004,7 @@ impl<'a> Parser<'a> {
1200012004
if self.parse_keyword(Keyword::INSERT) {
1200112005
Ok(Query {
1200212006
with,
12003-
body: self.parse_insert_setexpr_boxed()?,
12007+
body: self.parse_insert_setexpr_boxed(self.get_current_token().clone())?,
1200412008
order_by: None,
1200512009
limit_clause: None,
1200612010
fetch: None,
@@ -12014,7 +12018,7 @@ impl<'a> Parser<'a> {
1201412018
} else if self.parse_keyword(Keyword::UPDATE) {
1201512019
Ok(Query {
1201612020
with,
12017-
body: self.parse_update_setexpr_boxed()?,
12021+
body: self.parse_update_setexpr_boxed(self.get_current_token().clone())?,
1201812022
order_by: None,
1201912023
limit_clause: None,
1202012024
fetch: None,
@@ -12028,7 +12032,7 @@ impl<'a> Parser<'a> {
1202812032
} else if self.parse_keyword(Keyword::DELETE) {
1202912033
Ok(Query {
1203012034
with,
12031-
body: self.parse_delete_setexpr_boxed()?,
12035+
body: self.parse_delete_setexpr_boxed(self.get_current_token().clone())?,
1203212036
limit_clause: None,
1203312037
order_by: None,
1203412038
fetch: None,
@@ -15470,15 +15474,18 @@ impl<'a> Parser<'a> {
1547015474
}
1547115475

1547215476
/// Parse an REPLACE statement
15473-
pub fn parse_replace(&mut self) -> Result<Statement, ParserError> {
15477+
pub fn parse_replace(
15478+
&mut self,
15479+
replace_token: TokenWithSpan,
15480+
) -> Result<Statement, ParserError> {
1547415481
if !dialect_of!(self is MySqlDialect | GenericDialect) {
1547515482
return parser_err!(
1547615483
"Unsupported statement REPLACE",
1547715484
self.peek_token().span.start
1547815485
);
1547915486
}
1548015487

15481-
let mut insert = self.parse_insert()?;
15488+
let mut insert = self.parse_insert(replace_token)?;
1548215489
if let Statement::Insert(Insert { replace_into, .. }) = &mut insert {
1548315490
*replace_into = true;
1548415491
}
@@ -15489,12 +15496,15 @@ impl<'a> Parser<'a> {
1548915496
/// Parse an INSERT statement, returning a `Box`ed SetExpr
1549015497
///
1549115498
/// This is used to reduce the size of the stack frames in debug builds
15492-
fn parse_insert_setexpr_boxed(&mut self) -> Result<Box<SetExpr>, ParserError> {
15493-
Ok(Box::new(SetExpr::Insert(self.parse_insert()?)))
15499+
fn parse_insert_setexpr_boxed(
15500+
&mut self,
15501+
insert_token: TokenWithSpan,
15502+
) -> Result<Box<SetExpr>, ParserError> {
15503+
Ok(Box::new(SetExpr::Insert(self.parse_insert(insert_token)?)))
1549415504
}
1549515505

1549615506
/// Parse an INSERT statement
15497-
pub fn parse_insert(&mut self) -> Result<Statement, ParserError> {
15507+
pub fn parse_insert(&mut self, insert_token: TokenWithSpan) -> Result<Statement, ParserError> {
1549815508
let or = self.parse_conflict_clause();
1549915509
let priority = if !dialect_of!(self is MySqlDialect | GenericDialect) {
1550015510
None
@@ -15663,6 +15673,7 @@ impl<'a> Parser<'a> {
1566315673
};
1566415674

1566515675
Ok(Statement::Insert(Insert {
15676+
insert_token: insert_token.into(),
1566615677
or,
1566715678
table: table_object,
1566815679
table_alias,
@@ -15754,11 +15765,14 @@ impl<'a> Parser<'a> {
1575415765
/// Parse an UPDATE statement, returning a `Box`ed SetExpr
1575515766
///
1575615767
/// This is used to reduce the size of the stack frames in debug builds
15757-
fn parse_update_setexpr_boxed(&mut self) -> Result<Box<SetExpr>, ParserError> {
15758-
Ok(Box::new(SetExpr::Update(self.parse_update()?)))
15768+
fn parse_update_setexpr_boxed(
15769+
&mut self,
15770+
update_token: TokenWithSpan,
15771+
) -> Result<Box<SetExpr>, ParserError> {
15772+
Ok(Box::new(SetExpr::Update(self.parse_update(update_token)?)))
1575915773
}
1576015774

15761-
pub fn parse_update(&mut self) -> Result<Statement, ParserError> {
15775+
pub fn parse_update(&mut self, update_token: TokenWithSpan) -> Result<Statement, ParserError> {
1576215776
let or = self.parse_conflict_clause();
1576315777
let table = self.parse_table_and_joins()?;
1576415778
let from_before_set = if self.parse_keyword(Keyword::FROM) {
@@ -15793,6 +15807,7 @@ impl<'a> Parser<'a> {
1579315807
None
1579415808
};
1579515809
Ok(Update {
15810+
update_token: update_token.into(),
1579615811
table,
1579715812
assignments,
1579815813
from,

‎tests/sqlparser_common.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,7 @@ fn parse_update_set_from() {
456456
assert_eq!(
457457
stmt,
458458
Statement::Update(Update {
459+
update_token: AttachedToken::empty(),
459460
table: TableWithJoins {
460461
relation: table_from_name(ObjectName::from(vec![Ident::new("t1")])),
461462
joins: vec![],
@@ -551,6 +552,7 @@ fn parse_update_with_table_alias() {
551552
returning,
552553
or: None,
553554
limit: None,
555+
update_token: _,
554556
}) => {
555557
assert_eq!(
556558
TableWithJoins {

‎tests/sqlparser_mysql.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2632,6 +2632,7 @@ fn parse_update_with_joins() {
26322632
returning,
26332633
or: None,
26342634
limit: None,
2635+
update_token: _,
26352636
}) => {
26362637
assert_eq!(
26372638
TableWithJoins {

0 commit comments

Comments
 (0)