Fix using coroutine criteria executor with Hibernate (#2775)

(cherry picked from commit 2fab536)

Author: dstepanov

data-hibernate-jpa/src/main/java/io/micronaut/data/hibernate/operations/HibernateJpaOperations.java

@@ -49,6 +49,7 @@
 import io.micronaut.data.operations.reactive.ReactiveRepositoryOperations;
 import io.micronaut.data.runtime.convert.DataConversionService;
 import io.micronaut.data.runtime.operations.ExecutorAsyncOperations;
+import io.micronaut.data.runtime.operations.ExecutorAsyncOperationsSupportingCriteria;
 import io.micronaut.data.runtime.operations.ExecutorReactiveOperations;
 import io.micronaut.transaction.TransactionOperations;
 import jakarta.inject.Named;
@@ -568,7 +569,8 @@ public ExecutorAsyncOperations async() {
             synchronized (this) { // double check
                 executorAsyncOperations = this.asyncOperations;
                 if (executorAsyncOperations == null) {
-                    executorAsyncOperations = new ExecutorAsyncOperations(
+                    executorAsyncOperations = new ExecutorAsyncOperationsSupportingCriteria(
+                        this,
                         this,
                         executorService != null ? executorService : newLocalThreadPool()
                     );

data-runtime/src/main/java/io/micronaut/data/runtime/intercept/criteria/async/AbstractAsyncSpecificationInterceptor.java

@@ -56,13 +56,15 @@ public abstract class AbstractAsyncSpecificationInterceptor<T, R> extends Abstra
      */
     protected AbstractAsyncSpecificationInterceptor(RepositoryOperations operations) {
         super(operations);
-        if (operations instanceof AsyncCapableRepository) {
-            this.asyncOperations = ((AsyncCapableRepository) operations).async();
+        if (operations instanceof AsyncCapableRepository asyncCapableRepository) {
+            this.asyncOperations = asyncCapableRepository.async();
         } else {
             throw new DataAccessException("Datastore of type [" + operations.getClass() + "] does not support asynchronous operations");
         }
         if (operations instanceof AsyncCriteriaRepositoryOperations asyncCriteriaRepositoryOperations) {
             asyncCriteriaOperations = asyncCriteriaRepositoryOperations;
+        } else if (asyncOperations instanceof AsyncCriteriaRepositoryOperations asyncCriteriaRepositoryOperations) {
+            asyncCriteriaOperations = asyncCriteriaRepositoryOperations;
         } else if (operations instanceof AsyncCriteriaCapableRepository repository) {
             asyncCriteriaOperations = repository.async();
         } else {

data-runtime/src/main/java/io/micronaut/data/runtime/operations/ExecutorAsyncOperations.java

@@ -15,6 +15,7 @@
  */
 package io.micronaut.data.runtime.operations;
 
+import io.micronaut.core.annotation.Internal;
 import io.micronaut.core.annotation.NonNull;
 import io.micronaut.core.propagation.PropagatedContext;
 import io.micronaut.core.util.ArgumentUtils;
@@ -62,11 +63,12 @@ public ExecutorAsyncOperations(@NonNull RepositoryOperations operations, @NonNul
         this.executor = executor;
     }
 
-    private <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier) {
+    @Internal
+    <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier) {
         CompletableFuture<T> cf = new CompletableFuture<>();
         PropagatedContext propagatedContext = PropagatedContext.getOrEmpty();
         CompletableFuture.supplyAsync(PropagatedContext.wrapCurrent(supplier), executor).whenComplete((value, throwable) -> {
-            try (PropagatedContext.Scope scope = propagatedContext.propagate()) {
+            try (PropagatedContext.Scope ignore = propagatedContext.propagate()) {
                 if (throwable != null) {
                     cf.completeExceptionally(throwable);
                 } else {

data-runtime/src/main/java/io/micronaut/data/runtime/operations/ExecutorAsyncOperationsSupportingCriteria.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2017-2024 original authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.micronaut.data.runtime.operations;
+
+import io.micronaut.core.annotation.Internal;
+import io.micronaut.data.operations.CriteriaRepositoryOperations;
+import io.micronaut.data.operations.RepositoryOperations;
+import io.micronaut.data.operations.async.AsyncCriteriaRepositoryOperations;
+import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaDelete;
+import jakarta.persistence.criteria.CriteriaQuery;
+import jakarta.persistence.criteria.CriteriaUpdate;
+
+import java.util.List;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.Executor;
+
+/**
+ * A variation of {@link ExecutorAsyncOperations} that supports {@link AsyncCriteriaRepositoryOperations}.
+ * @author Denis Stepanov
+ */
+@Internal
+public final class ExecutorAsyncOperationsSupportingCriteria extends ExecutorAsyncOperations implements AsyncCriteriaRepositoryOperations {
+
+    private final CriteriaRepositoryOperations criteriaRepositoryOperations;
+
+    /**
+     * Default constructor.
+     *
+     * @param operations                   The target operations
+     * @param criteriaRepositoryOperations The criteria operations
+     * @param executor                     The executor to use.
+     */
+    public ExecutorAsyncOperationsSupportingCriteria(RepositoryOperations operations,
+                                                     CriteriaRepositoryOperations criteriaRepositoryOperations,
+                                                     Executor executor) {
+        super(operations, executor);
+        this.criteriaRepositoryOperations = criteriaRepositoryOperations;
+    }
+
+    @Override
+    public CriteriaBuilder getCriteriaBuilder() {
+        return criteriaRepositoryOperations.getCriteriaBuilder();
+    }
+
+    @Override
+    public <R> CompletionStage<R> findOne(CriteriaQuery<R> query) {
+        return supplyAsync(() -> criteriaRepositoryOperations.findOne(query));
+    }
+
+    @Override
+    public <T> CompletionStage<List<T>> findAll(CriteriaQuery<T> query) {
+        return supplyAsync(() -> criteriaRepositoryOperations.findAll(query));
+    }
+
+    @Override
+    public CompletionStage<Number> updateAll(CriteriaUpdate<Number> query) {
+        return supplyAsync(() -> criteriaRepositoryOperations.updateAll(query).orElse(null));
+    }
+
+    @Override
+    public CompletionStage<Number> deleteAll(CriteriaDelete<Number> query) {
+        return supplyAsync(() -> criteriaRepositoryOperations.deleteAll(query).orElse(null));
+    }
+}

doc-examples/hibernate-example-kotlin/src/main/kotlin/example/ParentSuspendRepository.kt

@@ -4,11 +4,12 @@ import io.micronaut.core.annotation.NonNull
 import io.micronaut.data.annotation.Join
 import io.micronaut.data.annotation.Repository
 import io.micronaut.data.repository.GenericRepository
+import io.micronaut.data.repository.jpa.kotlin.CoroutineJpaSpecificationExecutor
 import java.util.*
 import jakarta.transaction.Transactional
 
 @Repository
-interface ParentSuspendRepository : GenericRepository<Parent, Int> {
+interface ParentSuspendRepository : GenericRepository<Parent, Int>, CoroutineJpaSpecificationExecutor<Parent> {
 
     @Join(value = "children", type = Join.Type.FETCH)
     suspend fun findById(id: Int): Optional<Parent>

doc-examples/hibernate-example-kotlin/src/test/kotlin/example/HibernateTxTest.kt

@@ -1,5 +1,8 @@
 package example
 
+import io.micronaut.data.repository.jpa.criteria.PredicateSpecification
+import io.micronaut.data.repository.jpa.criteria.QuerySpecification
+import io.micronaut.data.runtime.criteria.get
 import io.micronaut.test.extensions.junit5.annotation.MicronautTest
 import org.junit.jupiter.api.Assertions.assertTrue
 import jakarta.inject.Inject
@@ -154,4 +157,19 @@ class HibernateTxTest {
         Assertions.assertEquals(0, service.countForCustomDb())
     }
 
+    @Test
+    @Order(12)
+    fun coroutineCriteria() {
+        runBlocking {
+            val parent = Parent("abc", Collections.emptyList())
+            val saved = repositorySuspended.save(parent)
+
+            val found1 = repositorySuspended.findOne(QuerySpecification { root, query, criteriaBuilder -> criteriaBuilder.equal(root.get<String>("name"), "abc")})
+            Assertions.assertEquals(found1!!.id, saved.id)
+
+            val found2 = repositorySuspended.findOne(PredicateSpecification { root, criteriaBuilder ->  criteriaBuilder.equal(root.get<String>("name"), "abc")})
+            Assertions.assertEquals(found2!!.id, saved.id)
+        }
+    }
+
 }