feat: Add setLockMode, forUpdate, and forShare query methods

docs/basic_queries.md

@@ -80,6 +80,60 @@ Some operations that do require the column context are also available outside th
 - `applyProjection`
 - `minBy` / `maxBy`
 
+## Lock Modes
+
+Yawn supports pessimistic locking through the `setLockMode` (or the `forUpdate` and `forShare` helpers). These are useful when you need to prevent concurrent
+modifications to rows you're reading.
+
+### forUpdate (PESSIMISTIC_WRITE)
+
+Use `forUpdate()` when you intend to update the selected rows. This adds `FOR UPDATE` to the SQL query, acquiring an exclusive lock:
+
+```kotlin
+val results = yawn.query(BookTable) { books ->
+    addEq(books.name, "The Hobbit")
+}
+    .forUpdate()
+    .list()
+```
+
+This is commonly used in event consumers and batch update operations where you need to prevent concurrent modifications.
+
+### forShare (PESSIMISTIC_READ)
+
+Use `forShare()` for "find or create" patterns where you need to prevent concurrent creates but allow concurrent reads. This adds `FOR SHARE` to the SQL query:
+
+```kotlin
+val existing = yawn.query(BookTable) { books ->
+    addEq(books.token, bookToken)
+}
+    .forShare()
+    .uniqueResult()
+
+if (existing == null) {
+    // Safe to create - other transactions will wait
+    session.save(DbBook(...))
+}
+```
+
+### Explicit Lock Mode
+
+You can also use `setLockMode` directly with any `YawnLockMode` value:
+
+```kotlin
+val results = yawn.query(BookTable) { books ->
+    addEq(books.name, "The Hobbit")
+}
+    .setLockMode(YawnLockMode.PESSIMISTIC_WRITE)
+    .list()
+```
+
+Available lock modes:
+
+- `YawnLockMode.NONE` - No lock (default behavior)
+- `YawnLockMode.PESSIMISTIC_READ` - Shared lock (`FOR SHARE`)
+- `YawnLockMode.PESSIMISTIC_WRITE` - Exclusive lock (`FOR UPDATE`)
+
 ## Terminal Operations
 
 Terminal operations such as `uniqueResult` or `list` can only be called outside the lambda, and they do not return the builder as they terminate the query.

yawn-api/src/main/kotlin/com/faire/yawn/criteria/builder/BaseTypeSafeCriteriaBuilder.kt

@@ -4,6 +4,7 @@ import com.faire.yawn.YawnTableDef
 import com.faire.yawn.pagination.Page
 import com.faire.yawn.pagination.PageNumber
 import com.faire.yawn.query.CompiledYawnQuery
+import com.faire.yawn.query.YawnLockMode
 import com.faire.yawn.query.YawnQuery
 import com.faire.yawn.query.YawnQueryFactory
 import com.faire.yawn.query.YawnQueryHint
@@ -212,4 +213,45 @@ abstract class BaseTypeSafeCriteriaBuilder<
         query.queryHints.add(YawnQueryHint(hint))
         return builderReturn()
     }
+
+    /**
+     * Sets the lock mode for this query.
+     *
+     * Lock modes control how the database handles concurrent access to the selected rows.
+     * When set, the generated SQL will include the appropriate locking clause.
+     *
+     * @param lockMode The lock mode to use.
+     * @return this builder for chaining.
+     * @see YawnLockMode
+     */
+    fun setLockMode(lockMode: YawnLockMode): CRITERIA {
+        query.lockMode = lockMode
+        return builderReturn()
+    }
+
+    /**
+     * Locks selected rows with a PESSIMISTIC_WRITE lock (SQL: `FOR UPDATE`).
+     *
+     * Use when you intend to update the selected rows and need to prevent concurrent
+     * modifications. This is the strongest lock mode and will block other transactions
+     * from reading (with locks) or writing to the locked rows until this transaction completes.
+     *
+     * @return this builder for chaining.
+     */
+    fun forUpdate(): CRITERIA {
+        return setLockMode(YawnLockMode.PESSIMISTIC_WRITE)
+    }
+
+    /**
+     * Locks selected rows with a PESSIMISTIC_READ lock (SQL: `FOR SHARE`).
+     *
+     * Use for "find or create" patterns where you need to prevent concurrent creates
+     * but allow concurrent reads. This lock allows other transactions to read the rows
+     * but prevents them from acquiring an exclusive lock until this transaction completes.
+     *
+     * @return this builder for chaining.
+     */
+    fun forShare(): CRITERIA {
+        return setLockMode(YawnLockMode.PESSIMISTIC_READ)
+    }
 }

