Feature/refactor and add remaining queryspec

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",

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",

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);
     }

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);
-
   }
 
 }

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

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;
 

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) {

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 {