Add support to simple locking clauses.

src/Database/Esqueleto.hs

@@ -39,7 +39,7 @@ module Database.Esqueleto
 
     -- * @esqueleto@'s Language
     Esqueleto( where_, on, groupBy, orderBy, rand, asc, desc, limit, offset
-             , distinct, distinctOn, don, distinctOnOrderBy, having
+             , distinct, distinctOn, don, distinctOnOrderBy, having, locking
              , sub_select, sub_selectDistinct, (^.), (?.)
              , val, isNothing, just, nothing, joinV, countRows, count, not_
              , (==.), (>=.), (>.), (<=.), (<.), (!=.), (&&.), (||.)
@@ -61,6 +61,7 @@ module Database.Esqueleto
   , ValueList(..)
   , OrderBy
   , DistinctOn
+  , LockingKind(..)
     -- ** Joins
   , InnerJoin(..)
   , CrossJoin(..)

src/Database/Esqueleto/Internal/Language.hs

@@ -30,6 +30,7 @@ module Database.Esqueleto.Internal.Language
   , DistinctOn
   , Update
   , Insertion
+  , LockingKind(..)
     -- * The guts
   , JoinKind(..)
   , IsJoinKind(..)
@@ -273,6 +274,15 @@ class (Functor query, Applicative query, Monad query) =>
   -- /Since: 1.2.2/
   having :: expr (Value Bool) -> query ()
 
+  -- | Add a locking clause to the query.  Please read
+  -- 'LockingKind' documentation and your RDBMS manual.
+  --
+  -- If multiple calls to 'locking' are made on the same query,
+  -- the last one is used.
+  --
+  -- /Since: 2.2.7/
+  locking :: LockingKind -> query ()
+
   -- | Execute a subquery @SELECT@ in an expression.  Returns a
   -- simple value so should be used only when the @SELECT@ query
   -- is guaranteed to return just one row.
@@ -703,6 +713,31 @@ data Update typ
 data Insertion a
 
 
+-- | Different kinds of locking clauses supported by 'locking'.
+--
+-- Note that each RDBMS has different locking support.  The
+-- constructors of this datatype specify only the /syntax/ of the
+-- locking mechanism, not its /semantics/.  For example, even
+-- though both MySQL and PostgreSQL support 'ForUpdate', there
+-- are no guarantees that they will behave the same.
+--
+-- /Since: 2.2.7/
+data LockingKind =
+    ForUpdate
+    -- ^ @FOR UPDATE@ syntax.  Supported by MySQL, Oracle and
+    -- PostgreSQL.
+    --
+    -- /Since: 2.2.7/
+  | ForShare
+    -- ^ @FOR SHARE@ syntax.  Supported by PostgreSQL.
+    --
+    -- /Since: 2.2.7/
+  | LockInShareMode
+    -- ^ @LOCK IN SHARE MODE@ syntax.  Supported by MySQL.
+    --
+    -- /Since: 2.2.7/
+
+
 -- | @FROM@ clause: bring entities into scope.
 --
 -- This function internally uses two type classes in order to

src/Database/Esqueleto/Internal/Sql.hs

@@ -58,7 +58,7 @@ import Control.Monad.Trans.Class (lift)
 import qualified Control.Monad.Trans.Reader as R
 import Data.Int (Int64)
 import Data.List (intersperse)
-import Data.Monoid (Monoid(..), (<>))
+import Data.Monoid (Last(..), Monoid(..), (<>))
 import Data.Proxy (Proxy(..))
 import Database.Esqueleto.Internal.PersistentImport
 import Database.Persist.Sql.Util (
@@ -111,12 +111,13 @@ data SideData = SideData { sdDistinctClause :: !DistinctClause
                          , sdHavingClause   :: !HavingClause
                          , sdOrderByClause  :: ![OrderByClause]
                          , sdLimitClause    :: !LimitClause
+                         , sdLockingClause  :: !LockingClause
                          }
 
 instance Monoid SideData where
-  mempty = SideData mempty mempty mempty mempty mempty mempty mempty mempty
+  mempty = SideData mempty mempty mempty mempty mempty mempty mempty mempty mempty
-  SideData d f s w g h o l `mappend` SideData d' f' s' w' g' h' o' l' =
+  SideData d f s w g h o l k `mappend` SideData d' f' s' w' g' h' o' l' k' =
-    SideData (d <> d') (f <> f') (s <> s') (w <> w') (g <> g') (h <> h') (o <> o') (l <> l')
+    SideData (d <> d') (f <> f') (s <> s') (w <> w') (g <> g') (h <> h') (o <> o') (l <> l') (k <> k')
 
 
 -- | The @DISTINCT@ "clause".
@@ -211,6 +212,10 @@ instance Monoid LimitClause where
     -- "reversed" arguments.
 
 
+-- | A locking clause.
+type LockingClause = Last LockingKind
+
+
 ----------------------------------------------------------------------
 
 
@@ -392,6 +397,8 @@ instance Esqueleto SqlQuery SqlExpr SqlBackend where
 
   having expr = Q $ W.tell mempty { sdHavingClause = Where expr }
 
+  locking kind = Q $ W.tell mempty { sdLockingClause = Last (Just kind) }
+
   orderBy exprs = Q $ W.tell mempty { sdOrderByClause = exprs }
   asc  = EOrderBy ASC
   desc = EOrderBy DESC
@@ -956,7 +963,8 @@ toRawSql mode (conn, firstIdentState) query =
                groupByClause
                havingClause
                orderByClauses
-               limitClause = sd
+               limitClause
+               lockingClause = sd
       -- Pass the finalIdentState (containing all identifiers
       -- that were used) to the subsequent calls.  This ensures
       -- that no name clashes will occur on subqueries that may
@@ -972,6 +980,7 @@ toRawSql mode (conn, firstIdentState) query =
       , makeHaving     info havingClause
       , makeOrderBy    info orderByClauses
       , makeLimit      info limitClause orderByClauses
+      , makeLocking         lockingClause
       ]
 
 
@@ -1112,6 +1121,15 @@ makeLimit (conn, _) (Limit ml mo) orderByClauses =
   in (TLB.fromText limitRaw, mempty)
 
 
+makeLocking :: LockingClause -> (TLB.Builder, [PersistValue])
+makeLocking = flip (,) [] . maybe mempty toTLB . getLast
+  where
+    toTLB ForUpdate       = "\nFOR UPDATE"
+    toTLB ForShare        = "\nFOR SHARE"
+    toTLB LockInShareMode = "\nLOCK IN SHARE MODE"
+
+
+
 parens :: TLB.Builder -> TLB.Builder
 parens b = "(" <> (b <> ")")
 

test/Test.hs

@@ -26,7 +26,9 @@ import Control.Monad.IO.Class (MonadIO(liftIO))
 import Control.Monad.Logger (MonadLogger(..), runStderrLoggingT, runNoLoggingT)
 import Control.Monad.Trans.Control (MonadBaseControl(..))
 import Control.Monad.Trans.Reader (ReaderT)
+import Data.Char (toLower, toUpper)
 import Data.List (sortBy)
+import Data.Monoid ((<>))
 import Data.Ord (comparing)
 import Database.Esqueleto
 #if   defined (WITH_POSTGRESQL)
@@ -48,9 +50,10 @@ import Database.Persist.TH
 import Test.Hspec
 
 import qualified Control.Monad.Trans.Resource as R
-import qualified Data.Set as S
 import qualified Data.List as L
-import Data.Char (toLower, toUpper)
+import qualified Data.Set as S
+import qualified Data.Text.Lazy.Builder as TLB
+import qualified Database.Esqueleto.Internal.Sql as EI
 
 
 -- Test schema
@@ -1268,6 +1271,48 @@ main = do
 
           liftIO $ ret `shouldBe` [ Value (3) ]
 
+    describe "locking" $ do
+      -- The locking clause is the last one, so try to use many
+      -- others to test if it's at the right position.  We don't
+      -- care about the text of the rest, nor with the RDBMS'
+      -- reaction to the clause.
+      let sanityCheck kind syntax = do
+            let complexQuery =
+                  from $ \(p1 `InnerJoin` p2) -> do
+                  on (p1 ^. PersonName ==. p2 ^. PersonName)
+                  where_ (p1 ^. PersonFavNum >. val 2)
+                  orderBy [desc (p2 ^. PersonAge)]
+                  limit 3
+                  offset 9
+                  groupBy (p1 ^. PersonId)
+                  having (countRows <. val (0 :: Int))
+                  return (p1, p2)
+                queryWithClause1 = do
+                  r <- complexQuery
+                  locking kind
+                  return r
+                queryWithClause2 = do
+                  locking ForUpdate
+                  r <- complexQuery
+                  locking ForShare
+                  locking kind
+                  return r
+                queryWithClause3 = do
+                  locking kind
+                  complexQuery
+                toText conn q =
+                  let (tlb, _) = EI.toRawSql EI.SELECT (conn, EI.initialIdentState) q
+                  in TLB.toLazyText tlb
+            [complex, with1, with2, with3] <-
+              runNoLoggingT $ withConn $ \conn -> return $
+                map (toText conn) [complexQuery, queryWithClause1, queryWithClause2, queryWithClause3]
+            let expected = complex <> "\n" <> syntax
+            (with1, with2, with3) `shouldBe` (expected, expected, expected)
+
+      it "looks sane for ForUpdate"       $ sanityCheck ForUpdate       "FOR UPDATE"
+      it "looks sane for ForShare"        $ sanityCheck ForShare        "FOR SHARE"
+      it "looks sane for LockInShareMode" $ sanityCheck LockInShareMode "LOCK IN SHARE MODE"
+
 
 ----------------------------------------------------------------------
 
@@ -1321,23 +1366,29 @@ verbose = True
 
 
 run_worker :: RunDbMonad m => SqlPersistT (R.ResourceT m) a -> m a
-run_worker act =
+run_worker act = withConn $ runSqlConn (migrateIt >> act)
+
+
+migrateIt :: RunDbMonad m => SqlPersistT (R.ResourceT m) ()
+migrateIt = do
+  void $ runMigrationSilent migrateAll
+#if defined (WITH_POSTGRESQL) || defined (WITH_MYSQL)
+  cleanDB
+#endif
+
+
+withConn :: RunDbMonad m => (SqlBackend -> R.ResourceT m a) -> m a
+withConn =
   R.runResourceT .
 #if defined(WITH_POSTGRESQL)
-  withPostgresqlConn "host=localhost port=5432 user=test dbname=test" .
+  withPostgresqlConn "host=localhost port=5432 user=test dbname=test"
 #elif defined (WITH_MYSQL)
   withMySQLConn defaultConnectInfo
     { connectHost     = "localhost"
     , connectUser     = "test"
     , connectPassword = "test"
     , connectDatabase = "test"
-    } .
+    }
 #else
-  withSqliteConn ":memory:" .
+  withSqliteConn ":memory:"
-#endif
-  runSqlConn .
-#if defined (WITH_POSTGRESQL) || defined (WITH_MYSQL)
-  (runMigrationSilent migrateAll >>) $ (cleanDB >> act)
-#else
-  (runMigrationSilent migrateAll >>) $ act
 #endif