issues-142 full support lock in select (SeaQL#289)

* issues-142 full support lock in select

* issues-142 Fix struct name

* issues-142 Fix test

* issues-280 Revert t previous commit

src/backend/query_builder.rs

@@ -400,19 +400,37 @@ pub trait QueryBuilder: QuotedBuilder {
     /// Translate [`LockType`] into SQL statement.
     fn prepare_select_lock(
         &self,
-        select_lock: &LockType,
+        lock: &LockClause,
         sql: &mut SqlWriter,
-        _collector: &mut dyn FnMut(Value),
+        collector: &mut dyn FnMut(Value),
     ) {
         write!(
             sql,
-            "{}",
-            match select_lock {
-                LockType::Shared => "FOR SHARE",
-                LockType::Exclusive => "FOR UPDATE",
+            "FOR {}",
+            match lock.r#type {
+                LockType::Update => "UPDATE",
+                LockType::NoKeyUpdate => "NO KEY UPDATE",
+                LockType::Share => "SHARE",
+                LockType::KeyShare => "KEY SHARE",
             }
         )
         .unwrap();
+        if !lock.tables.is_empty() {
+            write!(sql, " OF ").unwrap();
+            lock.tables.iter().fold(true, |first, table_ref| {
+                if !first {
+                    write!(sql, ", ").unwrap();
+                }
+                self.prepare_table_ref(table_ref, sql, collector);
+                false
+            });
+        }
+        if let Some(behavior) = lock.behavior {
+            match behavior {
+                LockBehavior::Nowait => write!(sql, " NOWAIT").unwrap(),
+                LockBehavior::SkipLocked => write!(sql, " SKIP LOCKED").unwrap(),
+            }
+        }
     }
 
     /// Translate [`SelectExpr`] into SQL statement.

src/backend/sqlite/query.rs

@@ -7,7 +7,7 @@ impl QueryBuilder for SqliteQueryBuilder {
 
     fn prepare_select_lock(
         &self,
-        _select_lock: &LockType,
+        _select_lock: &LockClause,
         _sql: &mut SqlWriter,
         _collector: &mut dyn FnMut(Value),
     ) {

src/query/select.rs

@@ -51,7 +51,7 @@ pub struct SelectStatement {
     pub(crate) orders: Vec<OrderExpr>,
     pub(crate) limit: Option<Value>,
     pub(crate) offset: Option<Value>,
-    pub(crate) lock: Option<LockType>,
+    pub(crate) lock: Option<LockClause>,
     pub(crate) window: Option<(DynIden, WindowStatement)>,
 }
 
@@ -92,8 +92,24 @@ pub struct JoinExpr {
 /// List of lock types that can be used in select statement
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum LockType {
-    Shared,
-    Exclusive,
+    Update,
+    NoKeyUpdate,
+    Share,
+    KeyShare,
+}
+
+/// List of lock behavior can be used in select statement
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LockBehavior {
+    Nowait,
+    SkipLocked,
+}
+
+#[derive(Debug, Clone)]
+pub struct LockClause {
+    pub(crate) r#type: LockType,
+    pub(crate) tables: Vec<TableRef>,
+    pub(crate) behavior: Option<LockBehavior>,
 }
 
 /// List of union types that can be used in union clause
@@ -1702,7 +1718,7 @@ impl SelectStatement {
     ///     .column(Char::Character)
     ///     .from(Char::Table)
     ///     .and_where(Expr::col(Char::FontId).eq(5))
-    ///     .lock(LockType::Exclusive)
+    ///     .lock(LockType::Update)
     ///     .to_owned();
     ///
     /// assert_eq!(
@@ -1718,8 +1734,133 @@ impl SelectStatement {
     ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 "#
     /// );
     /// ```
-    pub fn lock(&mut self, lock_type: LockType) -> &mut Self {
-        self.lock = Some(lock_type);
+    pub fn lock(&mut self, r#type: LockType) -> &mut Self {
+        self.lock = Some(LockClause {
+            r#type,
+            tables: Vec::new(),
+            behavior: None,
+        });
+        self
+    }
+
+    /// Row locking with tables (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_with_tables(LockType::Update, vec![Glyph::Table])
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 FOR UPDATE OF `glyph`"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 FOR UPDATE OF "glyph""#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 "#
+    /// );
+    /// ```
+    pub fn lock_with_tables<T, I>(&mut self, r#type: LockType, tables: I) -> &mut Self
+    where
+        T: IntoTableRef,
+        I: IntoIterator<Item = T>,
+    {
+        self.lock = Some(LockClause {
+            r#type,
+            tables: tables.into_iter().map(|t| t.into_table_ref()).collect(),
+            behavior: None,
+        });
+        self
+    }
+
+    /// Row locking with behavior (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_with_behavior(LockType::Update, LockBehavior::Nowait)
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 FOR UPDATE NOWAIT"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 FOR UPDATE NOWAIT"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 "#
+    /// );
+    /// ```
+    pub fn lock_with_behavior(&mut self, r#type: LockType, behavior: LockBehavior) -> &mut Self {
+        self.lock = Some(LockClause {
+            r#type,
+            tables: Vec::new(),
+            behavior: Some(behavior),
+        });
+        self
+    }
+
+    /// Row locking with tables and behavior (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_with_tables_behavior(LockType::Update, vec![Glyph::Table], LockBehavior::Nowait)
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT `character` FROM `character` WHERE `font_id` = 5 FOR UPDATE OF `glyph` NOWAIT"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 FOR UPDATE OF "glyph" NOWAIT"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT "character" FROM "character" WHERE "font_id" = 5 "#
+    /// );
+    /// ```
+    pub fn lock_with_tables_behavior<T, I>(
+        &mut self,
+        r#type: LockType,
+        tables: I,
+        behavior: LockBehavior,
+    ) -> &mut Self
+    where
+        T: IntoTableRef,
+        I: IntoIterator<Item = T>,
+    {
+        self.lock = Some(LockClause {
+            r#type,
+            tables: tables.into_iter().map(|t| t.into_table_ref()).collect(),
+            behavior: Some(behavior),
+        });
         self
     }
 
@@ -1751,8 +1892,7 @@ impl SelectStatement {
     /// );
     /// ```
     pub fn lock_shared(&mut self) -> &mut Self {
-        self.lock = Some(LockType::Shared);
-        self
+        self.lock(LockType::Share)
     }
 
     /// Exclusive row locking (if supported).
@@ -1783,8 +1923,7 @@ impl SelectStatement {
     /// );
     /// ```
     pub fn lock_exclusive(&mut self) -> &mut Self {
-        self.lock = Some(LockType::Exclusive);
-        self
+        self.lock(LockType::Update)
     }
 
     /// Union with another SelectStatement that must have the same selected fields.