add all postgres row-level locking options

add FOR NO KEY UPDATE and FOR KEY SHARE
bitemyapp · Sep 13, 2024 · 6d8913d · 6d8913d
1 parent 30a5e80
commit 6d8913d
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 1 deletion.
diff --git a/changelog.md b/changelog.md
@@ -1,3 +1,9 @@
+3.5.11.3
+=======
+- @ac251
+    - [#xxx](fillmein)
+      - Add `forNoKeyUpdate` and `forKeyShare` locking kinds for postgres
+
 3.5.11.2
 ========
 - @arguri

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,22 @@ 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
+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
+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