Add cross-DAO atomic locking to LockedContext

src/main/java/io/appform/dropwizard/sharding/dao/LockedContext.java

@@ -152,6 +152,62 @@ public LockedContext<T> apply(Consumer<T> handler) {
         return this;
     }
 
+    /**
+     * Locks multiple rows of another {@link RelationalDao} entity using SELECT FOR UPDATE NOWAIT,
+     * applies a mutator to each, and persists the changes — all within the current transaction.
+     * <p>
+     * Each {@link DetachedCriteria} must match exactly one row (point lock, no gap locking).
+     * If any criteria matches no row or the NOWAIT lock fails, the entire transaction is rolled back.
+     *
+     * @param <U>          The type of the entity to be locked and mutated.
+     * @param relationalDao The DAO for the entity to lock.
+     * @param criteriaList  A list of criteria, each targeting one row.
+     * @param mutator       Applied to each locked entity; returns the mutated entity.
+     * @return This LockedContext for further chaining.
+     */
+    public <U> LockedContext<T> lockAndMutateEach(
+            RelationalDao<U> relationalDao,
+            List<DetachedCriteria> criteriaList,
+            UnaryOperator<U> mutator) {
+        return apply(parent -> {
+            try {
+                relationalDao.lockAndMutateEach(this, criteriaList, mutator);
+            } catch (RuntimeException e) {
+                throw e;
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    /**
+     * Locks a {@link LookupDao} entity by key using SELECT FOR UPDATE, applies a mutator,
+     * and persists the change — all within the current transaction.
+     * <p>
+     * The LookupDao entity must reside on the same shard as this LockedContext
+     * (e.g. both sharded on the same user ID).
+     *
+     * @param <U>     The type of the LookupDao entity.
+     * @param lookupDao The DAO for the entity to lock.
+     * @param key       The lookup key identifying the entity.
+     * @param mutator   Applied to the locked entity.
+     * @return This LockedContext for further chaining.
+     */
+    public <U> LockedContext<T> lockAndMutate(
+            LookupDao<U> lookupDao,
+            String key,
+            Mutator<U> mutator) {
+        return apply(parent -> {
+            try {
+                lookupDao.lockAndMutate(this, key, mutator);
+            } catch (RuntimeException e) {
+                throw e;
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
     /**
      * Generates entity of type {@code U} using entityGenerator and then persists them
      *

src/main/java/io/appform/dropwizard/sharding/dao/LookupDao.java

@@ -224,6 +224,24 @@ public LockedContext<T> lockAndGetExecutor(final String id) {
         return delegate.lockAndGetExecutor(dbNamespace, id);
     }
 
+    /**
+     * Locks an entity by its lookup key using SELECT FOR UPDATE, applies a mutator, and
+     * persists the change — all within the existing transaction of the provided {@link LockedContext}.
+     *
+     * @param <U>     The entity type of the parent LockedContext.
+     * @param context The LockedContext whose transaction is joined.
+     * @param key     The lookup key identifying the entity to lock.
+     * @param mutator The mutator to apply to the locked entity.
+     * @return The mutated entity.
+     * @throws javax.persistence.EntityNotFoundException if no entity is found for the given key.
+     */
+    public <U> T lockAndMutate(
+            final LockedContext<U> context,
+            final String key,
+            final LockedContext.Mutator<T> mutator) {
+        return delegate.lockAndMutate(context, key, mutator);
+    }
+
     public ReadOnlyContext<T> readOnlyExecutor(String id) {
         return readOnlyExecutor(id, x -> x);
     }

src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantLookupDao.java

@@ -68,6 +68,7 @@
 import org.hibernate.criterion.Projections;
 import org.hibernate.criterion.Restrictions;
 
+import javax.persistence.EntityNotFoundException;
 import javax.persistence.LockModeType;
 import javax.persistence.criteria.CriteriaBuilder;
 import javax.persistence.criteria.Root;
@@ -444,6 +445,49 @@ public LockedContext<T> lockAndGetExecutor(String tenantId, final String id) {
 
     }
 
+    /**
+     * Locks an entity by its lookup key using SELECT FOR UPDATE, applies a mutator, and
+     * persists the change — all within the existing transaction of the provided {@link LockedContext}.
+     * <p>
+     * This is intended for cross-DAO chaining: e.g. locking a LookupDao entity inside a
+     * RelationalDao LockedContext when both reside on the same shard.
+     * <p>
+     * SINGLE-SHARD ONLY: the entity identified by {@code key} must map to the same shard as
+     * {@code context}. This is enforced: if {@code key} hashes to a different shard than
+     * {@code context.getShardId()}, an {@link IllegalArgumentException} is thrown before any lock
+     * is taken. Atomicity cannot be guaranteed across shards (separate databases, separate
+     * transactions), so cross-shard use is unsupported.
+     *
+     * @param <U>     The entity type of the parent LockedContext.
+     * @param context The LockedContext whose transaction is joined.
+     * @param key     The lookup key identifying the entity to lock.
+     * @param mutator The mutator to apply to the locked entity.
+     * @return The mutated entity.
+     * @throws IllegalArgumentException if {@code key} maps to a different shard than {@code context}.
+     * @throws javax.persistence.EntityNotFoundException if no entity is found for the given key.
+     */
+    <U> T lockAndMutate(
+            final LockedContext<U> context,
+            final String key,
+            final LockedContext.Mutator<T> mutator) {
+        val tenantId = context.getTenantId();
+        Preconditions.checkArgument(daos.containsKey(tenantId), "Unknown tenant: " + tenantId);
+        final int keyShardId = shardCalculator.shardId(tenantId, key);
+        Preconditions.checkArgument(keyShardId == context.getShardId(),
+                "Cross-shard lockAndMutate not allowed: key '%s' maps to shard %s but LockedContext "
+                        + "is on shard %s. The LookupDao entity must be co-sharded with the "
+                        + "LockedContext (e.g. keyed on the same userId).",
+                key, keyShardId, context.getShardId());
+        final LookupDaoPriv dao = daos.get(tenantId).get(context.getShardId());
+        final T entity = dao.getLockedForWrite(key);
+        if (entity == null) {
+            throw new EntityNotFoundException("Entity not found for key: " + key);
+        }
+        mutator.mutator(entity);
+        dao.update(entity);
+        return entity;
+    }
+
     public ReadOnlyContext<T> readOnlyExecutor(String tenantId, String id) {
         return readOnlyExecutor(tenantId, id, x -> x);
     }

src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDao.java

@@ -78,6 +78,7 @@
 import org.hibernate.criterion.Restrictions;
 import org.hibernate.query.Query;
 
+import javax.persistence.EntityNotFoundException;
 import javax.persistence.Id;
 import javax.persistence.LockModeType;
 import javax.persistence.criteria.CriteriaBuilder;
@@ -86,6 +87,7 @@
 import java.lang.reflect.Field;
 import java.util.Collection;
 import java.util.Comparator;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -571,6 +573,53 @@ <U> boolean update(
         }
     }
 
+    /**
+     * Locks multiple rows matching individual criteria using SELECT FOR UPDATE NOWAIT,
+     * applies a mutator to each, and persists the changes — all within the existing
+     * transaction of the provided {@link LockedContext}.
+     *
+     * <p>
+     * SINGLE-SHARD ONLY: every row addressed by {@code criteriaList} must reside on the same
+     * shard as {@code context} (the shard whose session/transaction is already bound). This
+     * feature exists to mutate co-sharded rows atomically; atomicity cannot be guaranteed across
+     * shards because each shard is a separate database with its own transaction. The caller is
+     * responsible for ensuring the criteria only match rows on {@code context.getShardId()}.
+     *
+     * @param <U>          The entity type of the parent LockedContext.
+     * @param context      The LockedContext whose transaction is joined.
+     * @param criteriaList A list of {@link DetachedCriteria}, each of which MUST match exactly one
+     *                     row (point lock by unique key). Zero rows throws (entity not found);
+     *                     more than one row throws {@code NonUniqueResultException}. Both roll back
+     *                     the entire transaction.
+     * @param mutator      A function applied to each locked entity; returns the mutated entity.
+     * @return The list of mutated entities in criteria order.
+     * @throws javax.persistence.EntityNotFoundException if any criteria matches no row
+     * @throws RuntimeException if the lock cannot be acquired (e.g. NOWAIT contention)
+     */
+    <U> List<T> lockAndMutateEach(
+            final LockedContext<U> context,
+            final List<DetachedCriteria> criteriaList,
+            final UnaryOperator<T> mutator) {
+        val tenantId = context.getTenantId();
+        Preconditions.checkArgument(daos.containsKey(tenantId), "Unknown tenant: " + tenantId);
+        Preconditions.checkArgument(criteriaList != null && !criteriaList.isEmpty(),
+                "criteriaList must not be null or empty");
+        final RelationalDaoPriv dao = daos.get(tenantId).get(context.getShardId());
+        final List<T> results = new ArrayList<>();
+        for (final DetachedCriteria criteria : criteriaList) {
+            // Each criteria must resolve to exactly one row: getLockedForWrite uses uniqueResult(),
+            // so >1 match throws NonUniqueResultException and rolls back the whole transaction.
+            final T entity = dao.getLockedForWrite(criteria);
+            if (entity == null) {
+                throw new EntityNotFoundException("Entity not found for criteria: " + criteria);
+            }
+            final T mutated = mutator.apply(entity);
+            dao.update(entity, mutated);
+            results.add(mutated);
+        }
+        return results;
+    }
+
     <U> List<T> select(
             MultiTenantLookupDao.ReadOnlyContext<U> context,
             DetachedCriteria criteria,

src/main/java/io/appform/dropwizard/sharding/dao/RelationalDao.java

@@ -150,6 +150,26 @@ <U> void save(LockedContext<U> context, T entity, Function<T, T> handler) {
         delegate.save(context, entity, handler);
     }
 
+    /**
+     * Locks multiple rows matching individual criteria using SELECT FOR UPDATE NOWAIT,
+     * applies a mutator to each, and persists the changes — all within the existing
+     * transaction of the provided {@link LockedContext}.
+     *
+     * @param <U>          The entity type of the parent LockedContext.
+     * @param context      The LockedContext whose transaction is joined.
+     * @param criteriaList A list of {@link DetachedCriteria}, each expected to match exactly one row.
+     * @param mutator      A function applied to each locked entity; returns the mutated entity.
+     * @return The list of mutated entities in criteria order.
+     * @throws javax.persistence.EntityNotFoundException if any criteria matches no row
+     * @throws RuntimeException if the lock cannot be acquired (e.g. NOWAIT contention)
+     */
+    public <U> List<T> lockAndMutateEach(
+            final LockedContext<U> context,
+            final List<DetachedCriteria> criteriaList,
+            final UnaryOperator<T> mutator) {
+        return delegate.lockAndMutateEach(context, criteriaList, mutator);
+    }
+
 
     /**
      * Updates an entity within a locked context using a specific ID and an updater function.