feat: support sqlite insert or statement (apache#281)

Author: zhangli-pear

src/ast/mod.rs

@@ -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"),
+        }
+    }
+}

src/dialect/keywords.rs

@@ -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,

src/parser.rs

@@ -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,

tests/sqlparser_common.rs

@@ -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";