yawn-api/src/main/kotlin/com/faire/yawn/query/YawnLockMode.kt

@@ -0,0 +1,42 @@
+package com.faire.yawn.query
+
+/**
+ * Lock modes for Yawn queries, abstracting over SQL's locking mechanisms.
+ *
+ * These lock modes control how the database handles concurrent access to selected rows.
+ * When a lock mode is set, the generated SQL will include the appropriate locking clause
+ * (e.g., `FOR UPDATE` or `FOR SHARE`).
+ *
+ * @see org.hibernate.LockMode
+ */
+enum class YawnLockMode {
+    /**
+     * No lock. This is the default behavior.
+     */
+    NONE,
+
+    /**
+     * A shared lock (SQL: `FOR SHARE`).
+     *
+     * Use for "find or create" patterns where you need to prevent concurrent creates
+     * but allow concurrent reads. This lock allows other transactions to read the rows
+     * but prevents them from acquiring an exclusive lock until this transaction completes.
+     *
+     * Example use case: Multiple execution paths (e.g., event consumption, multiple jobs)
+     * are calling a "find or create" path. Using this lock ensures the entity is only
+     * created once by making concurrent transactions wait.
+     */
+    PESSIMISTIC_READ,
+
+    /**
+     * An exclusive lock (SQL: `FOR UPDATE`).
+     *
+     * Use when you intend to update the selected rows and want to prevent
+     * concurrent modifications. This is the strongest lock mode and will block
+     * other transactions from reading (with locks) or writing to the locked rows.
+     *
+     * Example use case: Event consumers that replicate data use this lock to prevent
+     * concurrent modifications during batch updates.
+     */
+    PESSIMISTIC_WRITE,
+}

yawn-api/src/main/kotlin/com/faire/yawn/query/YawnQuery.kt

