Skip to content

Commit

Permalink
feat: support sqlite insert or statement (andialbrecht#281)
Browse files Browse the repository at this point in the history
  • Loading branch information
zhangli-pear authored Feb 9, 2021
1 parent 07342d5 commit add8991
Show file tree
Hide file tree
Showing 4 changed files with 105 additions and 9 deletions.
48 changes: 41 additions & 7 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -508,6 +508,8 @@ pub enum Statement {
Query(Box<Query>),
/// INSERT
Insert {
/// Only for Sqlite
or: Option<SqliteOnConflict>,
/// TABLE
table_name: ObjectName,
/// COLUMNS
Expand Down Expand Up @@ -804,6 +806,7 @@ impl fmt::Display for Statement {
Ok(())
}
Statement::Insert {
or,
table_name,
overwrite,
partitioned,
Expand All @@ -812,13 +815,17 @@ impl fmt::Display for Statement {
source,
table,
} => {
write!(
f,
"INSERT {act}{tbl} {table_name} ",
table_name = table_name,
act = if *overwrite { "OVERWRITE" } else { "INTO" },
tbl = if *table { " TABLE" } else { "" }
)?;
if let Some(action) = or {
write!(f, "INSERT OR {} INTO {} ", action, table_name)?;
} else {
write!(
f,
"INSERT {act}{tbl} {table_name} ",
table_name = table_name,
act = if *overwrite { "OVERWRITE" } else { "INTO" },
tbl = if *table { " TABLE" } else { "" }
)?;
}
if !columns.is_empty() {
write!(f, "({}) ", display_comma_separated(columns))?;
}
Expand All @@ -832,6 +839,7 @@ impl fmt::Display for Statement {
}
write!(f, "{}", source)
}

Statement::Copy {
table_name,
columns,
Expand Down Expand Up @@ -1560,3 +1568,29 @@ impl fmt::Display for SetVariableValue {
}
}
}

/// Sqlite specific syntax
///
/// https://sqlite.org/lang_conflict.html
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum SqliteOnConflict {
Rollback,
Abort,
Fail,
Ignore,
Replace,
}

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"),
}
}
}
3 changes: 3 additions & 0 deletions src/dialect/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ macro_rules! define_keywords {

// The following keywords should be sorted to be able to match using binary search
define_keywords!(
ABORT,
ABS,
ACTION,
ADD,
Expand Down Expand Up @@ -202,6 +203,7 @@ define_keywords!(
EXTENDED,
EXTERNAL,
EXTRACT,
FAIL,
FALSE,
FETCH,
FIELDS,
Expand Down Expand Up @@ -233,6 +235,7 @@ define_keywords!(
HOUR,
IDENTITY,
IF,
IGNORE,
IN,
INDEX,
INDICATOR,
Expand Down
22 changes: 22 additions & 0 deletions src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,10 @@ impl<'a> Parser<'a> {
Keyword::DEALLOCATE => Ok(self.parse_deallocate()?),
Keyword::EXECUTE => Ok(self.parse_execute()?),
Keyword::PREPARE => Ok(self.parse_prepare()?),
Keyword::REPLACE if dialect_of!(self is SQLiteDialect ) => {
self.prev_token();
Ok(self.parse_insert()?)
}
_ => self.expected("an SQL statement", Token::Word(w)),
},
Token::LParen => {
Expand Down Expand Up @@ -2719,6 +2723,23 @@ 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 action = self.expect_one_of_keywords(&[Keyword::INTO, Keyword::OVERWRITE])?;
let overwrite = action == Keyword::OVERWRITE;
let local = self.parse_keyword(Keyword::LOCAL);
Expand Down Expand Up @@ -2758,6 +2779,7 @@ impl<'a> Parser<'a> {

let source = Box::new(self.parse_query()?);
Ok(Statement::Insert {
or,
table_name,
overwrite,
partitioned,
Expand Down
41 changes: 39 additions & 2 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,8 @@ use test_utils::{all_dialects, expr_from_projection, join, number, only, table,

use matches::assert_matches;
use sqlparser::ast::*;
use sqlparser::dialect::keywords::ALL_KEYWORDS;
use sqlparser::parser::ParserError;
use sqlparser::dialect::{keywords::ALL_KEYWORDS, SQLiteDialect};
use sqlparser::parser::{Parser, ParserError};

#[test]
fn parse_insert_values() {
Expand Down Expand Up @@ -97,6 +97,43 @@ fn parse_insert_invalid() {
);
}

#[test]
fn parse_insert_sqlite() {
let dialect = SQLiteDialect {};

let check = |sql: &str, expected_action: Option<SqliteOnConflict>| match Parser::parse_sql(
&dialect, &sql,
)
.unwrap()
.pop()
.unwrap()
{
Statement::Insert { or, .. } => assert_eq!(or, expected_action),
_ => panic!(sql.to_string()),
};

let sql = "INSERT INTO test_table(id) VALUES(1)";
check(sql, None);

let sql = "REPLACE INTO test_table(id) VALUES(1)";
check(sql, Some(SqliteOnConflict::Replace));

let sql = "INSERT OR REPLACE INTO test_table(id) VALUES(1)";
check(sql, Some(SqliteOnConflict::Replace));

let sql = "INSERT OR ROLLBACK INTO test_table(id) VALUES(1)";
check(sql, Some(SqliteOnConflict::Rollback));

let sql = "INSERT OR ABORT INTO test_table(id) VALUES(1)";
check(sql, Some(SqliteOnConflict::Abort));

let sql = "INSERT OR FAIL INTO test_table(id) VALUES(1)";
check(sql, Some(SqliteOnConflict::Fail));

let sql = "INSERT OR IGNORE INTO test_table(id) VALUES(1)";
check(sql, Some(SqliteOnConflict::Ignore));
}

#[test]
fn parse_update() {
let sql = "UPDATE t SET a = 1, b = 2, c = 3 WHERE d";
Expand Down

0 comments on commit add8991

Please sign in to comment.