Merge pull request SeaQL#117 from nappa85/master

Initial lock support
arpancodes · Aug 31, 2021 · 3312a1f · 3312a1f
2 parents e177960 + a38696f
commit 3312a1f
Show file tree

Hide file tree

Showing 3 changed files with 138 additions and 0 deletions.
diff --git a/src/backend/query_builder.rs b/src/backend/query_builder.rs
@@ -121,6 +121,11 @@ pub trait QueryBuilder: QuotedBuilder {
             write!(sql, " OFFSET ").unwrap();
             self.prepare_value(offset, sql, collector);
         }
+
+        if let Some(lock) = &select.lock {
+            write!(sql, " ").unwrap();
+            self.prepare_select_lock(lock, sql, collector);
+        }
     }
 
     /// Translate [`UpdateStatement`] into SQL statement.
@@ -338,6 +343,24 @@ pub trait QueryBuilder: QuotedBuilder {
         .unwrap();
     }
 
+    /// Translate [`LockType`] into SQL statement.
+    fn prepare_select_lock(
+        &self,
+        select_lock: &LockType,
+        sql: &mut SqlWriter,
+        _collector: &mut dyn FnMut(Value),
+    ) {
+        write!(
+            sql,
+            "{}",
+            match select_lock {
+                LockType::Shared => "FOR SHARE",
+                LockType::Exclusive => "FOR UPDATE",
+            }
+        )
+        .unwrap();
+    }
+
     /// Translate [`SelectExpr`] into SQL statement.
     fn prepare_select_expr(
         &self,

diff --git a/src/backend/sqlite/query.rs b/src/backend/sqlite/query.rs
@@ -4,4 +4,13 @@ impl QueryBuilder for SqliteQueryBuilder {
     fn char_length_function(&self) -> &str {
         "LENGTH"
     }
+
+    fn prepare_select_lock(
+        &self,
+        _select_lock: &LockType,
+        _sql: &mut SqlWriter,
+        _collector: &mut dyn FnMut(Value),
+    ) {
+         // SQLite doesn't supports row locking
+    }
 }
diff --git a/src/query/select.rs b/src/query/select.rs
@@ -50,6 +50,7 @@ pub struct SelectStatement {
     pub(crate) orders: Vec<OrderExpr>,
     pub(crate) limit: Option<Value>,
     pub(crate) offset: Option<Value>,
+    pub(crate) lock: Option<LockType>,
 }
 
 /// List of distinct keywords that can be used in select statement
@@ -104,6 +105,7 @@ impl SelectStatement {
             orders: Vec::new(),
             limit: None,
             offset: None,
+            lock: None,
         }
     }
 
@@ -120,6 +122,7 @@ impl SelectStatement {
             orders: std::mem::take(&mut self.orders),
             limit: self.limit.take(),
             offset: self.offset.take(),
+            lock: self.lock.take(),
         }
     }
 
@@ -1268,6 +1271,102 @@ impl SelectStatement {
         self.offset = None;
         self
     }
+
+    /// Row locking (if supported).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, tests_cfg::*};
+    ///
+    /// let query = Query::select()
+    ///     .column(Char::Character)
+    ///     .from(Char::Table)
+    ///     .and_where(Expr::col(Char::FontId).eq(5))
+    ///     .lock(LockType::Exclusive)
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 FOR UPDATE"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 FOR UPDATE"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 "#
+    /// );
+    /// ```
+    pub fn lock(&mut self, lock_type: LockType) -> &mut Self {
+        self.lock = Some(lock_type);
+        self
+    }
+
+    /// Shared row locking (if supported).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, tests_cfg::*};
+    ///
+    /// let query = Query::select()
+    ///     .column(Char::Character)
+    ///     .from(Char::Table)
+    ///     .and_where(Expr::col(Char::FontId).eq(5))
+    ///     .lock_shared()
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 FOR SHARE"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 FOR SHARE"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 "#
+    /// );
+    /// ```
+    pub fn lock_shared(&mut self) -> &mut Self {
+        self.lock = Some(LockType::Shared);
+        self
+    }
+
+    /// Exclusive row locking (if supported).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{*, tests_cfg::*};
+    ///
+    /// let query = Query::select()
+    ///     .column(Char::Character)
+    ///     .from(Char::Table)
+    ///     .and_where(Expr::col(Char::FontId).eq(5))
+    ///     .lock_exclusive()
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 FOR UPDATE"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 FOR UPDATE"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 "#
+    /// );
+    /// ```
+    pub fn lock_exclusive(&mut self) -> &mut Self {
+        self.lock = Some(LockType::Exclusive);
+        self
+    }
 }
 
 impl QueryStatementBuilder for SelectStatement {
@@ -1345,3 +1444,10 @@ impl ConditionalStatement for SelectStatement {
         self
     }
 }
+
+/// List of lock types that can be used in select statement
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LockType {
+    Shared,
+    Exclusive,
+}