@@ -21,6 +21,7 @@ data class YawnQuery<SOURCE : Any, T : Any>(
 
     var offset: Int? = null,
     var maxResults: Int? = null,
+    var lockMode: YawnLockMode? = null,
 ) : YawnCriteriaQuery<SOURCE, T> {
 
     override fun addCriterion(criterion: YawnQueryCriterion<SOURCE>) {

yawn-database-test/src/main/kotlin/com/faire/yawn/setup/hibernate/YawnTestCompiledQuery.kt

@@ -4,7 +4,7 @@ import com.faire.yawn.query.CompiledYawnQuery
 import org.hibernate.Criteria
 
 internal class YawnTestCompiledQuery<T>(
-    private val rawQuery: Criteria,
+    val rawQuery: Criteria,
 ) : CompiledYawnQuery<T> {
     override fun list(): List<T> {
         @Suppress("UNCHECKED_CAST")

yawn-database-test/src/main/kotlin/com/faire/yawn/setup/hibernate/YawnTestQueryFactory.kt

@@ -3,9 +3,11 @@ package com.faire.yawn.setup.hibernate
 import com.faire.yawn.YawnTableDef
 import com.faire.yawn.query.CompiledYawnQuery
 import com.faire.yawn.query.YawnCompilationContext
+import com.faire.yawn.query.YawnLockMode
 import com.faire.yawn.query.YawnQuery
 import com.faire.yawn.query.YawnQueryFactory
 import com.faire.yawn.query.YawnQueryRestriction.And
+import org.hibernate.LockMode
 import org.hibernate.Session
 
 internal class YawnTestQueryFactory(
@@ -60,6 +62,20 @@ internal class YawnTestQueryFactory(
             rawQuery.setProjection(hibernateProjection)
         }
 
+        val lockMode = query.lockMode
+        if (lockMode != null) {
+            rawQuery.setLockMode(lockMode.toHibernateLockMode())
+        }
+
         return YawnTestCompiledQuery(rawQuery)
     }
 }
+
+/**
+ * Maps [YawnLockMode] to Hibernate's [LockMode].
+ */
+private fun YawnLockMode.toHibernateLockMode(): LockMode = when (this) {
+    YawnLockMode.NONE -> LockMode.NONE
+    YawnLockMode.PESSIMISTIC_READ -> LockMode.PESSIMISTIC_READ
+    YawnLockMode.PESSIMISTIC_WRITE -> LockMode.PESSIMISTIC_WRITE
+}

yawn-database-test/src/test/kotlin/com/faire/yawn/database/YawnLockModeTest.kt

@@ -0,0 +1,111 @@
+package com.faire.yawn.database
+
+import com.faire.yawn.YawnTableDef
+import com.faire.yawn.criteria.builder.TypeSafeCriteriaBuilder
+import com.faire.yawn.query.YawnLockMode
+import com.faire.yawn.setup.entities.BaseEntity
+import com.faire.yawn.setup.entities.BookTable
+import com.faire.yawn.setup.hibernate.YawnTestCompiledQuery
+import org.assertj.core.api.Assertions.assertThat
+import org.hibernate.LockMode
+import org.hibernate.internal.CriteriaImpl
+import org.junit.jupiter.api.Test
+
+internal class YawnLockModeTest : BaseYawnDatabaseTest() {
+    @Test
+    fun `forUpdate sets PESSIMISTIC_WRITE lock mode`() {
+        transactor.open { session ->
+            val compiledQuery = getCompiledQuery {
+                session.query(BookTable) { books ->
+                    addEq(books.name, "The Hobbit")
+                }.forUpdate()
+            }
+
+            assertThat(compiledQuery.hibernateCriteriaImpl.lockModes.values).containsOnly(LockMode.PESSIMISTIC_WRITE)
+            assertThat(compiledQuery.list().single().name).isEqualTo("The Hobbit")
+        }
+    }
+
+    @Test
+    fun `forShare sets PESSIMISTIC_READ lock mode`() {
+        transactor.open { session ->
+            val compiledQuery = getCompiledQuery {
+                session.query(BookTable) { books ->
+                    addEq(books.name, "The Hobbit")
+                }.forShare()
+            }
+
+            assertThat(compiledQuery.hibernateCriteriaImpl.lockModes.values).containsOnly(LockMode.PESSIMISTIC_READ)
+            assertThat(compiledQuery.list().single().name).isEqualTo("The Hobbit")
+        }
+    }
+
+    @Test
+    fun `setLockMode with PESSIMISTIC_WRITE`() {
+        transactor.open { session ->
+            val compiledQuery = getCompiledQuery {
+                session.query(BookTable) { books ->
+                    addEq(books.name, "The Hobbit")
+                }.setLockMode(YawnLockMode.PESSIMISTIC_WRITE)
+            }
+
+            assertThat(compiledQuery.hibernateCriteriaImpl.lockModes.values).containsOnly(LockMode.PESSIMISTIC_WRITE)
+            assertThat(compiledQuery.list().single().name).isEqualTo("The Hobbit")
+        }
+    }
+
+    @Test
+    fun `setLockMode with PESSIMISTIC_READ`() {
+        transactor.open { session ->
+            val compiledQuery = getCompiledQuery {
+                session.query(BookTable) { books ->
+                    addEq(books.name, "The Hobbit")
+                }.setLockMode(YawnLockMode.PESSIMISTIC_READ)
+            }
+
+            assertThat(compiledQuery.hibernateCriteriaImpl.lockModes.values).containsOnly(LockMode.PESSIMISTIC_READ)
+            assertThat(compiledQuery.list().single().name).isEqualTo("The Hobbit")
+        }
+    }
+
+    @Test
+    fun `setLockMode with NONE explicitly set`() {
+        transactor.open { session ->
+            val compiledQuery = getCompiledQuery {
+                session.query(BookTable) { books ->
+                    addEq(books.name, "The Hobbit")
+                }.setLockMode(YawnLockMode.NONE)
+            }
+
+            assertThat(compiledQuery.hibernateCriteriaImpl.lockModes.values).containsOnly(LockMode.NONE)
+            assertThat(compiledQuery.list().single().name).isEqualTo("The Hobbit")
+        }
+    }
+
+    @Test
+    fun `forUpdate can be combined with maxResults and offset`() {
+        transactor.open { session ->
+            val compiledQuery = getCompiledQuery {
+                session.query(BookTable) { books ->
+                    orderAsc(books.name)
+                }
+                    .forUpdate()
+                    .maxResults(2)
+                    .offset(1)
+            }
+
+            assertThat(compiledQuery.hibernateCriteriaImpl.lockModes.values).containsOnly(LockMode.PESSIMISTIC_WRITE)
+            assertThat(compiledQuery.list().map { it.name })
+                .containsExactly("Lord of the Rings", "The Emperor's New Clothes")
+        }
+    }
+
+    private fun <T : BaseEntity<T>, DEF : YawnTableDef<T, T>> getCompiledQuery(
+        queryProvider: () -> TypeSafeCriteriaBuilder<T, DEF>,
+    ): YawnTestCompiledQuery<T> {
+        return queryProvider().compile() as YawnTestCompiledQuery<T>
+    }
+}
+
+private val YawnTestCompiledQuery<*>.hibernateCriteriaImpl: CriteriaImpl
+    get() = rawQuery as CriteriaImpl