mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-04 06:18:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

support sqlite's OR clauses in update statements (#1530)

Browse source
This commit is contained in:
Ophir LOJKINE 2024-11-18 13:30:53 +01:00 • committed by GitHub
parent a67a4f3cbe
commit 4c629e8520
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 59 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

4
src/ast/dml.rs
View file

@ -505,8 +505,8 @@ impl Display for Insert {
self.table_name.to_string()
};
if let Some(action) = self.or {
write!(f, "INSERT OR {action} INTO {table_name} ")?;
if let Some(on_conflict) = self.or {
write!(f, "INSERT {on_conflict} INTO {table_name} ")?;
} else {
write!(
f,

19
src/ast/mod.rs
View file

@ -2396,6 +2396,8 @@ pub enum Statement {
selection: Option<Expr>,
/// RETURNING
returning: Option<Vec<SelectItem>>,
/// SQLite-specific conflict resolution clause
or: Option<SqliteOnConflict>,
},
/// ```sql
/// DELETE
@ -3691,8 +3693,13 @@ impl fmt::Display for Statement {
from,
selection,
returning,
or,
} => {
write!(f, "UPDATE {table}")?;
write!(f, "UPDATE ")?;
if let Some(or) = or {
write!(f, "{or} ")?;
}
write!(f, "{table}")?;
if !assignments.is_empty() {
write!(f, " SET {}", display_comma_separated(assignments))?;
}
@ -6304,11 +6311,11 @@ impl fmt::Display for SqliteOnConflict {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
use SqliteOnConflict::*;
match self {
Rollback => write!(f, "ROLLBACK"),
Abort => write!(f, "ABORT"),
Fail => write!(f, "FAIL"),
Ignore => write!(f, "IGNORE"),
Replace => write!(f, "REPLACE"),
Rollback => write!(f, "OR ROLLBACK"),
Abort => write!(f, "OR ABORT"),
Fail => write!(f, "OR FAIL"),
Ignore => write!(f, "OR IGNORE"),
Replace => write!(f, "OR REPLACE"),
}
}
}

39
src/parser/mod.rs
View file

@ -11042,24 +11042,7 @@ impl<'a> Parser<'a> {
/// Parse an INSERT statement
pub fn parse_insert(&mut self) -> Result<Statement, ParserError> {
let or = if !dialect_of!(self is SQLiteDialect) {
None
} else if self.parse_keywords(&[Keyword::OR, Keyword::REPLACE]) {
Some(SqliteOnConflict::Replace)
} else if self.parse_keywords(&[Keyword::OR, Keyword::ROLLBACK]) {
Some(SqliteOnConflict::Rollback)
} else if self.parse_keywords(&[Keyword::OR, Keyword::ABORT]) {
Some(SqliteOnConflict::Abort)
} else if self.parse_keywords(&[Keyword::OR, Keyword::FAIL]) {
Some(SqliteOnConflict::Fail)
} else if self.parse_keywords(&[Keyword::OR, Keyword::IGNORE]) {
Some(SqliteOnConflict::Ignore)
} else if self.parse_keyword(Keyword::REPLACE) {
Some(SqliteOnConflict::Replace)
} else {
None
};
let or = self.parse_conflict_clause();
let priority = if !dialect_of!(self is MySqlDialect | GenericDialect) {
None
} else if self.parse_keyword(Keyword::LOW_PRIORITY) {
@ -11218,6 +11201,24 @@ impl<'a> Parser<'a> {
}
}
fn parse_conflict_clause(&mut self) -> Option<SqliteOnConflict> {
if self.parse_keywords(&[Keyword::OR, Keyword::REPLACE]) {
Some(SqliteOnConflict::Replace)
} else if self.parse_keywords(&[Keyword::OR, Keyword::ROLLBACK]) {
Some(SqliteOnConflict::Rollback)
} else if self.parse_keywords(&[Keyword::OR, Keyword::ABORT]) {
Some(SqliteOnConflict::Abort)
} else if self.parse_keywords(&[Keyword::OR, Keyword::FAIL]) {
Some(SqliteOnConflict::Fail)
} else if self.parse_keywords(&[Keyword::OR, Keyword::IGNORE]) {
Some(SqliteOnConflict::Ignore)
} else if self.parse_keyword(Keyword::REPLACE) {
Some(SqliteOnConflict::Replace)
} else {
None
}
}
pub fn parse_insert_partition(&mut self) -> Result<Option<Vec<Expr>>, ParserError> {
if self.parse_keyword(Keyword::PARTITION) {
self.expect_token(&Token::LParen)?;
@ -11253,6 +11254,7 @@ impl<'a> Parser<'a> {
}
pub fn parse_update(&mut self) -> Result<Statement, ParserError> {
let or = self.parse_conflict_clause();
let table = self.parse_table_and_joins()?;
self.expect_keyword(Keyword::SET)?;
let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
@ -11279,6 +11281,7 @@ impl<'a> Parser<'a> {
from,
selection,
returning,
or,
})
}

21
tests/sqlparser_common.rs
View file

@ -443,6 +443,7 @@ fn parse_update_set_from() {
])),
}),
returning: None,
or: None,
}
);
}
@ -457,6 +458,7 @@ fn parse_update_with_table_alias() {
from: _from,
selection,
returning,
or: None,
} => {
assert_eq!(
TableWithJoins {
@ -505,6 +507,25 @@ fn parse_update_with_table_alias() {
}
}
#[test]
fn parse_update_or() {
let expect_or_clause = |sql: &str, expected_action: SqliteOnConflict| match verified_stmt(sql) {
Statement::Update { or, .. } => assert_eq!(or, Some(expected_action)),
other => unreachable!("Expected update with or, got {:?}", other),
};
expect_or_clause(
"UPDATE OR REPLACE t SET n = n + 1",
SqliteOnConflict::Replace,
);
expect_or_clause(
"UPDATE OR ROLLBACK t SET n = n + 1",
SqliteOnConflict::Rollback,
);
expect_or_clause("UPDATE OR ABORT t SET n = n + 1", SqliteOnConflict::Abort);
expect_or_clause("UPDATE OR FAIL t SET n = n + 1", SqliteOnConflict::Fail);
expect_or_clause("UPDATE OR IGNORE t SET n = n + 1", SqliteOnConflict::Ignore);
}
#[test]
fn parse_select_with_table_alias_as() {
// AS is optional

1
tests/sqlparser_mysql.rs
View file

@ -1970,6 +1970,7 @@ fn parse_update_with_joins() {
from: _from,
selection,
returning,
or: None,
} => {
assert_eq!(
TableWithJoins {

1
tests/sqlparser_sqlite.rs
View file

@ -465,6 +465,7 @@ fn parse_update_tuple_row_values() {
assert_eq!(
sqlite().verified_stmt("UPDATE x SET (a, b) = (1, 2)"),
Statement::Update {
or: None,
assignments: vec![Assignment {
target: AssignmentTarget::Tuple(vec![
ObjectName(vec![Ident::new("a"),]),