diff --git a/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantLookupDao.java b/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantLookupDao.java
index e237ce5e..2f933e80 100644
--- a/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantLookupDao.java
+++ b/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantLookupDao.java
@@ -774,9 +774,9 @@ public <U> U run(String tenantId, DetachedCriteria criteria,
                 .boxed()
                 .collect(Collectors.toMap(Function.identity(), shardId -> {
                     final LookupDaoPriv dao = daos.get(tenantId).get(shardId);
-                    OpContext<List<T>> opContext = RunWithCriteria.<List<T>>builder()
+                    OpContext<List<T>> opContext = RunWithCriteria.<List<T>, DetachedCriteria>builder()
                             .handler(dao::run)
-                            .detachedCriteria(criteria)
+                            .criteria(criteria)
                             .build();
                     return transactionExecutor.get(tenantId).execute(dao.sessionFactory,
                             true, "run",
diff --git a/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDao.java b/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDao.java
index 6cae95e1..21018dd0 100644
--- a/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDao.java
+++ b/src/main/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDao.java
@@ -162,6 +162,11 @@ T get(DetachedCriteria criteria) {
             return uniqueResult(criteria.getExecutableCriteria(currentSession()));
         }
 
+        T get(final QuerySpec<T, T> querySpec) {
+            val q = InternalUtils.createQuery(currentSession(), entityClass, querySpec);
+            return uniqueResult(q.setLockMode(LockModeType.NONE));
+        }
+
         T getLocked(Object lookupKey, UnaryOperator<Criteria> criteriaUpdater, LockMode lockMode) {
             Criteria criteria = criteriaUpdater.apply(currentSession()
                     .createCriteria(entityClass)
@@ -235,6 +240,17 @@ List run(DetachedCriteria criteria) {
                     .list();
         }
 
+        /**
+         * Run a query inside this shard using QuerySpec and return the matching list.
+         *
+         * @param querySpec QuerySpec defining the query criteria.
+         * @return List of elements or empty list if none found
+         */
+        List<T> run(QuerySpec<T, T> querySpec) {
+            val query = InternalUtils.createQuery(currentSession(), entityClass, querySpec);
+            return list(query);
+        }
+
         long count(final DetachedCriteria criteria) {
             return (long) criteria.getExecutableCriteria(currentSession())
                     .setProjection(Projections.rowCount())
@@ -424,7 +440,7 @@ public Optional<T> createOrUpdate(String tenantId,
         Preconditions.checkArgument(daos.containsKey(tenantId), "Unknown tenant: " + tenantId);
         int shardId = shardCalculator.shardId(tenantId, parentKey);
         RelationalDaoPriv dao = daos.get(tenantId).get(shardId);
-        val opContext = CreateOrUpdate.<T>builder()
+        val opContext = CreateOrUpdate.<T, DetachedCriteria>builder()
                 .criteria(selectionCriteria)
                 .getLockedForWrite(dao::getLockedForWrite)
                 .entityGenerator(entityGenerator)
@@ -441,6 +457,31 @@ public Optional<T> createOrUpdate(String tenantId,
                 shardId));
     }
 
+    public Optional<T> createOrUpdate(String tenantId,
+                                  final String parentKey,
+                                  final QuerySpec<T, T> querySpec,
+                                  final UnaryOperator<T> updater,
+                                  final Supplier<T> entityGenerator) {
+        Preconditions.checkArgument(daos.containsKey(tenantId), "Unknown tenant: " + tenantId);
+        int shardId = shardCalculator.shardId(tenantId, parentKey);
+        RelationalDaoPriv dao = daos.get(tenantId).get(shardId);
+        val opContext = CreateOrUpdate.<T, QuerySpec<T, T>>builder()
+                .criteria(querySpec)
+                .getLockedForWrite(dao::getLockedForWrite)
+                .entityGenerator(entityGenerator)
+                .saver(dao::save)
+                .mutator(updater)
+                .updater(dao::update)
+                .getter(dao::get)
+                .build();
+        return Optional.of(transactionExecutor.get(tenantId).execute(
+                dao.sessionFactory,
+                false,
+                "createOrUpdate",
+                opContext,
+                shardId));
+    }
+
     public <U> void save(LockedContext<U> context, T entity) {
         val tenantId = context.getTenantId();
         RelationalDaoPriv dao = daos.get(tenantId).get(context.getShardId());
@@ -814,8 +855,49 @@ public <U> U run(String tenantId, DetachedCriteria criteria,
                 .boxed()
                 .collect(Collectors.toMap(Function.identity(), shardId -> {
                     final RelationalDaoPriv dao = daos.get(tenantId).get(shardId);
-                    OpContext<List> opContext = RunWithCriteria.<List>builder()
-                            .detachedCriteria(criteria).handler(dao::run).build();
+                    OpContext<List> opContext = RunWithCriteria.<List, DetachedCriteria>builder()
+                            .criteria(criteria).handler(dao::run).build();
+                    return transactionExecutor.get(tenantId).execute(dao.sessionFactory,
+                            true,
+                            "run",
+                            opContext,
+                            shardId);
+                }));
+        return translator.apply(output);
+    }
+
+    /**
+     * Run arbitrary read-only queries on all shards using QuerySpec and return results.
+     *
+     * @param tenantId  The tenant ID associated with the entity.
+     * @param querySpec The QuerySpec defining query criteria. Typically, a grouping or counting query
+     * @return A map of shard vs result-list
+     */
+    public Map<Integer, List<T>> run(String tenantId, QuerySpec<T, T> querySpec) {
+        return run(tenantId, querySpec, Function.identity());
+    }
+
+
+    /**
+     * Run read-only queries on all shards using QuerySpec and transform them into required types
+     *
+     * @param tenantId   The tenant ID associated with the entity.
+     * @param querySpec  The QuerySpec defining query criteria. Typically, a grouping or counting query
+     * @param translator A method to transform results to required type
+     * @param <U>        Return type
+     * @return Translated result
+     */
+    public <U> U run(String tenantId, QuerySpec<T, T> querySpec,
+                     Function<Map<Integer, List<T>>, U> translator) {
+        Preconditions.checkArgument(daos.containsKey(tenantId), "Unknown tenant: " + tenantId);
+        val output = IntStream.range(0, daos.get(tenantId).size())
+                .boxed()
+                .collect(Collectors.toMap(Function.identity(), shardId -> {
+                    final RelationalDaoPriv dao = daos.get(tenantId).get(shardId);
+                    OpContext<List<T>> opContext = RunWithCriteria.<List<T>, QuerySpec<T, T>>builder()
+                            .criteria(querySpec)
+                            .handler(dao::run)
+                            .build();
                     return transactionExecutor.get(tenantId).execute(dao.sessionFactory,
                             true,
                             "run",
diff --git a/src/main/java/io/appform/dropwizard/sharding/dao/RelationalDao.java b/src/main/java/io/appform/dropwizard/sharding/dao/RelationalDao.java
index dd03820c..fa0ec155 100644
--- a/src/main/java/io/appform/dropwizard/sharding/dao/RelationalDao.java
+++ b/src/main/java/io/appform/dropwizard/sharding/dao/RelationalDao.java
@@ -142,6 +142,14 @@ public Optional<T> createOrUpdate(
         return delegate.createOrUpdate(tenantId, parentKey, selectionCriteria, updater, entityGenerator);
     }
 
+    public Optional<T> createOrUpdate(
+            final String parentKey,
+            final QuerySpec<T, T> querySpec,
+            final UnaryOperator<T> updater,
+            final Supplier<T> entityGenerator) {
+        return delegate.createOrUpdate(tenantId, parentKey, querySpec, updater, entityGenerator);
+    }
+
     public <U> void save(LockedContext<U> context, T entity) {
         delegate.save(context, entity);
     }
@@ -362,6 +370,28 @@ public <U> U run(DetachedCriteria criteria, Function<Map<Integer, List>, U> tran
         return delegate.run(tenantId, criteria, translator);
     }
 
+    /**
+     * Run arbitrary read-only queries on all shards using QuerySpec and return results.
+     *
+     * @param querySpec The QuerySpec defining query criteria. Typically, a grouping or counting query
+     * @return A map of shard vs result-list
+     */
+    public Map<Integer, List<T>> run(QuerySpec<T, T> querySpec) {
+        return delegate.run(tenantId, querySpec);
+    }
+
+    /**
+     * Run read-only queries on all shards using QuerySpec and transform them into required types
+     *
+     * @param querySpec  The QuerySpec defining query criteria. Typically, a grouping or counting query
+     * @param translator A method to transform results to required type
+     * @param <U>        Return type
+     * @return Translated result
+     */
+    public <U> U run(QuerySpec<T, T> querySpec, Function<Map<Integer, List<T>>, U> translator) {
+        return delegate.run(tenantId, querySpec, translator);
+    }
+
     public <U> U runInSession(String id, Function<Session, U> handler) {
         return delegate.runInSession(tenantId, id, handler);
     }
diff --git a/src/main/java/io/appform/dropwizard/sharding/dao/operations/OpContext.java b/src/main/java/io/appform/dropwizard/sharding/dao/operations/OpContext.java
index c3c7543c..14adc408 100644
--- a/src/main/java/io/appform/dropwizard/sharding/dao/operations/OpContext.java
+++ b/src/main/java/io/appform/dropwizard/sharding/dao/operations/OpContext.java
@@ -55,7 +55,7 @@ public interface OpContextVisitor<P> {
 
     <T> P visit(RunInSession<T> opContext);
 
-    <T> P visit(RunWithCriteria<T> opContext);
+    <T, C> P visit(RunWithCriteria<T, C> opContext);
 
     P visit(DeleteByLookupKey opContext);
 
@@ -65,12 +65,11 @@ public interface OpContextVisitor<P> {
 
     <T> P visit(CreateOrUpdateByLookupKey<T> opContext);
 
-    <T> P visit(CreateOrUpdate<T> opContext);
+    <T, C> P visit(CreateOrUpdate<T, C> opContext);
 
     <T, U> P visit(CreateOrUpdateInLockedContext<T, U> opContext);
 
     <T, R> P visit(Select<T, R> opContext);
-
   }
 
 }
diff --git a/src/main/java/io/appform/dropwizard/sharding/dao/operations/RunWithCriteria.java b/src/main/java/io/appform/dropwizard/sharding/dao/operations/RunWithCriteria.java
index 47c8a2e1..e505c21a 100644
--- a/src/main/java/io/appform/dropwizard/sharding/dao/operations/RunWithCriteria.java
+++ b/src/main/java/io/appform/dropwizard/sharding/dao/operations/RunWithCriteria.java
@@ -4,27 +4,32 @@
 import lombok.Data;
 import lombok.NonNull;
 import org.hibernate.Session;
-import org.hibernate.criterion.DetachedCriteria;
 
 import java.util.function.Function;
 
 /**
- * Run a query with given criteria inside this shard and returns resulting list.
+ * Run a query inside this shard and return resulting list.
+ * <p>
+ * This operation is generic over the criteria type, supporting both legacy Hibernate API
+ * (DetachedCriteria) and modern JPA Criteria API (QuerySpec), as well as any future
+ * criteria types.
  *
  * @param <T> Return type on performing the operation.
+ * @param <C> Type of criteria used to query (DetachedCriteria, QuerySpec, etc.).
  */
 @Data
 @Builder
-public class RunWithCriteria<T> extends OpContext<T> {
+public class RunWithCriteria<T, C> extends OpContext<T> {
 
   @NonNull
-  private Function<DetachedCriteria, T> handler;
+  private C criteria;
+
   @NonNull
-  private DetachedCriteria detachedCriteria;
+  private Function<C, T> handler;
 
   @Override
   public T apply(Session session) {
-    return handler.apply(detachedCriteria);
+    return handler.apply(criteria);
   }
 
   @Override
diff --git a/src/main/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdate.java b/src/main/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdate.java
index 89cabd71..a5575ceb 100644
--- a/src/main/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdate.java
+++ b/src/main/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdate.java
@@ -20,16 +20,17 @@
  * Else create the entity using the given @Supplier entityGenerator.
  *
  * @param <T> Type of entity on which operation being performed.
+ * @param <C> Type of criteria used to query the entity (DetachedCriteria or QuerySpec).
  */
 @Data
 @Builder
-public class CreateOrUpdate<T> extends OpContext<T> {
+public class CreateOrUpdate<T, C> extends OpContext<T> {
 
-  @NonNull DetachedCriteria criteria;
+  @NonNull C criteria;
   UnaryOperator<T> mutator;
   Supplier<T> entityGenerator;
-  private Function<DetachedCriteria, T> getLockedForWrite;
-  private Function<DetachedCriteria, T> getter;
+  private Function<C, T> getLockedForWrite;
+  private Function<C, T> getter;
   private Function<T, T> saver;
   private BiConsumer<T, T> updater;
 
diff --git a/src/main/java/io/appform/dropwizard/sharding/observers/bucket/BucketKeyPersistor.java b/src/main/java/io/appform/dropwizard/sharding/observers/bucket/BucketKeyPersistor.java
index 4a04950c..0abfda6b 100644
--- a/src/main/java/io/appform/dropwizard/sharding/observers/bucket/BucketKeyPersistor.java
+++ b/src/main/java/io/appform/dropwizard/sharding/observers/bucket/BucketKeyPersistor.java
@@ -166,7 +166,7 @@ public <T> Void visit(RunInSession<T> runInSession) {
     }
 
     @Override
-    public <T> Void visit(RunWithCriteria<T> runWithCriteria) {
+    public <T, C> Void visit(RunWithCriteria<T, C> runWithCriteria) {
         return null;
     }
 
@@ -216,7 +216,7 @@ public <T> Void visit(CreateOrUpdateByLookupKey<T> createOrUpdateByLookupKey) {
     }
 
     @Override
-    public <T> Void visit(CreateOrUpdate<T> createOrUpdate) {
+    public <T, C> Void visit(CreateOrUpdate<T, C> createOrUpdate) {
         final var oldMutator = createOrUpdate.getMutator();
         createOrUpdate.setMutator(result -> {
             if (result != null) {
diff --git a/src/test/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDaoTest.java b/src/test/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDaoTest.java
index 8ea7da44..bf348e3f 100644
--- a/src/test/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDaoTest.java
+++ b/src/test/java/io/appform/dropwizard/sharding/dao/MultiTenantRelationalDaoTest.java
@@ -171,6 +171,36 @@ public void testCreateOrUpdate() throws Exception {
     assertEquals("Some Other Text", updated.getValue());
   }
 
+  @Test
+  public void testCreateOrUpdateWithQuerySpec() throws Exception {
+    // Test creation path
+    val saved = relationalWithAIDao.createOrUpdate("TENANT1", "parent",
+            (QuerySpec<RelationalEntityWithAIKey, RelationalEntityWithAIKey>) (root, query, cb) ->
+                query.where(cb.equal(root.get("key"), "testIdQuerySpec")),
+            e -> e.setValue("Some Other Text"),
+            () -> RelationalEntityWithAIKey.builder()
+                .key("testIdQuerySpec")
+                .value("Some New Text")
+                .build())
+        .orElse(null);
+    assertNotNull(saved);
+    assertEquals("Some New Text", saved.getValue());
+
+    // Test update path
+    val updated = relationalWithAIDao.createOrUpdate("TENANT1", "parent",
+            (QuerySpec<RelationalEntityWithAIKey, RelationalEntityWithAIKey>) (root, query, cb) ->
+                query.where(cb.equal(root.get("key"), "testIdQuerySpec")),
+            e -> e.setValue("Some Other Text"),
+            () -> RelationalEntityWithAIKey.builder()
+                .key("testIdQuerySpec")
+                .value("Some New Text")
+                .build())
+        .orElse(null);
+    assertNotNull(updated);
+    assertEquals(saved.getId(), updated.getId());
+    assertEquals("Some Other Text", updated.getValue());
+  }
+
   @Test
   public void testUpdateUsingQuery() throws Exception {
     val relationalKey = UUID.randomUUID().toString();
@@ -402,6 +432,36 @@ public void testMultiShardRun() {
             .collect(Collectors.toSet()));
   }
 
+  @Test
+  public void testMultiShardRunWithQuerySpec() {
+    val ids = new HashSet<String>();
+    IntStream.range(1, 1_000)
+        .forEach(i -> {
+          try {
+            val id = Integer.toString(i);
+            ids.add(id);
+            relationalDao.save("TENANT1", UUID.randomUUID().toString(),
+                RelationalEntity.builder()
+                    .key(id)
+                    .value("abcd" + i)
+                    .build());
+          } catch (Exception e) {
+            throw new RuntimeException(e);
+          }
+        });
+
+    // QuerySpec equivalent of DetachedCriteria.forClass(RelationalEntity.class) -- select all
+    final QuerySpec<RelationalEntity, RelationalEntity> querySpec = (root, query, cb) -> { };
+
+    assertEquals(ids,
+        relationalDao.run("TENANT1", querySpec)
+            .values()
+            .stream()
+            .flatMap(Collection::stream)
+            .map(v -> ((RelationalEntity) v).getKey())
+            .collect(Collectors.toSet()));
+  }
+
 
   @Test
   public void testPersistenceAndQueryOnSameShard() throws Exception {
diff --git a/src/test/java/io/appform/dropwizard/sharding/dao/RelationalDaoTest.java b/src/test/java/io/appform/dropwizard/sharding/dao/RelationalDaoTest.java
index b7915191..d54c872b 100644
--- a/src/test/java/io/appform/dropwizard/sharding/dao/RelationalDaoTest.java
+++ b/src/test/java/io/appform/dropwizard/sharding/dao/RelationalDaoTest.java
@@ -164,6 +164,36 @@ public void testCreateOrUpdate() throws Exception {
         assertEquals("Some Other Text", updated.getValue());
     }
 
+    @Test
+    public void testCreateOrUpdateWithQuerySpec() throws Exception {
+        // Test creation path
+        val saved = relationalWithAIDao.createOrUpdate("parent",
+                                                       (QuerySpec<RelationalEntityWithAIKey, RelationalEntityWithAIKey>) (root, query, cb) ->
+                                                           query.where(cb.equal(root.get("key"), "testIdQuerySpec")),
+                                                       e -> e.setValue("Some Other Text"),
+                                                       () -> RelationalEntityWithAIKey.builder()
+                                                               .key("testIdQuerySpec")
+                                                               .value("Some New Text")
+                                                               .build())
+                .orElse(null);
+        assertNotNull(saved);
+        assertEquals("Some New Text", saved.getValue());
+
+        // Test update path
+        val updated = relationalWithAIDao.createOrUpdate("parent",
+                                                         (QuerySpec<RelationalEntityWithAIKey, RelationalEntityWithAIKey>) (root, query, cb) ->
+                                                             query.where(cb.equal(root.get("key"), "testIdQuerySpec")),
+                                                         e -> e.setValue("Some Other Text"),
+                                                         () -> RelationalEntityWithAIKey.builder()
+                                                                 .key("testIdQuerySpec")
+                                                                 .value("Some New Text")
+                                                                 .build())
+                .orElse(null);
+        assertNotNull(updated);
+        assertEquals(saved.getId(), updated.getId());
+        assertEquals("Some Other Text", updated.getValue());
+    }
+
     @Test
     public void testUpdateUsingQuery() throws Exception {
         val relationalKey = UUID.randomUUID().toString();
@@ -396,6 +426,36 @@ public void testMultiShardRun() {
                 .collect(Collectors.toSet()));
     }
 
+    @Test
+    public void testRunWithQuerySpec() {
+        val ids = new HashSet<String>();
+        IntStream.range(1, 1_000)
+                .forEach(i -> {
+                    try {
+                        val id = Integer.toString(i);
+                        ids.add(id);
+                        relationalDao.save(UUID.randomUUID().toString(),
+                                           RelationalEntity.builder()
+                                                   .key(id)
+                                                   .value("abcd" + i)
+                                                   .build());
+                    }
+                    catch (Exception e) {
+                        throw new RuntimeException(e);
+                    }
+                });
+
+        // QuerySpec equivalent of DetachedCriteria.forClass(RelationalEntity.class) -- select all
+        final QuerySpec<RelationalEntity, RelationalEntity> querySpec = (root, query, cb) -> { };
+
+        assertEquals(ids, relationalDao.run(querySpec)
+                .values()
+                .stream()
+                .flatMap(Collection::stream)
+                .map(v -> ((RelationalEntity)v).getKey())
+                .collect(Collectors.toSet()));
+    }
+
 
     @Test
     public void testPersistenceAndQueryOnSameShard() throws Exception {
diff --git a/src/test/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdateTest.java b/src/test/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdateTest.java
index 6c5df41b..2fb952f2 100644
--- a/src/test/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdateTest.java
+++ b/src/test/java/io/appform/dropwizard/sharding/dao/operations/relationaldao/CreateOrUpdateTest.java
@@ -2,6 +2,7 @@
 
 import io.appform.dropwizard.sharding.dao.operations.LambdaTestUtils;
 import io.appform.dropwizard.sharding.dao.testdata.entities.Order;
+import io.appform.dropwizard.sharding.query.QuerySpec;
 import lombok.val;
 import org.hibernate.Session;
 import org.hibernate.criterion.DetachedCriteria;
@@ -28,7 +29,7 @@ public void testCreateOrUpdate_creation() {
 
     Order o = Order.builder().id(123).customerId("C1").build();
 
-    val createOrUpdate = CreateOrUpdate.<Order>builder()
+    val createOrUpdate = CreateOrUpdate.<Order, DetachedCriteria>builder()
         .criteria(DetachedCriteria.forClass(Order.class))
         .getLockedForWrite(s -> null)
         .entityGenerator(() -> o)
@@ -54,7 +55,7 @@ public void testCreateOrUpdate_updation() {
 
     Order o = Order.builder().id(123).customerId("C1").build();
 
-    val createOrUpdate = CreateOrUpdate.<Order>builder()
+    val createOrUpdate = CreateOrUpdate.<Order, DetachedCriteria>builder()
         .criteria(DetachedCriteria.forClass(Order.class))
         .getLockedForWrite(s -> o)
         .entityGenerator(() -> o)
@@ -71,4 +72,91 @@ public void testCreateOrUpdate_updation() {
         .accept(Mockito.any(Order.class),
             ArgumentMatchers.argThat((Order x) -> x.getCustomerId().equals("C2")));
   }
+
+  @Test
+  public void testCreateOrUpdateWithQuerySpec_creation() {
+
+    Function<Order, Order> spiedSaver = LambdaTestUtils.spiedFunction((o) -> o);
+    BiConsumer<Order, Order> spiedUpdater = LambdaTestUtils.spiedBiConsumer((o1, o2) -> {
+    });
+
+    Order o = Order.builder().id(123).customerId("C1").build();
+
+    // QuerySpec that matches the order
+    QuerySpec<Order, Order> querySpec = (root, query, cb) -> query.where(cb.equal(root.get("id"), 123));
+
+    val createOrUpdate = CreateOrUpdate.<Order, QuerySpec<Order, Order>>builder()
+        .criteria(querySpec)
+        .getLockedForWrite(s -> null)
+        .entityGenerator(() -> o)
+        .saver(spiedSaver)
+        .updater(spiedUpdater)
+        .mutator(o1 -> o.setCustomerId("C2"))
+        .getter(s -> o)
+        .build();
+
+    Order result = createOrUpdate.apply(session);
+    Assertions.assertEquals(result, o);
+    Mockito.verify(spiedSaver, Mockito.times(1)).apply(Mockito.any(Order.class));
+    Mockito.verify(spiedUpdater, Mockito.times(0))
+        .accept(Mockito.any(Order.class), Mockito.any(Order.class));
+  }
+
+  @Test
+  public void testCreateOrUpdateWithQuerySpec_updation() {
+
+    Function<Order, Order> spiedSaver = LambdaTestUtils.spiedFunction((o) -> o);
+    BiConsumer<Order, Order> spiedUpdater = LambdaTestUtils.spiedBiConsumer((o1, o2) -> {
+    });
+
+    Order o = Order.builder().id(123).customerId("C1").build();
+
+    // QuerySpec that matches the order
+    QuerySpec<Order, Order> querySpec = (root, query, cb) -> query.where(cb.equal(root.get("id"), 123));
+
+    val createOrUpdate = CreateOrUpdate.<Order, QuerySpec<Order, Order>>builder()
+        .criteria(querySpec)
+        .getLockedForWrite(s -> o)
+        .entityGenerator(() -> o)
+        .saver(spiedSaver)
+        .updater(spiedUpdater)
+        .mutator(o1 -> o.setCustomerId("C2"))
+        .getter(s -> o)
+        .build();
+
+    createOrUpdate.apply(session);
+
+    Mockito.verify(spiedSaver, Mockito.times(0)).apply(Mockito.any(Order.class));
+    Mockito.verify(spiedUpdater, Mockito.times(1))
+        .accept(Mockito.any(Order.class),
+            ArgumentMatchers.argThat((Order x) -> x.getCustomerId().equals("C2")));
+  }
+
+  @Test
+  public void testCreateOrUpdateWithQuerySpec_nullEntityGenerator() {
+
+    Function<Order, Order> spiedSaver = LambdaTestUtils.spiedFunction((o) -> o);
+    BiConsumer<Order, Order> spiedUpdater = LambdaTestUtils.spiedBiConsumer((o1, o2) -> {
+    });
+
+    // QuerySpec that matches the order
+    QuerySpec<Order, Order> querySpec = (root, query, cb) -> query.where(cb.equal(root.get("id"), 123));
+
+    val createOrUpdate = CreateOrUpdate.<Order, QuerySpec<Order, Order>>builder()
+        .criteria(querySpec)
+        .getLockedForWrite(s -> null)
+        .entityGenerator(() -> null)
+        .saver(spiedSaver)
+        .updater(spiedUpdater)
+        .mutator(o1 -> o1.setCustomerId("C2"))
+        .getter(s -> null)
+        .build();
+
+    Order result = createOrUpdate.apply(session);
+
+    Assertions.assertNull(result);
+    Mockito.verify(spiedSaver, Mockito.times(0)).apply(Mockito.any(Order.class));
+    Mockito.verify(spiedUpdater, Mockito.times(0))
+        .accept(Mockito.any(Order.class), Mockito.any(Order.class));
+  }
 }
\ No newline at end of file