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

feat: support sqlite insert or statement (#281)

Browse source
This commit is contained in:
zhangli-pear 2021-02-10 04:04:54 +08:00 • committed by GitHub
parent 07342d5853
commit add8991144
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 105 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

48
src/ast/mod.rs
View file

@ -508,6 +508,8 @@ pub enum Statement {
Query(Box<Query>),
/// INSERT
Insert {
/// Only for Sqlite
or: Option<SqliteOnConflict>,
/// TABLE
table_name: ObjectName,
/// COLUMNS
@ -804,6 +806,7 @@ impl fmt::Display for Statement {
Ok(())
}
Statement::Insert {
or,
table_name,
overwrite,
partitioned,
@ -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))?;
}
@ -832,6 +839,7 @@ impl fmt::Display for Statement {
}
write!(f, "{}", source)
}
Statement::Copy {
table_name,
columns,
@ -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
src/dialect/keywords.rs
View file

@ -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,
@ -202,6 +203,7 @@ define_keywords!(
EXTENDED,
EXTERNAL,
EXTRACT,
FAIL,
FALSE,
FETCH,
FIELDS,
@ -233,6 +235,7 @@ define_keywords!(
HOUR,
IDENTITY,
IF,
IGNORE,
IN,
INDEX,
INDICATOR,

22
src/parser.rs
View file

@ -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 => {
@ -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);
@ -2758,6 +2779,7 @@ impl<'a> Parser<'a> {
let source = Box::new(self.parse_query()?);
Ok(Statement::Insert {
or,
table_name,
overwrite,
partitioned,

41
tests/sqlparser_common.rs
View file

@ -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() {
@ -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";