add all postgres row-level locking options

add FOR NO KEY UPDATE and FOR KEY SHARE
bitemyapp · Oct 24, 2024 · 6a30cb0 · 6a30cb0
1 parent d5df118
commit 6a30cb0
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 2 deletions.
diff --git a/changelog.md b/changelog.md
@@ -1,3 +1,9 @@
+3.5.13.0
+========
+- @ac251
+    - [#402](https://github.com/bitemyapp/esqueleto/pull/402)
+      - Add `forNoKeyUpdate` and `forKeyShare` locking kinds for postgres
+
 3.5.12.0
 ========
 - @csamak

diff --git a/esqueleto.cabal b/esqueleto.cabal
@@ -2,7 +2,7 @@ cabal-version: 1.12
 
 name:           esqueleto
 
-version:        3.5.12.0
+version:        3.5.13.0
 synopsis:       Type-safe EDSL for SQL queries on persistent backends.
 description:    @esqueleto@ is a bare bones, type-safe EDSL for SQL queries that works with unmodified @persistent@ SQL backends.  Its language closely resembles SQL, so you don't have to learn new concepts, just new syntax, and it's fairly easy to predict the generated SQL and optimize it for your backend. Most kinds of errors committed when writing SQL are caught as compile-time errors---although it is possible to write type-checked @esqueleto@ queries that fail at runtime.
                 .

diff --git a/src/Database/Esqueleto/Internal/Internal.hs b/src/Database/Esqueleto/Internal/Internal.hs
@@ -1485,7 +1485,9 @@ data PostgresLockingKind =
 -- Arranged in order of lock strength
 data PostgresRowLevelLockStrength =
     PostgresForUpdate
+    | PostgresForNoKeyUpdate
     | PostgresForShare
+    | PostgresForKeyShare
   deriving (Ord, Eq)
 
 data LockingOfClause where
@@ -3254,7 +3256,9 @@ makeLocking info (PostgresLockingClauses clauses) =
                     <> makeLockingBehavior (postgresOnLockedBehavior l)
             makeLockingStrength :: PostgresRowLevelLockStrength -> (TLB.Builder, [PersistValue])
             makeLockingStrength PostgresForUpdate = plain "FOR UPDATE"
+            makeLockingStrength PostgresForNoKeyUpdate = plain "FOR NO KEY UPDATE"
             makeLockingStrength PostgresForShare = plain "FOR SHARE"
+            makeLockingStrength PostgresForKeyShare = plain "FOR KEY SHARE"
 
             makeLockingBehavior :: OnLockedBehavior -> (TLB.Builder, [PersistValue])
             makeLockingBehavior NoWait = plain "NOWAIT"

diff --git a/src/Database/Esqueleto/PostgreSQL.hs b/src/Database/Esqueleto/PostgreSQL.hs
@@ -31,7 +31,9 @@ module Database.Esqueleto.PostgreSQL
     , wait
     , skipLocked
     , forUpdateOf
+    , forNoKeyUpdateOf
     , forShareOf
+    , forKeyShareOf
     , filterWhere
     , values
     -- * Internal
@@ -469,11 +471,26 @@ forUpdateOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
 forUpdateOf lockableEntities onLockedBehavior =
   putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForUpdate (Just $ LockingOfClause lockableEntities) onLockedBehavior]
 
+-- | `FOR NO KEY UPDATE OF` syntax for postgres locking
+-- allows locking of specific tables with a no key update lock in a view or join
+--
+-- @since 3.5.13.0
+forNoKeyUpdateOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
+forNoKeyUpdateOf lockableEntities onLockedBehavior =
+  putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForNoKeyUpdate (Just $ LockingOfClause lockableEntities) onLockedBehavior]
+
 -- | `FOR SHARE OF` syntax for postgres locking
 -- allows locking of specific tables with a share lock in a view or join
 --
 -- @since 3.5.9.0
-
 forShareOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
 forShareOf lockableEntities onLockedBehavior =
   putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForShare (Just $ LockingOfClause lockableEntities) onLockedBehavior]
+
+-- | `FOR KEY SHARE OF` syntax for postgres locking
+-- allows locking of specific tables with a key share lock in a view or join
+--
+-- @since 3.5.13.0
+forKeyShareOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
+forKeyShareOf lockableEntities onLockedBehavior =
+  putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForKeyShare (Just $ LockingOfClause lockableEntities) onLockedBehavior]
diff --git a/test/PostgreSQL/Test.hs b/test/PostgreSQL/Test.hs
@@ -1243,7 +1243,9 @@ testPostgresqlLocking = do
                     p <- Experimental.from $ table @Person
                     EP.forUpdateOf p EP.skipLocked
                     EP.forUpdateOf p EP.skipLocked
+                    EP.forNoKeyUpdateOf p EP.skipLocked
                     EP.forShareOf p EP.skipLocked
+                    EP.forKeyShareOf p EP.skipLocked
             conn <- ask
             let res1 = toText conn multipleLockingQuery
                 resExpected =
@@ -1253,7 +1255,9 @@ testPostgresqlLocking = do
                     ,"FROM \"Person\""
                     ,"FOR UPDATE OF \"Person\" SKIP LOCKED"
                     ,"FOR UPDATE OF \"Person\" SKIP LOCKED"
+                    ,"FOR NO KEY UPDATE OF \"Person\" SKIP LOCKED"
                     ,"FOR SHARE OF \"Person\" SKIP LOCKED"
+                    ,"FOR KEY SHARE OF \"Person\" SKIP LOCKED"
                     ]
 
             asserting $ res1 `shouldBe` resExpected