Added locking clause for postgres-only

Author: gammelalf

src/lib.rs

@@ -299,14 +299,17 @@ impl DBImpl {
     ) -> Select<'until_build, 'post_build> {
         Select {
             db_impl: *self,
-            join_tables: joins,
             resulting_columns: columns,
+            from_clause,
+            join_tables: joins,
+            order_by_clause,
+
             limit: None,
             offset: None,
-            from_clause,
             where_clause: None,
             distinct: false,
-            order_by_clause,
+            #[cfg(feature = "postgres-only")]
+            locking_clause: None,
         }
     }
 

src/select.rs

@@ -17,13 +17,54 @@ pub struct Select<'until_build, 'post_query> {
     pub(crate) resulting_columns: &'until_build [SelectColumnImpl<'until_build>],
     pub(crate) from_clause: &'until_build str,
     pub(crate) join_tables: &'until_build [JoinTableImpl<'until_build, 'post_query>],
-    pub(crate) where_clause: Option<&'until_build Condition<'post_query>>,
+    pub(crate) order_by_clause: &'until_build [OrderByEntry<'until_build>],
 
     // Set by builder
-    pub(crate) order_by_clause: &'until_build [OrderByEntry<'until_build>],
+    pub(crate) where_clause: Option<&'until_build Condition<'post_query>>,
     pub(crate) distinct: bool,
     pub(crate) limit: Option<u64>,
     pub(crate) offset: Option<u64>,
+
+    #[cfg(feature = "postgres-only")]
+    pub(crate) locking_clause: Option<LockingClause>,
+}
+
+/// A `SELECT` statement's locking clause
+#[cfg(feature = "postgres-only")]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
+pub struct LockingClause {
+    /// Strength of the row lock
+    pub strength: LockStrength,
+
+    /// How to acquire the row lock
+    pub acquire: LockAcquire,
+}
+
+#[cfg(feature = "postgres-only")]
+/// Strength of a `SELECT` statement's row lock
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
+pub enum LockStrength {
+    /// `SELECT ... FOR UPDATE;`
+    Update,
+    /// `SELECT ... FOR NO KEY UPDATE;`
+    NoKeyUpdate,
+    /// `SELECT ... FOR SHARE;`
+    Share,
+    /// `SELECT ... FOR KEY SHARE;`
+    KeyShare,
+}
+
+#[cfg(feature = "postgres-only")]
+/// How to acquire a `SELECT` statement's row lock
+#[derive(Default, Copy, Clone, Debug, Eq, PartialEq, Hash)]
+pub enum LockAcquire {
+    /// `SELECT ... FOR <strength>;`
+    #[default]
+    Wait,
+    /// `SELECT ... FOR <strength> NOWAIT;`
+    NoWait,
+    /// `SELECT ... FOR <strength> SKIP LOCKED;`
+    SkipLocked,
 }
 
 impl<'until_build, 'post_build> Select<'until_build, 'post_build> {
@@ -46,6 +87,13 @@ impl<'until_build, 'post_build> Select<'until_build, 'post_build> {
         self
     }
 
+    /// Set a lock on the resulting rows
+    #[cfg(feature = "postgres-only")]
+    pub fn locking_clause(mut self, locking: LockingClause) -> Self {
+        self.locking_clause = Some(locking);
+        self
+    }
+
     /// Build the select query
     pub fn build(self) -> (String, Vec<Value<'post_build>>) {
         let mut sql;
@@ -166,6 +214,26 @@ impl<'until_build, 'post_build> Select<'until_build, 'post_build> {
                     }
                 };
 
+                #[cfg(feature = "postgres-only")]
+                if let Some(locking) = self.locking_clause {
+                    write!(
+                        sql,
+                        " FOR {}{}",
+                        match locking.strength {
+                            LockStrength::Update => "UPDATE",
+                            LockStrength::NoKeyUpdate => "NO KEY UPDATE",
+                            LockStrength::Share => "SHARE",
+                            LockStrength::KeyShare => "KEY SHARE",
+                        },
+                        match locking.acquire {
+                            LockAcquire::Wait => "",
+                            LockAcquire::NoWait => " NOWAIT",
+                            LockAcquire::SkipLocked => " SKIP LOCKED",
+                        }
+                    )
+                    .unwrap();
+                }
+
                 write!(sql, ";").unwrap();
             }
         }