support for SQLite UPDATE ... LIMIT ...

Author: xtuc

src/ast/mod.rs

@@ -3238,6 +3238,8 @@ pub enum Statement {
         returning: Option<Vec<SelectItem>>,
         /// SQLite-specific conflict resolution clause
         or: Option<SqliteOnConflict>,
+        /// LIMIT
+        limit: Option<Expr>,
     },
     /// ```sql
     /// DELETE
@@ -4799,6 +4801,7 @@ impl fmt::Display for Statement {
                 selection,
                 returning,
                 or,
+                limit,
             } => {
                 f.write_str("UPDATE ")?;
                 if let Some(or) = or {
@@ -4832,6 +4835,10 @@ impl fmt::Display for Statement {
                     f.write_str("RETURNING")?;
                     indented_list(f, returning)?;
                 }
+                if let Some(limit) = limit {
+                    SpaceOrNewline.fmt(f)?;
+                    write!(f, "LIMIT {limit}")?;
+                }
                 Ok(())
             }
             Statement::Delete(delete) => delete.fmt(f),

src/ast/spans.rs

@@ -378,6 +378,7 @@ impl Spanned for Statement {
                 selection,
                 returning,
                 or: _,
+                limit: _,
             } => union_spans(
                 core::iter::once(table.span())
                     .chain(assignments.iter().map(|i| i.span()))

src/dialect/mod.rs

@@ -1163,6 +1163,16 @@ pub trait Dialect: Debug + Any {
     fn supports_interval_options(&self) -> bool {
         false
     }
+
+    /// Returns true if this dialect allows an optional ORDER BY and LIMIT clause
+    /// on UPDATE and DELETE statements
+    /// e.g. `UPDATE table SET value=1 LIMIT 1` (SQLite)
+    ///
+    /// For SQLite the compilation flag `SQLITE_ENABLE_UPDATE_DELETE_LIMIT` needs
+    /// to be set.
+    fn supports_update_delete_limit(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/parser/mod.rs

@@ -15116,13 +15116,20 @@ impl<'a> Parser<'a> {
         } else {
             None
         };
+        let limit =
+            if self.dialect.supports_update_delete_limit() && self.parse_keyword(Keyword::LIMIT) {
+                Some(self.parse_expr()?)
+            } else {
+                None
+            };
         Ok(Statement::Update {
             table,
             assignments,
             from,
             selection,
             returning,
             or,
+            limit,
         })
     }
 

tests/sqlparser_common.rs

@@ -515,6 +515,7 @@ fn parse_update_set_from() {
             }),
             returning: None,
             or: None,
+            limit: None
         }
     );
 
@@ -533,6 +534,7 @@ fn parse_update_with_table_alias() {
             selection,
             returning,
             or: None,
+            limit: None,
         } => {
             assert_eq!(
                 TableWithJoins {

tests/sqlparser_mysql.rs

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

tests/sqlparser_sqlite.rs

@@ -28,7 +28,7 @@ use test_utils::*;
 use sqlparser::ast::SelectItem::UnnamedExpr;
 use sqlparser::ast::Value::Placeholder;
 use sqlparser::ast::*;
-use sqlparser::dialect::{GenericDialect, SQLiteDialect};
+use sqlparser::dialect::{Dialect, GenericDialect, SQLiteDialect};
 use sqlparser::parser::{ParserError, ParserOptions};
 use sqlparser::tokenizer::Token;
 
@@ -485,7 +485,8 @@ fn parse_update_tuple_row_values() {
                 joins: vec![],
             },
             from: None,
-            returning: None
+            returning: None,
+            limit: None
         }
     );
 }
@@ -592,6 +593,41 @@ fn test_regexp_operator() {
     sqlite().verified_only_select(r#"SELECT count(*) FROM messages WHERE msg_text REGEXP '\d+'"#);
 }
 
+#[test]
+fn test_update_delete_limit() {
+    #[derive(Debug)]
+    struct SQLiteDialect;
+    impl Dialect for SQLiteDialect {
+        fn is_identifier_start(&self, ch: char) -> bool {
+            ch.is_ascii_lowercase() || ch.is_ascii_uppercase()
+        }
+
+        fn is_identifier_part(&self, ch: char) -> bool {
+            self.is_identifier_start(ch)
+        }
+
+        fn supports_update_delete_limit(&self) -> bool {
+            true
+        }
+    }
+
+    let sqlite = TestedDialects::new(vec![Box::new(SQLiteDialect {})]);
+
+    match sqlite.verified_stmt("UPDATE foo SET bar = 1 LIMIT 99") {
+        Statement::Update { limit, .. } => {
+            assert_eq!(limit, Some(Expr::value(number("99"))));
+        }
+        _ => unreachable!(),
+    }
+
+    match sqlite.verified_stmt("DELETE FROM foo LIMIT 99") {
+        Statement::Delete(Delete { limit, .. }) => {
+            assert_eq!(limit, Some(Expr::value(number("99"))));
+        }
+        _ => unreachable!(),
+    }
+}
+
 fn sqlite() -> TestedDialects {
     TestedDialects::new(vec![Box::new(SQLiteDialect {})])
 }