Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement ON CONFLICT and RETURNING #666

Merged
merged 5 commits into from Nov 11, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
58 changes: 55 additions & 3 deletions src/ast/mod.rs
Expand Up @@ -1037,6 +1037,8 @@ pub enum Statement {
/// whether the insert has the table keyword (Hive)
table: bool,
on: Option<OnInsert>,
/// RETURNING
returning: Option<Vec<SelectItem>>,
},
// TODO: Support ROW FORMAT
Directory {
Expand Down Expand Up @@ -1077,6 +1079,8 @@ pub enum Statement {
from: Option<TableWithJoins>,
/// WHERE
selection: Option<Expr>,
/// RETURNING
returning: Option<Vec<SelectItem>>,
},
/// DELETE
Delete {
Expand All @@ -1086,6 +1090,8 @@ pub enum Statement {
using: Option<TableFactor>,
/// WHERE
selection: Option<Expr>,
/// RETURNING
returning: Option<Vec<SelectItem>>,
},
/// CREATE VIEW
CreateView {
Expand Down Expand Up @@ -1633,6 +1639,7 @@ impl fmt::Display for Statement {
source,
table,
on,
returning,
} => {
if let Some(action) = or {
write!(f, "INSERT OR {} INTO {} ", action, table_name)?;
Expand Down Expand Up @@ -1660,10 +1667,14 @@ impl fmt::Display for Statement {
write!(f, "{}", source)?;

if let Some(on) = on {
write!(f, "{}", on)
} else {
Ok(())
write!(f, "{}", on)?;
}

if let Some(returning) = returning {
write!(f, " RETURNING {}", display_comma_separated(returning))?;
}

Ok(())
}

Statement::Copy {
Expand Down Expand Up @@ -1707,6 +1718,7 @@ impl fmt::Display for Statement {
assignments,
from,
selection,
returning,
} => {
write!(f, "UPDATE {}", table)?;
if !assignments.is_empty() {
Expand All @@ -1718,12 +1730,16 @@ impl fmt::Display for Statement {
if let Some(selection) = selection {
write!(f, " WHERE {}", selection)?;
}
if let Some(returning) = returning {
write!(f, " RETURNING {}", display_comma_separated(returning))?;
}
Ok(())
}
Statement::Delete {
table_name,
using,
selection,
returning,
} => {
write!(f, "DELETE FROM {}", table_name)?;
if let Some(using) = using {
Expand All @@ -1732,6 +1748,9 @@ impl fmt::Display for Statement {
if let Some(selection) = selection {
write!(f, " WHERE {}", selection)?;
}
if let Some(returning) = returning {
write!(f, " RETURNING {}", display_comma_separated(returning))?;
}
Ok(())
}
Statement::Close { cursor } => {
Expand Down Expand Up @@ -2416,6 +2435,21 @@ impl fmt::Display for Statement {
pub enum OnInsert {
/// ON DUPLICATE KEY UPDATE (MySQL when the key already exists, then execute an update instead)
DuplicateKeyUpdate(Vec<Assignment>),
/// ON CONFLICT is a PostgreSQL and Sqlite extension
OnConflict(OnConflict),
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct OnConflict {
pub conflict_target: Vec<Ident>,
pub action: OnConflictAction,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum OnConflictAction {
DoNothing,
DoUpdate(Vec<Assignment>),
}

impl fmt::Display for OnInsert {
Expand All @@ -2426,6 +2460,24 @@ impl fmt::Display for OnInsert {
" ON DUPLICATE KEY UPDATE {}",
display_comma_separated(expr)
),
Self::OnConflict(o) => write!(f, " {o}"),
}
}
}
impl fmt::Display for OnConflict {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, " ON CONFLICT")?;
if !self.conflict_target.is_empty() {
write!(f, "({})", display_comma_separated(&self.conflict_target))?;
}
write!(f, " {}", self.action)
}
}
impl fmt::Display for OnConflictAction {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Self::DoNothing => write!(f, "DO NOTHING"),
Self::DoUpdate(a) => write!(f, "DO UPDATE SET {}", display_comma_separated(a)),
}
}
}
Expand Down
4 changes: 4 additions & 0 deletions src/keywords.rs
Expand Up @@ -142,6 +142,7 @@ define_keywords!(
COMMITTED,
COMPUTE,
CONDITION,
CONFLICT,
CONNECT,
CONNECTION,
CONSTRAINT,
Expand Down Expand Up @@ -198,6 +199,7 @@ define_keywords!(
DISCONNECT,
DISTINCT,
DISTRIBUTE,
DO,
DOUBLE,
DOW,
DOY,
Expand Down Expand Up @@ -363,6 +365,7 @@ define_keywords!(
NOSCAN,
NOSUPERUSER,
NOT,
NOTHING,
NTH_VALUE,
NTILE,
NULL,
Expand Down Expand Up @@ -454,6 +457,7 @@ define_keywords!(
RESTRICT,
RESULT,
RETURN,
RETURNING,
RETURNS,
REVOKE,
RIGHT,
Expand Down
50 changes: 45 additions & 5 deletions src/parser.rs
Expand Up @@ -3733,10 +3733,17 @@ impl<'a> Parser<'a> {
None
};

let returning = if self.parse_keyword(Keyword::RETURNING) {
Some(self.parse_comma_separated(Parser::parse_select_item)?)
} else {
None
};

Ok(Statement::Delete {
table_name,
using,
selection,
returning,
})
}

Expand Down Expand Up @@ -4824,12 +4831,38 @@ impl<'a> Parser<'a> {

let source = Box::new(self.parse_query()?);
let on = if self.parse_keyword(Keyword::ON) {
self.expect_keyword(Keyword::DUPLICATE)?;
self.expect_keyword(Keyword::KEY)?;
self.expect_keyword(Keyword::UPDATE)?;
let l = self.parse_comma_separated(Parser::parse_assignment)?;
if self.parse_keyword(Keyword::CONFLICT) {
let conflict_target =
self.parse_parenthesized_column_list(IsOptional::Optional)?;

Some(OnInsert::DuplicateKeyUpdate(l))
self.expect_keyword(Keyword::DO)?;
let action = if self.parse_keyword(Keyword::NOTHING) {
OnConflictAction::DoNothing
} else {
self.expect_keyword(Keyword::UPDATE)?;
self.expect_keyword(Keyword::SET)?;
let l = self.parse_comma_separated(Parser::parse_assignment)?;
OnConflictAction::DoUpdate(l)
};

Some(OnInsert::OnConflict(OnConflict {
conflict_target,
action,
}))
} else {
self.expect_keyword(Keyword::DUPLICATE)?;
self.expect_keyword(Keyword::KEY)?;
self.expect_keyword(Keyword::UPDATE)?;
let l = self.parse_comma_separated(Parser::parse_assignment)?;

Some(OnInsert::DuplicateKeyUpdate(l))
}
} else {
None
};

let returning = if self.parse_keyword(Keyword::RETURNING) {
Some(self.parse_comma_separated(Parser::parse_select_item)?)
} else {
None
};
Expand All @@ -4845,6 +4878,7 @@ impl<'a> Parser<'a> {
source,
table,
on,
returning,
})
}
}
Expand All @@ -4863,11 +4897,17 @@ impl<'a> Parser<'a> {
} else {
None
};
let returning = if self.parse_keyword(Keyword::RETURNING) {
Some(self.parse_comma_separated(Parser::parse_select_item)?)
} else {
None
};
Ok(Statement::Update {
table,
assignments,
from,
selection,
returning,
})
}

Expand Down
6 changes: 6 additions & 0 deletions tests/sqlparser_common.rs
Expand Up @@ -195,6 +195,7 @@ fn parse_update_with_table_alias() {
assignments,
from: _from,
selection,
returning,
} => {
assert_eq!(
TableWithJoins {
Expand Down Expand Up @@ -231,6 +232,7 @@ fn parse_update_with_table_alias() {
}),
selection
);
assert_eq!(None, returning);
}
_ => unreachable!(),
}
Expand Down Expand Up @@ -278,6 +280,7 @@ fn parse_where_delete_statement() {
table_name,
using,
selection,
returning,
} => {
assert_eq!(
TableFactor::Table {
Expand All @@ -298,6 +301,7 @@ fn parse_where_delete_statement() {
},
selection.unwrap(),
);
assert_eq!(None, returning);
}
_ => unreachable!(),
}
Expand All @@ -313,6 +317,7 @@ fn parse_where_delete_with_alias_statement() {
table_name,
using,
selection,
returning,
} => {
assert_eq!(
TableFactor::Table {
Expand Down Expand Up @@ -353,6 +358,7 @@ fn parse_where_delete_with_alias_statement() {
},
selection.unwrap(),
);
assert_eq!(None, returning);
}
_ => unreachable!(),
}
Expand Down
2 changes: 2 additions & 0 deletions tests/sqlparser_mysql.rs
Expand Up @@ -815,6 +815,7 @@ fn parse_update_with_joins() {
assignments,
from: _from,
selection,
returning,
} => {
assert_eq!(
TableWithJoins {
Expand Down Expand Up @@ -870,6 +871,7 @@ fn parse_update_with_joins() {
}),
selection
);
assert_eq!(None, returning);
}
_ => unreachable!(),
}
Expand Down